57 lines
1.3 KiB
Go
57 lines
1.3 KiB
Go
package config
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"time"
|
|
|
|
"yh_web/server/pkg/logger"
|
|
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"go.mongodb.org/mongo-driver/v2/mongo"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
)
|
|
|
|
var MongoClient *mongo.Client
|
|
|
|
// ConnectMongoDB 连接本地 MongoDB
|
|
func ConnectMongoDB(uri string) error {
|
|
clientOpts := options.Client().ApplyURI(uri)
|
|
client, err := mongo.Connect(clientOpts)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// 验证连接
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
var result bson.M
|
|
if err = client.Database("admin").RunCommand(ctx, bson.D{{Key: "ping", Value: 1}}).Decode(&result); err != nil {
|
|
return err
|
|
}
|
|
|
|
MongoClient = client
|
|
log.Println("MongoDB 连接成功")
|
|
logger.Log("config/database", "MongoDB 连接成功")
|
|
return nil
|
|
}
|
|
|
|
// GetDB 获取指定数据库；未连接 MongoDB 时返回 nil
|
|
func GetDB(name string) *mongo.Database {
|
|
if MongoClient == nil {
|
|
return nil
|
|
}
|
|
return MongoClient.Database(name)
|
|
}
|
|
|
|
// CloseMongoDB 关闭连接
|
|
func CloseMongoDB() {
|
|
if MongoClient != nil {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
_ = MongoClient.Disconnect(ctx)
|
|
log.Println("MongoDB 连接已关闭")
|
|
logger.Log("config/database", "MongoDB 连接已关闭")
|
|
}
|
|
}
|