eta_api/utils/mgodb/mgo_config.go

package mgodb

import (
	"context"
	"eta_gn/eta_api/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"
	"sync"
)

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

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"`
}

func MgoNewClient(mgoConfig MgoConfig) *qmgo.Client {
	startedCommands := sync.Map{} // map[int64]bson.Raw

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

	ops := options.ClientOptions{ClientOptions: &mgoptions.ClientOptions{}}
	ops.SetMonitor(cmdMonitor)

	ctx := context.Background()



	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,
		},
		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{
		Coll:             collName,
		ConnectTimeoutMS: &ConnectTimeoutMS,
		SocketTimeoutMS:  &SocketTimeoutMS,
		MaxPoolSize:      &MaxPoolSize,
		MinPoolSize:      &MinPoolSize,
	}
	cli, err = qmgo.Open(ctx, config)
	return
}