base-golang/pkg/database/mongo.go

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
}