95 lines
2.4 KiB
Go
95 lines
2.4 KiB
Go
package mongo
|
||
|
||
import (
|
||
"context"
|
||
"git.hlsq.asia/alex/service-common/config"
|
||
"git.hlsq.asia/alex/service-common/log"
|
||
"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"
|
||
"time"
|
||
)
|
||
|
||
// Colors
|
||
const (
|
||
Reset = "\033[0m"
|
||
BrightGreen = "\033[92m"
|
||
)
|
||
|
||
type DBName string
|
||
|
||
var clients = make(map[DBName]*mongo.Client)
|
||
|
||
func Init(cfg map[string]*config.MongoConfig) error {
|
||
for name, oneConfig := range cfg {
|
||
if client, err := initOneClient(oneConfig); err != nil {
|
||
return err
|
||
} else {
|
||
clients[DBName(name)] = client
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func initOneClient(cfg *config.MongoConfig) (*mongo.Client, error) {
|
||
opts := options.Client().
|
||
ApplyURI(cfg.URI).
|
||
SetMonitor(&event.CommandMonitor{
|
||
Started: func(ctx context.Context, startedEvent *event.CommandStartedEvent) {
|
||
log.Infof(BrightGreen+"[MongoDB] Started, Command: %v"+Reset, startedEvent.Command.String())
|
||
},
|
||
Failed: func(ctx context.Context, failedEvent *event.CommandFailedEvent) {
|
||
log.Errorf(BrightGreen+"[MongoDB] [%v] Failed, Failure: %v"+Reset, failedEvent.Duration, failedEvent.Failure)
|
||
},
|
||
})
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
|
||
client, err := mongo.Connect(ctx, opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err = client.Ping(ctx, readpref.Primary()); err != nil {
|
||
_ = client.Disconnect(context.Background())
|
||
return nil, err
|
||
}
|
||
|
||
return client, nil
|
||
}
|
||
|
||
// GetClient 返回 mongo.Client，你可以通过 .Database("xxx") 获取具体数据库
|
||
func GetClient(dbName DBName) *mongo.Client {
|
||
if c, ok := clients[dbName]; ok {
|
||
return c
|
||
}
|
||
log.Errorf("mongo client %s not found", dbName)
|
||
return nil
|
||
}
|
||
|
||
// GetDB 是便捷方法，直接返回 *mongo.Database（假设 dbName 对应数据库名）
|
||
// 如果你的配置中 dbName 和实际数据库名一致，可这样用
|
||
func GetDB(dbName DBName) *mongo.Database {
|
||
client := GetClient(dbName)
|
||
if client == nil {
|
||
return nil
|
||
}
|
||
// 假设配置中的 key 就是数据库名；若需分离，可在 config 中加字段
|
||
return client.Database(string(dbName))
|
||
}
|
||
|
||
func Close() error {
|
||
if clients == nil {
|
||
return nil
|
||
}
|
||
for name, client := range clients {
|
||
if client != nil {
|
||
if err := client.Disconnect(context.Background()); err != nil {
|
||
log.Errorf("close mongo client %s error: %v", name, err)
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
}
|