宇恒一号官网

server/config/constants.go

@@ -0,0 +1,4 @@
+package config
+
+// DBName 数据库名，可由环境变量 MONGODB_DB 覆盖
+var DBName = "yxd-agent-testing"

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 连接已关闭")
+	}
+}

server/config/db_structure.go

@@ -0,0 +1,78 @@
+package config
+
+import (
+	"context"
+	"time"
+
+	"go.mongodb.org/mongo-driver/v2/bson"
+)
+
+// DBStructure 数据库结构信息
+type DBStructure struct {
+	Databases []DatabaseInfo `json:"databases"`
+}
+
+// DatabaseInfo 数据库信息
+type DatabaseInfo struct {
+	Name        string          `json:"name"`
+	Collections []CollInfo      `json:"collections"`
+}
+
+// CollInfo 集合信息
+type CollInfo struct {
+	Name    string      `json:"name"`
+	Count   int64       `json:"count"`
+	Sample  interface{} `json:"sample,omitempty"` // 采样一条文档看结构
+}
+
+// GetDBStructure 获取 MongoDB 数据结构
+func GetDBStructure() (*DBStructure, error) {
+	if MongoClient == nil {
+		return nil, nil
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// 列出所有数据库（排除系统库）
+	dbs, err := MongoClient.ListDatabaseNames(ctx, bson.M{
+		"name": bson.M{"$nin": []string{"admin", "config", "local"}},
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	result := &DBStructure{Databases: make([]DatabaseInfo, 0, len(dbs))}
+
+	for _, dbName := range dbs {
+		db := MongoClient.Database(dbName)
+		colls, err := db.ListCollectionNames(ctx, bson.M{})
+		if err != nil {
+			continue
+		}
+
+		dbInfo := DatabaseInfo{
+			Name:        dbName,
+			Collections: make([]CollInfo, 0, len(colls)),
+		}
+
+		for _, collName := range colls {
+			coll := db.Collection(collName)
+			count, _ := coll.CountDocuments(ctx, bson.M{})
+
+			// 采样一条文档
+			var sample bson.M
+			_ = coll.FindOne(ctx, bson.M{}).Decode(&sample)
+
+			dbInfo.Collections = append(dbInfo.Collections, CollInfo{
+				Name:   collName,
+				Count:  count,
+				Sample: sample,
+			})
+		}
+
+		result.Databases = append(result.Databases, dbInfo)
+	}
+
+	return result, nil
+}