package mongo import ( "context" "sync" "time" "github.com/fiskerinc/cloud-services/pkg/logger" "github.com/fiskerinc/cloud-services/pkg/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) }