common增加nats消息队列支持
This commit is contained in:
749
mongo/mongo.go
749
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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user