first commit

pkg/database/mongo.go

@@ -0,0 +1,77 @@
+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
+}