mongodb/mongodb.go

package mongodb

import (
	"context"
	"fmt"
	"strings"

	"git.selly.red/Selly-Modules/logger"
	"go.mongodb.org/mongo-driver/event"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
)

// Config ...
type Config struct {
	Host    string
	DBName  string
	Monitor *event.CommandMonitor

	TLS        *ConnectTLSOpts
	Standalone *ConnectStandaloneOpts
}

// ConnectTLSOpts ...
type ConnectTLSOpts struct {
	ReplSet             string
	CaFile              string
	CertKeyFile         string
	CertKeyFilePassword string
	ReadPreferenceMode  string
}

// ConnectStandaloneOpts ...
type ConnectStandaloneOpts struct {
	AuthMechanism string
	AuthSource    string
	Username      string
	Password      string
}

var db *mongo.Database

// Connect to mongo server
func Connect(cfg Config) (*mongo.Database, error) {
	if cfg.TLS != nil && cfg.TLS.ReplSet != "" {
		return connectWithTLS(cfg)
	}
	connectOptions := options.ClientOptions{}
	opts := cfg.Standalone
	// Set auth if existed
	if opts != nil && opts.Username != "" && opts.Password != "" {
		connectOptions.Auth = &options.Credential{
			AuthMechanism: opts.AuthMechanism,
			AuthSource:    opts.AuthSource,
			Username:      opts.Username,
			Password:      opts.Password,
		}
	}
	if cfg.Monitor != nil {
		connectOptions.SetMonitor(cfg.Monitor)
	}

	// Connect
	client, err := mongo.Connect(context.Background(), connectOptions.ApplyURI(cfg.Host))
	if err != nil {
		logger.Error("connect", logger.LogData{
			Source:  "mongodb.Connect",
			Message: err.Error(),
			Data: logger.Map{
				"host": cfg.Host,
			},
		})
		return nil, err
	}

	fmt.Printf("⚡️[mongodb]: connected to %s/%s \n", cfg.Host, cfg.DBName)

	// Set data
	db = client.Database(cfg.DBName)
	return db, nil
}

func connectWithTLS(cfg Config) (*mongo.Database, error) {
	ctx := context.Background()
	opts := cfg.TLS

	caFile, err := initFileFromBase64String("ca.pem", opts.CaFile)
	if err != nil {
		return nil, err
	}
	certFile, err := initFileFromBase64String("cert.pem", opts.CertKeyFile)
	if err != nil {
		return nil, err
	}
	pwd := base64DecodeToString(opts.CertKeyFilePassword)
	uri := getURIWithTLS(cfg, caFile.Name(), certFile.Name(), pwd)
	readPref := getReadPref(opts.ReadPreferenceMode)
	clientOpts := options.Client().SetReadPreference(readPref).SetReplicaSet(opts.ReplSet).ApplyURI(uri)
	if cfg.Monitor != nil {
		clientOpts.SetMonitor(cfg.Monitor)
	}
	client, err := mongo.Connect(ctx, clientOpts)
	if err != nil {
		return nil, err
	}
	if err := client.Ping(ctx, readpref.SecondaryPreferred()); err != nil {
		return nil, err
	}
	db = client.Database(cfg.DBName)

	fmt.Printf("⚡️[mongodb/tls]: connected to %s/%s \n", cfg.Host, cfg.DBName)
	return db, err
}

func getURIWithTLS(cfg Config, caFilePath, certFilePath, pwd string) string {
	host := cfg.Host
	if strings.Contains(host, "?") {
		host += "&"
	} else {
		if !strings.HasSuffix(host, "/") {
			host += "/?"
		} else {
			host += "?"
		}
	}
	s := "%stls=true&tlsCAFile=./%s&tlsCertificateKeyFile=./%s&tlsCertificateKeyFilePassword=%s&authMechanism=MONGODB-X509"
	uri := fmt.Sprintf(s, host, caFilePath, certFilePath, pwd)
	return uri
}

// GetInstance ...
func GetInstance() *mongo.Database {
	return db
}

func getReadPref(mode string) *readpref.ReadPref {
	m, err := readpref.ModeFromString(mode)
	if err != nil {
		m = readpref.SecondaryPreferredMode
	}
	readPref, err := readpref.New(m)
	if err != nil {
		fmt.Println("mongodb.getReadPref err: ", err, m)
	}
	return readPref
}