package database import ( "context" "fmt" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/readpref" "time" ) var _ MongoDB = (*mongoDB)(nil) type MongoDB interface { i() GetDB() *mongo.Database Close() error } type MongoDBConfig struct { Addr string `yaml:"addr"` User string `yaml:"user"` Pass string `yaml:"pass"` Name string `yaml:"name"` Timeout time.Duration `yaml:"timeout"` } type mongoDB struct { client *mongo.Client db *mongo.Database timeout time.Duration } func (m *mongoDB) i() {} func NewMongoDB(cfg MongoDBConfig) (MongoDB, error) { timeout := cfg.Timeout * time.Second connectCtx, connectCancelFunc := context.WithTimeout(context.Background(), timeout) defer connectCancelFunc() var auth string if len(cfg.User) > 0 && len(cfg.Pass) > 0 { auth = fmt.Sprintf("%s:%s@", cfg.User, cfg.Pass) } client, err := mongo.Connect(connectCtx, options.Client().ApplyURI( fmt.Sprintf("mongodb://%s%s", auth, cfg.Addr), )) if err != nil { return nil, err } pingCtx, pingCancelFunc := context.WithTimeout(context.Background(), timeout) defer pingCancelFunc() err = client.Ping(pingCtx, readpref.Primary()) if err != nil { return nil, err } return &mongoDB{ client: client, db: client.Database(cfg.Name), timeout: timeout, }, nil } func (m *mongoDB) GetDB() *mongo.Database { return m.db } func (m *mongoDB) Close() error { disconnectCtx, disconnectCancelFunc := context.WithTimeout(context.Background(), m.timeout) defer disconnectCancelFunc() err := m.client.Disconnect(disconnectCtx) if err != nil { return err } return nil }