131 lines
3.1 KiB
Go
131 lines
3.1 KiB
Go
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)
|
|
}
|