Initial cloud-services repo - gateway service + pkg modules

pkg/mongo/client.go

@@ -0,0 +1,130 @@
+package mongo
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"fiskerinc.com/modules/logger"
+	"fiskerinc.com/modules/utils/envtool"
+
+	"github.com/pkg/errors"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+
+	mongotrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/go.mongodb.org/mongo-driver/mongo"
+)
+
+var (
+	conn_str   = envtool.GetEnv("MONGO_CONN_STR", "mongodb://REPLACE_ME:REPLACE_ME@localhost:27017/?authSource=admin&w=majority")
+	StandardDB = envtool.GetEnv("MONGO_DB_NAME", "db")
+	ODXDB      = envtool.GetEnv("MONGO_ODX_DB_NAME", "odx_db")
+	timeout    = time.Duration(envtool.GetEnvInt("MONGO_CLIENT_TIMEOUT", 60)) * time.Second
+)
+
+// Supply the database to connect to, either mongo.StandardDB or mongo.ODXDB
+func NewClient(dbToConnectTo string) (Client, error) {
+	var conn Client
+	client, database, err := NewMongoConnection(dbToConnectTo)
+	if err != nil {
+		return conn, err
+	}
+	conn = &Conn{
+		conn:     client,
+		database: database,
+	}
+
+	return conn, nil
+}
+
+// Sets up a mongo connection, if connection string is empty, defaults to the env connection string
+func NewMongoConnection(databaseName string) (client *mongo.Client, db *mongo.Database, err error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	logger.Info().Msgf("Connection string %s", conn_str)
+
+	serverAPIOptions := options.ServerAPI(options.ServerAPIVersion1)
+	opts := options.Client().ApplyURI(conn_str).SetServerAPIOptions(serverAPIOptions)
+	opts.Monitor = mongotrace.NewMonitor()
+	client, err = mongo.Connect(ctx, opts)
+	if err != nil {
+		return nil, nil, errors.WithStack(err)
+	}
+	db = client.Database(databaseName)
+	return
+}
+
+type Client interface {
+	GetVehicles() VehiclesCollectionInterface
+	SetVehicles(v VehiclesCollectionInterface)
+
+	GetFleets() FleetsCollectionInterface
+	SetFleets(f FleetsCollectionInterface)
+
+	Collection(name string) *mongo.Collection
+	Close() error
+
+	Ping(ctx context.Context) error
+}
+
+type Conn struct {
+	conn     *mongo.Client
+	database *mongo.Database
+
+	vehicles     VehiclesCollectionInterface
+	vehiclesOnce sync.Once
+
+	fleets     FleetsCollectionInterface
+	fleetsOnce sync.Once
+}
+
+func (c *Conn) GetVehicles() VehiclesCollectionInterface {
+	c.vehiclesOnce.Do(func() {
+		if c.vehicles == nil {
+			c.vehicles = NewVehiclesCollection(
+				NewCollection(c.Collection("vehicles")),
+			)
+		}
+	})
+	return c.vehicles
+}
+
+func (c *Conn) SetVehicles(v VehiclesCollectionInterface) {
+	c.vehicles = v
+}
+
+func (c *Conn) GetFleets() FleetsCollectionInterface {
+	c.fleetsOnce.Do(func() {
+		if c.fleets == nil {
+			c.fleets = NewFleetsCollection(
+				NewCollection(c.Collection("fleets")),
+			)
+		}
+	})
+	return c.fleets
+}
+
+func (c *Conn) SetFleets(f FleetsCollectionInterface) {
+	c.fleets = f
+}
+
+func (c *Conn) Collection(name string) *mongo.Collection {
+	return c.database.Collection(name)
+}
+
+func (c *Conn) Close() error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	err := c.conn.Disconnect(ctx)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+func (c *Conn) Ping(ctx context.Context) error {
+	return c.conn.Ping(ctx, nil)
+}