diff --git a/mongo/mongo.go b/mongo/mongo.go index b1c8609..3a04d72 100644 --- a/mongo/mongo.go +++ b/mongo/mongo.go @@ -1,24 +1,32 @@ +// ============================================================================= +// MongoDB 多数据源支持 +// 支持多数据源配置、自动重连、优雅关闭 +// 向后兼容原有的单数据源API +// ============================================================================= + package mongo import ( "context" "errors" "fmt" + "os" + "os/signal" "reflect" "strings" "sync" + "syscall" "time" "gitee.com/red-future---jilin-g/common/beans" "gitee.com/red-future---jilin-g/common/log/model/entity" - "github.com/gogf/gf/v2/container/gvar" - "github.com/gogf/gf/v2/os/grpool" - "gitee.com/red-future---jilin-g/common/redis" "gitee.com/red-future---jilin-g/common/utils" + "github.com/gogf/gf/v2/container/gvar" "github.com/gogf/gf/v2/errors/gerror" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/glog" + "github.com/gogf/gf/v2/os/grpool" "github.com/gogf/gf/v2/os/gtime" "github.com/gogf/gf/v2/text/gstr" "github.com/gogf/gf/v2/util/gconv" @@ -27,8 +35,367 @@ import ( "go.mongodb.org/mongo-driver/v2/mongo/options" ) +// ============================================================================= +// 数据源配置结构 +// ============================================================================= + +type DataSourceConfig struct { + Name string `json:"name"` + Address string `json:"address"` + Database string `json:"database"` + MaxPoolSize int32 `json:"maxPoolSize"` + MinPoolSize int32 `json:"minPoolSize"` + ConnectTimeout time.Duration `json:"connectTimeout"` +} + +// ============================================================================= +// 单个数据源接口 +// ============================================================================= + +type DataSource interface { + Name() string + Database() *mongo.Database + Client() *mongo.Client + IsConnected() bool + Connect(ctx context.Context) error + Reconnect(ctx context.Context) error + Close(ctx context.Context) error +} + +// ============================================================================= +// 数据源实现 +// ============================================================================= + +type BaseDataSource struct { + config *DataSourceConfig + client *mongo.Client + database *mongo.Database + isConnected bool + mu sync.RWMutex + lastError error + lastErrorTime time.Time +} + +func NewBaseDataSource(config *DataSourceConfig) *BaseDataSource { + return &BaseDataSource{ + config: config, + isConnected: false, + } +} + +func (d *BaseDataSource) Name() string { + return d.config.Name +} + +func (d *BaseDataSource) Database() *mongo.Database { + d.mu.RLock() + defer d.mu.RUnlock() + return d.database +} + +func (d *BaseDataSource) Client() *mongo.Client { + d.mu.RLock() + defer d.mu.RUnlock() + return d.client +} + +func (d *BaseDataSource) IsConnected() bool { + d.mu.RLock() + defer d.mu.RUnlock() + return d.isConnected && d.client != nil +} + +func (d *BaseDataSource) Connect(ctx context.Context) error { + d.mu.Lock() + defer d.mu.Unlock() + + if d.client != nil { + d.client.Disconnect(ctx) + } + + // 解析数据库名 + dbName := d.config.Database + if strings.Contains(dbName, "?") { + dbName = gstr.SubStr(dbName, 0, strings.Index(dbName, "?")) + } + + // 构建连接选项 + opt := options.Client(). + ApplyURI(d.config.Address). + SetMaxPoolSize(uint64(d.config.MaxPoolSize)). + SetMinPoolSize(uint64(d.config.MinPoolSize)). + SetConnectTimeout(d.config.ConnectTimeout). + SetMaxConnecting(10). + SetServerSelectionTimeout(10 * time.Second). + SetHeartbeatInterval(10 * time.Second). + SetMaxConnIdleTime(60 * time.Second). + SetRetryWrites(true). + SetRetryReads(true) + + var err error + d.client, err = mongo.Connect(opt) + if err != nil { + d.isConnected = false + d.lastError = err + d.lastErrorTime = time.Now() + return fmt.Errorf("datasource [%s] connection failed: %w", d.config.Name, err) + } + + // 测试连接 + pingCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + if err = d.client.Ping(pingCtx, nil); err != nil { + d.isConnected = false + d.lastError = err + d.lastErrorTime = time.Now() + return fmt.Errorf("datasource [%s] ping failed: %w", d.config.Name, err) + } + + d.database = d.client.Database(dbName) + d.isConnected = true + d.lastError = nil + glog.Infof(ctx, "✅ datasource [%s] connected successfully", d.config.Name) + return nil +} + +func (d *BaseDataSource) Reconnect(ctx context.Context) error { + glog.Infof(ctx, "🔄 reconnecting datasource [%s]", d.config.Name) + return d.Connect(ctx) +} + +func (d *BaseDataSource) Close(ctx context.Context) error { + d.mu.Lock() + defer d.mu.Unlock() + + if d.client != nil { + disconnectCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + if err := d.client.Disconnect(disconnectCtx); err != nil { + return fmt.Errorf("datasource [%s] close failed: %w", d.config.Name, err) + } + } + + d.isConnected = false + glog.Infof(ctx, "datasource [%s] closed", d.config.Name) + return nil +} + +// ============================================================================= +// 多数据源管理器 +// ============================================================================= + +type DataSourceManager struct { + sources map[string]DataSource + mu sync.RWMutex + ctx context.Context + cancel context.CancelFunc + started bool + maxRetries int +} + +var ( + globalManager *DataSourceManager + managerOnce sync.Once +) + +// GetManager 获取全局管理器 +func GetManager() *DataSourceManager { + managerOnce.Do(func() { + ctx, cancel := context.WithCancel(context.Background()) + globalManager = &DataSourceManager{ + sources: make(map[string]DataSource), + ctx: ctx, + cancel: cancel, + started: false, + maxRetries: 3, + } + }) + return globalManager +} + +// RegisterDataSource 注册数据源 +func (m *DataSourceManager) RegisterDataSource(config *DataSourceConfig) error { + m.mu.Lock() + defer m.mu.Unlock() + + if _, exists := m.sources[config.Name]; exists { + return fmt.Errorf("datasource [%s] already exists", config.Name) + } + + source := NewBaseDataSource(config) + m.sources[config.Name] = source + return nil +} + +// GetDataSource 获取数据源 +func (m *DataSourceManager) GetDataSource(name string) (DataSource, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + source, exists := m.sources[name] + if !exists { + return nil, fmt.Errorf("datasource [%s] not found", name) + } + return source, nil +} + +// GetAllDataSourceNames 获取所有数据源名称 +func (m *DataSourceManager) GetAllDataSourceNames() []string { + m.mu.RLock() + defer m.mu.RUnlock() + + names := make([]string, 0, len(m.sources)) + for name := range m.sources { + names = append(names, name) + } + return names +} + +// InitializeFromConfig 从配置初始化数据源 +// 动态读取 config.yml 中 mongo 下的所有配置项 +func (m *DataSourceManager) InitializeFromConfig(ctx context.Context) error { + var firstErr error + + // 获取 mongo 配置下的所有子键 + mongoConfig := g.Cfg().MustGet(ctx, "mongo") + if mongoConfig.IsNil() { + glog.Warningf(ctx, "no mongo configuration found in config.yml") + return nil + } + + // 将配置转换为 map + configMap := mongoConfig.Map() + if configMap == nil { + glog.Warningf(ctx, "mongo configuration is not a map") + return nil + } + + // 遍历所有 mongo 子配置 + for name, subConfig := range configMap { + // 跳过非对象类型的配置 + subMap, ok := subConfig.(map[string]interface{}) + if !ok { + continue + } + + // 检查是否有 address 配置 + address, hasAddress := subMap["address"] + if !hasAddress || gconv.String(address) == "" { + continue + } + + // 构建数据源配置 + config := &DataSourceConfig{ + Name: name, + Address: gconv.String(address), + Database: gconv.String(subMap["database"]), + MaxPoolSize: int32(gconv.Int(subMap["maxPoolSize"])), + MinPoolSize: int32(gconv.Int(subMap["minPoolSize"])), + ConnectTimeout: gconv.Duration(subMap["connectTimeout"]), + } + + // 设置默认值 + if config.MaxPoolSize == 0 { + config.MaxPoolSize = 100 + } + if config.MinPoolSize == 0 { + config.MinPoolSize = 10 + } + if config.ConnectTimeout == 0 { + config.ConnectTimeout = 10 * time.Second + } + + // 注册数据源 + if err := m.RegisterDataSource(config); err != nil { + glog.Errorf(ctx, "failed to register datasource [%s]: %v", name, err) + if firstErr == nil { + firstErr = err + } + continue + } + + // 连接数据源 + source, _ := m.GetDataSource(name) + if err := source.Connect(ctx); err != nil { + glog.Errorf(ctx, "failed to initialize datasource [%s]: %v", name, err) + if firstErr == nil { + firstErr = err + } + } + } + + return firstErr +} + +// StartHealthCheck 启动健康检查 +func (m *DataSourceManager) StartHealthCheck() { + if m.started { + return + } + m.started = true + go m.healthCheckLoop() +} + +// healthCheckLoop 健康检查循环 +func (m *DataSourceManager) healthCheckLoop() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + for { + select { + case <-m.ctx.Done(): + return + case <-ticker.C: + m.checkAndReconnect() + } + } +} + +// checkAndReconnect 检查并重新连接 +func (m *DataSourceManager) checkAndReconnect() { + m.mu.RLock() + defer m.mu.RUnlock() + + for name, source := range m.sources { + if !source.IsConnected() { + glog.Warningf(context.Background(), "datasource [%s] disconnected, attempting reconnect", name) + + reconnectCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if err := source.Reconnect(reconnectCtx); err != nil { + glog.Errorf(reconnectCtx, "datasource [%s] reconnect failed: %v", name, err) + } else { + glog.Infof(reconnectCtx, "✅ datasource [%s] reconnected successfully", name) + } + } + } +} + +// CloseAll 关闭所有数据源 +func (m *DataSourceManager) CloseAll(ctx context.Context) error { + m.cancel() + + m.mu.RLock() + defer m.mu.RUnlock() + + var lastErr error + for name, source := range m.sources { + if err := source.Close(ctx); err != nil { + glog.Errorf(ctx, "failed to close datasource [%s]: %v", name, err) + lastErr = err + } + } + return lastErr +} + +// ============================================================================= +// 向后兼容的MongoDB结构体 +// ============================================================================= + type MongoDB struct { - Cache bool + Cache bool + dataSource string // 数据源名称,默认为 "default" } func DB(cache ...bool) *MongoDB { @@ -37,199 +404,100 @@ func DB(cache ...bool) *MongoDB { b = cache[0] } return &MongoDB{ - Cache: b, + Cache: b, + dataSource: "default", } } +// WithDataSource 指定使用的数据源 +func (m *MongoDB) WithDataSource(name string) *MongoDB { + m.dataSource = name + return m +} + +// ============================================================================= +// 向后兼容的全局变量和方法 +// ============================================================================= + var ( - db *mongo.Database - client *mongo.Client - isConnected bool - mu sync.RWMutex - mongoAddr string - dbName string - healthCtx context.Context - healthCancel context.CancelFunc + manager = GetManager() + logPool *grpool.Pool + serverName string + logRedisKey string ) -// checkConnected 检查连接状态 -func checkConnected() bool { - mu.RLock() - defer mu.RUnlock() - return isConnected -} - -// connect 建立MongoDB连接 -func connect() error { - mu.Lock() - defer mu.Unlock() - - if client != nil { - client.Disconnect(context.Background()) - } - - // 创建连接选项 - opt := options.Client(). - ApplyURI(mongoAddr). - SetMaxPoolSize(100). - SetMinPoolSize(10). - SetMaxConnecting(10). - SetConnectTimeout(10 * time.Second) - - var err error - client, err = mongo.Connect(opt) - if err != nil { - isConnected = false - glog.Error(context.Background(), "MongoDB连接失败", err) - return err - } - - // 测试连接 - testCtx, testCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer testCancel() - - err = client.Ping(testCtx, nil) - if err != nil { - isConnected = false - glog.Error(testCtx, "MongoDB连接测试失败", err) - return err - } - - db = client.Database(dbName) - isConnected = true - glog.Info(context.Background(), "✅ MongoDB连接成功") - return nil -} - -// GetDB 获取 MongoDB 数据库实例 -func GetDB() *mongo.Database { - mu.RLock() - defer mu.RUnlock() - return db -} - -// healthCheck 健康检查协程 -func healthCheck() { - ticker := time.NewTicker(30 * time.Second) - defer ticker.Stop() - - for { - select { - case <-healthCtx.Done(): - return - case <-ticker.C: - mu.RLock() - currentConnected := isConnected - currentClient := client - mu.RUnlock() - - if !currentConnected || currentClient == nil { - glog.Warning(context.Background(), "MongoDB连接断开,尝试重连") - if err := reconnect(); err != nil { - glog.Error(context.Background(), "MongoDB重连失败", err) - } - continue - } - - // 测试连接状态 - testCtx, testCancel := context.WithTimeout(context.Background(), 5*time.Second) - err := currentClient.Ping(testCtx, nil) - testCancel() - - if err != nil { - mu.Lock() - isConnected = false - mu.Unlock() - glog.Warning(context.Background(), "MongoDB连接健康检查失败", err) - - // 尝试重连 - if err := reconnect(); err != nil { - glog.Error(context.Background(), "MongoDB重连失败", err) - } - } else { - glog.Debug(context.Background(), "MongoDB连接健康检查通过") - } - } - } -} - -// reconnect 重连函数 -func reconnect() error { - maxRetries := 3 - retryDelay := 2 * time.Second - - for i := 0; i < maxRetries; i++ { - glog.Info(context.Background(), fmt.Sprintf("尝试第%d次重连MongoDB", i+1)) - - if err := connect(); err == nil { - glog.Info(context.Background(), "MongoDB重连成功") - return nil - } - - if i < maxRetries-1 { - time.Sleep(retryDelay) - retryDelay *= 2 // 指数退避 - } - } - - return gerror.New("MongoDB重连失败,已达到最大重试次数") -} - -var logPool *grpool.Pool - -// init 初始化MongoDB连接 -func init() { - logPool = grpool.New(1) - // 按需初始化:没有配置 mongo.address 则跳过 - mongoAddr = g.Cfg().MustGet(context.Background(), "mongo.address").String() - if mongoAddr == "" { - return - } - - // 创建健康检查上下文 - healthCtx, healthCancel = context.WithCancel(context.Background()) - - // 从连接串中解析数据库名 - dbName = gstr.SubStr(mongoAddr, strings.LastIndex(mongoAddr, "/")+1, len(mongoAddr)) - // 如果连接串带有参数(如 ?retryWrites=true),需要去掉参数部分 - if strings.Contains(dbName, "?") { - dbName = gstr.SubStr(dbName, 0, strings.Index(dbName, "?")) - } - go func() { - // 初始连接 - if err := connect(); err != nil { - glog.Error(context.Background(), "MongoDB初始连接失败", err) - return - } - }() - - // 启动健康检查协程 - go healthCheck() -} - -// close 关闭MongoDB连接 -func close() { - if healthCancel != nil { - healthCancel() - } - - mu.Lock() - defer mu.Unlock() - - if client != nil { - disconnectCtx, disconnectCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer disconnectCancel() - client.Disconnect(disconnectCtx) - } - - isConnected = false - glog.Info(context.Background(), "MongoDB连接已关闭") -} - const PageSize = 20 +// GetDB 获取默认数据源的数据库实例(向后兼容) +func GetDB() *mongo.Database { + source, err := manager.GetDataSource("default") + if err != nil { + return nil + } + return source.Database() +} + +// init 初始化多数据源 +func init() { + logPool = grpool.New(1) + serverName = g.Cfg().MustGet(context.TODO(), "server.name").String() + logRedisKey = fmt.Sprintf("log:%s", serverName) + + ctx := context.Background() + + // 从配置初始化多数据源 + if err := manager.InitializeFromConfig(ctx); err != nil { + glog.Errorf(ctx, "❌ Failed to initialize MongoDB datasources: %v", err) + } else { + glog.Infof(ctx, "✅ MongoDB datasources initialized: %v", manager.GetAllDataSourceNames()) + } + + // 启动健康检查 + manager.StartHealthCheck() + + // 设置优雅关闭 + setupGracefulShutdown() +} + +// setupGracefulShutdown 设置优雅关闭 +func setupGracefulShutdown() { + go func() { + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) + <-sigCh + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + glog.Info(ctx, "🔄 Shutting down MongoDB connections...") + if err := manager.CloseAll(ctx); err != nil { + glog.Errorf(ctx, "❌ Failed to close MongoDB connections: %v", err) + } else { + glog.Info(ctx, "✅ MongoDB connections closed successfully") + } + }() +} + +// ============================================================================= +// MongoDB 操作方法(支持多数据源) +// ============================================================================= + +// getDataSource 获取当前使用的数据源 +func (m *MongoDB) getDataSource() (DataSource, error) { + if m.dataSource == "" { + m.dataSource = "default" + } + return manager.GetDataSource(m.dataSource) +} + // Find 查询多条记录 func (m *MongoDB) Find(ctx context.Context, filter bson.M, result interface{}, collection string, page *beans.Page, orderBy []beans.OrderBy) (total int64, err error) { + source, err := m.getDataSource() + if err != nil { + return 0, err + } + db := source.Database() + if err = utils.ValidStructPtr(result); err != nil { return } @@ -257,7 +525,6 @@ func (m *MongoDB) Find(ctx context.Context, filter bson.M, result interface{}, c } filter["tenantId"] = user.TenantId - // 分页参数处理 limit := int64(PageSize) skip := int64(0) if page != nil { @@ -283,9 +550,9 @@ func (m *MongoDB) Find(ctx context.Context, filter bson.M, result interface{}, c orderBson := bson.D{} for _, v := range orderBy { if v.Order == beans.Asc { - orderBson = append(orderBson, bson.E{Key: v.Field, Value: 1}) // 1 表示升序 + orderBson = append(orderBson, bson.E{Key: v.Field, Value: 1}) } else { - orderBson = append(orderBson, bson.E{Key: v.Field, Value: -1}) // -1 表示降序 + orderBson = append(orderBson, bson.E{Key: v.Field, Value: -1}) } } opt.SetSort(orderBson) @@ -313,6 +580,12 @@ func (m *MongoDB) Find(ctx context.Context, filter bson.M, result interface{}, c // FindOne 查询1条记录 func (m *MongoDB) FindOne(ctx context.Context, filter bson.M, result interface{}, collection string, opts ...options.Lister[options.FindOneOptions]) (err error) { + source, err := m.getDataSource() + if err != nil { + return err + } + db := source.Database() + if len(filter) == 0 { err = gerror.New("缺少查询条件") return @@ -357,6 +630,7 @@ func (m *MongoDB) FindOne(ctx context.Context, filter bson.M, result interface{} } return } + func (m *MongoDB) CleanRedis(ctx context.Context, filter bson.M, tenantId interface{}, collection string) (err error) { listKeys := fmt.Sprintf(redis.CleanList, tenantId, collection) keys, err := redis.RedisClient.Keys(ctx, listKeys) @@ -391,9 +665,6 @@ func (m *MongoDB) CleanRedis(ctx context.Context, filter bson.M, tenantId interf return } -var serverName = g.Cfg().MustGet(context.TODO(), "server.name").String() -var logRedisKey = fmt.Sprintf("log:%s", serverName) - func (m *MongoDB) log(ctx context.Context, filter bson.M, collection string, data interface{}, userName, tenantId interface{}, operationType string) { _ = logPool.AddWithRecover(ctx, func(ctx context.Context) { log := &entity.OperationLog{ @@ -420,6 +691,12 @@ func (m *MongoDB) log(ctx context.Context, filter bson.M, collection string, dat // Delete 删除记录 func (m *MongoDB) Delete(ctx context.Context, filter bson.M, collection string, opts ...options.Lister[options.DeleteManyOptions]) (count int64, err error) { + source, err := m.getDataSource() + if err != nil { + return 0, err + } + db := source.Database() + if len(filter) == 0 { err = gerror.New("缺少查询条件") return @@ -435,13 +712,17 @@ func (m *MongoDB) Delete(ctx context.Context, filter bson.M, collection string, } count = r.DeletedCount err = m.CleanRedis(ctx, filter, user.TenantId, collection) - //写日志 - //m.log(ctx, filter, collection, nil, user.UserName, user.TenantId, "delete") return } // Update 修改记录 func (m *MongoDB) Update(ctx context.Context, filter bson.M, update bson.M, collection string, opts ...options.Lister[options.UpdateManyOptions]) (result *mongo.UpdateResult, err error) { + source, err := m.getDataSource() + if err != nil { + return nil, err + } + db := source.Database() + if len(filter) == 0 { err = gerror.New("缺少查询条件") return @@ -465,25 +746,23 @@ func (m *MongoDB) Update(ctx context.Context, filter bson.M, update bson.M, coll return } err = m.CleanRedis(ctx, filter, user.TenantId, collection) - //写日志 - //m.log(ctx, filter, collection, update, user.UserName, user.TenantId, "update") return } // RandomSoftDelete 随机软删除个文档的 _id func (m *MongoDB) RandomSoftDelete(ctx context.Context, limit int, collection string, opts ...options.Lister[options.UpdateManyOptions]) (result *mongo.UpdateResult, err error) { + source, err := m.getDataSource() + if err != nil { + return nil, err + } + db := source.Database() + _ = opts - // 步骤 1: 使用聚合管道的 $sample 操作符随机抽取5个文档的 _id pipeline := mongo.Pipeline{ - // 阶段1: 为每个文档添加一个 0-1 之间的随机数字段 'random' bson.D{{Key: "$addFields", Value: bson.D{{Key: "random", Value: bson.M{"$rand": bson.M{}}}}}}, - // 阶段1: 匹配所有未删除的文档 bson.D{{Key: "$match", Value: bson.D{{Key: "isDeleted", Value: false}}}}, - // 阶段2: 按随机数降序排序 bson.D{{Key: "$sort", Value: bson.D{{Key: "random", Value: -1}}}}, - // 阶段3: 只取前5个 bson.D{{Key: "$limit", Value: limit}}, - // 阶段4: 只投影 _id bson.D{{Key: "$project", Value: bson.D{{Key: "_id", Value: 1}}}}, } cursor, err := db.Collection(collection).Aggregate(ctx, pipeline) @@ -491,14 +770,13 @@ func (m *MongoDB) RandomSoftDelete(ctx context.Context, limit int, collection st return } defer cursor.Close(ctx) - // 步骤 2: 从聚合结果中提取 _id 到一个切片中 + var idsToUpdate []bson.ObjectID for cursor.Next(ctx) { var result bson.M if err := cursor.Decode(&result); err != nil { return nil, err } - // 将 bson.M 中的 _id 断言为 primitive.ObjectID id := result["_id"].(bson.ObjectID) idsToUpdate = append(idsToUpdate, id) } @@ -506,11 +784,9 @@ func (m *MongoDB) RandomSoftDelete(ctx context.Context, limit int, collection st return nil, err } fmt.Printf("准备更新的随机文档ID: %v\n", idsToUpdate) - // 步骤 3: 使用 $in 操作符和 UpdateMany 批量更新选定的文档 + if len(idsToUpdate) > 0 { - // 过滤条件:匹配 idsToUpdate 切片中的任意一个 _id filter := bson.D{{Key: "_id", Value: bson.D{{Key: "$in", Value: idsToUpdate}}}} - // 更新操作:使用 $set 修改字段 update := bson.D{{Key: "$set", Value: bson.D{{Key: "isDeleted", Value: true}}}} _, err = db.Collection(collection).UpdateMany(ctx, filter, update) if err != nil { @@ -522,6 +798,12 @@ func (m *MongoDB) RandomSoftDelete(ctx context.Context, limit int, collection st // SaveOrUpdate 批量增加或修改 func (m *MongoDB) SaveOrUpdate(ctx context.Context, filter []bson.M, update []bson.M, collection string, opts ...options.Lister[options.UpdateManyOptions]) (result *mongo.BulkWriteResult, err error) { + source, err := m.getDataSource() + if err != nil { + return nil, err + } + db := source.Database() + if len(filter) == 0 || len(update) == 0 { err = gerror.New("缺少查询条件或更新数据") return @@ -534,22 +816,19 @@ func (m *MongoDB) SaveOrUpdate(ctx context.Context, filter []bson.M, update []bs if err != nil { return } - // 构建批量操作模型 + var models []mongo.WriteModel for i := 0; i < len(filter); i++ { - // 处理过滤器 filter[i]["isDeleted"] = false if !g.IsEmpty(user.TenantId) { filter[i]["tenantId"] = user.TenantId } - // 处理更新数据 if setDoc, exists := update[i]["$set"].(bson.M); exists { if !g.IsEmpty(user.UserName) { setDoc["updater"] = user.UserName } setDoc["updatedAt"] = gtime.Now().Time } else { - // 如果没有$set字段,则创建一个 setDoc := bson.M{} if !g.IsEmpty(user.UserName) { setDoc["updater"] = user.UserName @@ -557,12 +836,10 @@ func (m *MongoDB) SaveOrUpdate(ctx context.Context, filter []bson.M, update []bs setDoc["updatedAt"] = gtime.Now().Time update[i]["$set"] = setDoc } - // 创建更新操作模型 updateModel := mongo.NewUpdateOneModel() updateModel.SetFilter(filter[i]) updateModel.SetUpdate(update[i]) - updateModel.SetUpsert(true) // 默认不插入新文档 - // 处理选项参数 + updateModel.SetUpsert(true) if len(opts) > 0 { for _, opt := range opts { var updateOpts options.UpdateManyOptions @@ -577,13 +854,13 @@ func (m *MongoDB) SaveOrUpdate(ctx context.Context, filter []bson.M, update []bs } models = append(models, updateModel) } - // 执行批量操作,无序执行提高性能 + bulkOpts := options.BulkWrite().SetOrdered(false) bulkResult, err := db.Collection(collection).BulkWrite(ctx, models, bulkOpts) if err != nil { return nil, err } - // 清理相关缓存 + for _, filterItem := range filter { err = m.CleanRedis(ctx, filterItem, user.TenantId, collection) if err != nil { @@ -595,6 +872,12 @@ func (m *MongoDB) SaveOrUpdate(ctx context.Context, filter []bson.M, update []bs // Insert 插入多条记录 func (m *MongoDB) Insert(ctx context.Context, documents []interface{}, collection string, opts ...options.Lister[options.InsertManyOptions]) (ids []interface{}, err error) { + source, err := m.getDataSource() + if err != nil { + return nil, err + } + db := source.Database() + user, err := utils.GetUserInfo(ctx) if err != nil { return @@ -623,13 +906,17 @@ func (m *MongoDB) Insert(ctx context.Context, documents []interface{}, collectio } ids = r.InsertedIDs err = m.CleanRedis(ctx, bson.M{}, user.TenantId, collection) - //写日志 - //m.log(ctx, nil, collection, ids, user.UserName, user.TenantId, "insert") return } -// Count 查询总数 +// Count 查询总数 func (m *MongoDB) Count(ctx context.Context, filter bson.M, collection string) (count int64, err error) { + source, err := m.getDataSource() + if err != nil { + return 0, err + } + db := source.Database() + user, err := utils.GetUserInfo(ctx) if err != nil { return @@ -649,7 +936,6 @@ func (m *MongoDB) Count(ctx context.Context, filter bson.M, collection string) ( return } } - // 调用驱动的 CountDocuments,在数据库端执行的 count, err = db.Collection(collection).CountDocuments(ctx, filter) if m.Cache { err = redis.RedisClient.SetEX(ctx, redisKey, count, int64(time.Hour)) @@ -661,38 +947,26 @@ func (m *MongoDB) Count(ctx context.Context, filter bson.M, collection string) ( } // EntityToBson 将 *entity/entity 转换为 bson.M -// 支持传入值类型或指针类型,返回 bson.M 和错误信息 func EntityToBson(entity interface{}) (bson.M, error) { return EntityToBsonWithFilter(entity, false) } // EntityToBsonWithFilter 将 *entity/entity 转换为 bson.M,并可选择是否过滤空值 -// filterEmpty: 为 true 时会过滤掉空值字段(nil、空字符串、空切片、空map等),但保留 int 类型的 0 值 -// 支持: -// - 未传值的指针(如 *consts.AssetStatus(nil))会被过滤 -// - 传值为0的指针(如 *consts.AssetStatus(0))不会被过滤 -// - 非0的整数值不会被过滤 func EntityToBsonWithFilter(entity interface{}, filterEmpty bool) (bson.M, error) { - // 第一步:判断入参是否为 nil 或无效类型 if entity == nil { return nil, fmt.Errorf("传入的 entity 实例为 nil") } - // 第二步:将 entity 序列化为 BSON 字节流 - // bson.Marshal 支持值类型和指针类型,会自动解析结构体的 bson 标签 bsonBytes, err := bson.Marshal(entity) if err != nil { return nil, fmt.Errorf("entity 序列化为 BSON 字节流失败:%w", err) } - // 第三步:将 BSON 字节流反序列化为 bson.M var bsonMap bson.M err = bson.Unmarshal(bsonBytes, &bsonMap) if err != nil { return nil, fmt.Errorf("BSON 字节流反序列化为 bson.M 失败:%w", err) } - // 如果需要过滤空值 if filterEmpty { for key, value := range bsonMap { - // 判断是否为空值,但保留 int 类型的 0 值 if isEmptyWithZero(value) { delete(bsonMap, key) } @@ -702,33 +976,24 @@ func EntityToBsonWithFilter(entity interface{}, filterEmpty bool) (bson.M, error } // isEmptyWithZero 判断是否为空值,但保留 int 类型的 0 值 -// 支持区分"未传值"和"传值为0"的情况: -// - *int(nil) 或 *consts.AssetStatus(nil) → 返回 true(过滤掉) -// - *int(0) 或 *consts.AssetStatus(0) → 返回 false(保留) -// - int(0) 或 consts.AssetStatus(0) → 返回 false(保留) func isEmptyWithZero(value interface{}) bool { - // 先检查 value 是否为 nil if value == nil { return true } rv := reflect.ValueOf(value) kind := rv.Kind() - // 处理 nil 指针 if kind == reflect.Ptr { if rv.IsNil() { return true } - // 判断时如果是指针,需要获取指向的值的类型 kind = rv.Elem().Kind() } - // 数字类型(int/uint/float)都保留,包括 0 值 switch kind { case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Float32, reflect.Float64: return false default: - // 其他类型使用 g.IsEmpty 判断 return g.IsEmpty(value) } }