eta_index_lib/utils/mgodb/mgo_config.go

package mgodb

import (
	"context"
	"eta/eta_index_lib/utils"
	"github.com/qiniu/qmgo"
	"github.com/qiniu/qmgo/options"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/event"
	mgoptions "go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/writeconcern"
	"sync"
)

var (
	ConnectTimeoutMS = int64(3000)
	SocketTimeoutMS  = int64(500000)
	MaxPoolSize      = uint64(50)
	MinPoolSize      = uint64(0)
)

//func init() {
//	fmt.Println("start MgoNewClient")
//	MgoNewClient()
//	fmt.Println("end MgoNewClient")
//}

type MgoConfig struct {
	Host          string `json:"host"`
	Url           string `json:"url"`
	Port          string `json:"port"`
	Username      string `json:"username"`
	Password      string `json:"password"`
	AuthSource    string `json:"authSource"`
	AuthMechanism string `json:"auth_mechanism"`
	Database      string `json:"database"`
}

// MgoNewClient
// @Description: 创建一个mongodb客户端链接
// @author: Roc
// @datetime 2024-04-25 14:01:14
// @return *qmgo.Client
func MgoNewClient(mgoConfig MgoConfig) *qmgo.Client {
	// 创建cmdMonitor，用于打印SQL
	//startedCommands := make(map[int64]bson.Raw)
	startedCommands := sync.Map{} // map[int64]bson.Raw

	cmdMonitor := &event.CommandMonitor{
		Started: func(_ context.Context, evt *event.CommandStartedEvent) {
			startedCommands.Store(evt.RequestID, evt.Command)
			//startedCommands[evt.RequestID] = evt.Command
		},
		Succeeded: func(_ context.Context, evt *event.CommandSucceededEvent) {
			//log.Printf("Command: %v Reply: %v\n",
			//	startedCommands[evt.RequestID],
			//	evt.Reply,
			//)
			var commands bson.Raw
			v, ok := startedCommands.Load(evt.RequestID)
			if ok {
				commands = v.(bson.Raw)
			}
			//utils.MongoLog.Info("\n【MongoDB】[%.3fms] [%v] %v \n", float64(evt.Duration)/1e6, commands, evt.Reply)
			utils.MongoLog.Info("\n【MongoDB】[%.3fms] [%v] \n", float64(evt.Duration)/1e6, commands)
		},
		Failed: func(_ context.Context, evt *event.CommandFailedEvent) {
			//log.Printf("Command: %v Failure: %v\n",
			//	startedCommands[evt.RequestID],
			//	evt.Failure,
			//)
			var commands bson.Raw
			v, ok := startedCommands.Load(evt.RequestID)
			if ok {
				commands = v.(bson.Raw)
			}
			utils.MongoLog.Info("\n【MongoDB】[%.3fms] [%v] \n %v \n", float64(evt.Duration)/1e6, commands, evt.Failure)
		},
	}

	// 创建options
	ops := options.ClientOptions{ClientOptions: &mgoptions.ClientOptions{
		WriteConcern: &writeconcern.WriteConcern{W: "majority"},
	}}
	ops.SetMonitor(cmdMonitor)

	ctx := context.Background()

	//mongodb+srv://myDatabaseUser:D1fficultP%40ssw0rd@mongodb0.example.com/?authSource=admin&replicaSet=myRepl

	//var mongoUrl string
	//if mgoConfig.Password != "" {
	//	mongoUrl = "mongodb://" + mgoConfig.Username + ":" + mgoConfig.Password + "@" +
	//		mgoConfig.Host + ":" + mgoConfig.Port
	//
	//	if mgoConfig.AuthSource != `` {
	//		mongoUrl += mongoUrl + "?&authSource=" + mgoConfig.AuthSource
	//	}
	//} else {
	//	mongoUrl = "mongodb://" + mgoConfig.Host + ":" + mgoConfig.Port
	//}

	// 创建一个数据库链接
	client, err := qmgo.NewClient(ctx, &qmgo.Config{
		Uri: mgoConfig.Url,
		Auth: &qmgo.Credential{
			AuthMechanism: mgoConfig.AuthMechanism,
			AuthSource:    mgoConfig.AuthSource,
			Username:      mgoConfig.Username,
			Password:      mgoConfig.Password,
			//PasswordSet:   false,
		},
		Database:         mgoConfig.Database,
		ConnectTimeoutMS: &ConnectTimeoutMS,
		SocketTimeoutMS:  &SocketTimeoutMS,
		MaxPoolSize:      &MaxPoolSize,
		MinPoolSize:      &MinPoolSize,
	}, ops)

	if err != nil {
		panic("MongoDB连接异常：" + err.Error())
	}

	return client
}

func MgoGetColl(collName string) (cli *qmgo.QmgoClient, err error) {
	ctx := context.Background()
	config := &qmgo.Config{
		//Uri:              MGO_URL,
		//Database:         MGO_DB,
		Coll:             collName,
		ConnectTimeoutMS: &ConnectTimeoutMS,
		SocketTimeoutMS:  &SocketTimeoutMS,
		MaxPoolSize:      &MaxPoolSize,
		MinPoolSize:      &MinPoolSize,
	}
	cli, err = qmgo.Open(ctx, config)
	return
}