宇恒一号官网

server/config/database.go

@@ -0,0 +1,56 @@
+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 连接已关闭")
+	}
+}