package database import ( "context" "fmt" "math/rand" "sync" "time" "github.com/EthanCodeCraft/xlgo-core/config" "github.com/EthanCodeCraft/xlgo-core/logger" "go.uber.org/zap" "gorm.io/driver/mysql" "gorm.io/gorm" gormlogger "gorm.io/gorm/logger" ) var ( // DB 主库实例(写操作) DB *gorm.DB // DBRead 读库实例(读操作) DBRead *gorm.DB // replicas 从库列表 replicas []*gorm.DB // replicaMutex 从库选择锁 replicaMutex sync.Mutex ) // InitMySQL 初始化 MySQL 连接(带重试机制) func InitMySQL(cfg *config.Config) error { var err error // GORM 日志配置 var gormLogLevel gormlogger.LogLevel if cfg.IsDevelopment() { gormLogLevel = gormlogger.Info } else { gormLogLevel = gormlogger.Warn } gormConfig := &gorm.Config{ Logger: gormlogger.Default.LogMode(gormLogLevel), } // 重试配置 maxRetries := 5 retryDelay := time.Second for i := range maxRetries { // 连接主库 DB, err = gorm.Open(mysql.Open(cfg.Database.DSN()), gormConfig) if err == nil { sqlDB, err := DB.DB() if err == nil { sqlDB.SetMaxIdleConns(cfg.Database.MaxIdleConns) sqlDB.SetMaxOpenConns(cfg.Database.MaxOpenConns) sqlDB.SetConnMaxLifetime(time.Hour) if err := sqlDB.Ping(); err == nil { logger.Info("MySQL 主库连接成功", zap.String("host", cfg.Database.Host), zap.Int("port", cfg.Database.Port)) return nil } } } logger.Warnf("MySQL 连接失败,第 %d/%d 次重试: %v", i+1, maxRetries, err) time.Sleep(retryDelay) retryDelay *= 2 if retryDelay > 30*time.Second { retryDelay = 30 * time.Second } } return fmt.Errorf("MySQL 连接失败(重试 %d 次): %w", maxRetries, err) } // InitMySQLWithReplicas 初始化 MySQL 主从连接 // masterDSN: 主库连接字符串 // replicaDSNs: 从库连接字符串列表 func InitMySQLWithReplicas(cfg *config.Config, replicaDSNs []string) error { // 先初始化主库 if err := InitMySQL(cfg); err != nil { return err } // 初始化从库 if len(replicaDSNs) > 0 { var gormLogLevel gormlogger.LogLevel if cfg.IsDevelopment() { gormLogLevel = gormlogger.Info } else { gormLogLevel = gormlogger.Warn } gormConfig := &gorm.Config{ Logger: gormlogger.Default.LogMode(gormLogLevel), } for i, dsn := range replicaDSNs { replicaDB, err := gorm.Open(mysql.Open(dsn), gormConfig) if err != nil { logger.Warnf("MySQL 从库 %d 连接失败: %v", i+1, err) continue } sqlDB, err := replicaDB.DB() if err != nil { logger.Warnf("MySQL 从库 %d 获取连接池失败: %v", i+1, err) continue } sqlDB.SetMaxIdleConns(cfg.Database.MaxIdleConns) sqlDB.SetMaxOpenConns(cfg.Database.MaxOpenConns / 2) // 从库连接数可适当减少 sqlDB.SetConnMaxLifetime(time.Hour) if err := sqlDB.Ping(); err != nil { logger.Warnf("MySQL 从库 %d Ping 失败: %v", i+1, err) continue } replicas = append(replicas, replicaDB) logger.Info("MySQL 从库连接成功", zap.Int("index", i+1)) } // 设置默认读库 if len(replicas) > 0 { DBRead = replicas[0] } else { DBRead = DB // 无从库时使用主库 } } else { DBRead = DB // 无从库配置时使用主库 } return nil } // GetReadDB 获取读库实例(自动选择从库) func GetReadDB() *gorm.DB { if len(replicas) == 0 { return DB } replicaMutex.Lock() defer replicaMutex.Unlock() // 随机选择一个从库 idx := rand.Intn(len(replicas)) return replicas[idx] } // GetWriteDB 获取写库实例(主库) func GetWriteDB() *gorm.DB { return DB } // GetDB 获取数据库实例(默认主库,兼容旧代码) func GetDB() *gorm.DB { return DB } // GetReplicas 获取所有从库实例 func GetReplicas() []*gorm.DB { return replicas } // UseMaster 强制使用主库(用于事务或需要实时数据的场景) func UseMaster(ctx context.Context) context.Context { return context.WithValue(ctx, "db_mode", "master") } // UseReplica 强制使用从库(用于报表查询等场景) func UseReplica(ctx context.Context) context.Context { return context.WithValue(ctx, "db_mode", "replica") } // GetDBFromContext 根据上下文选择数据库 func GetDBFromContext(ctx context.Context) *gorm.DB { mode, ok := ctx.Value("db_mode").(string) if !ok { return GetReadDB() } switch mode { case "master": return DB case "replica": return GetReadDB() default: return GetReadDB() } } // DBResolver 数据库解析器(用于 GORM 钩子) type DBResolver struct{} // BeforeQuery 查询前钩子,自动路由到从库 func (r *DBResolver) BeforeQuery(db *gorm.DB) { // 如果在事务中,使用主库 if db.Statement.ConnPool != nil { return } // 检查上下文 ctx := db.Statement.Context if ctx != nil { mode, ok := ctx.Value("db_mode").(string) if ok && mode == "master" { return // 强制主库 } } // 读操作路由到从库 if len(replicas) > 0 && DBRead != nil { db.Statement.ConnPool = DBRead.Statement.ConnPool } } // AutoMigrate 自动迁移数据库表结构(由应用重写) func AutoMigrate() error { logger.Info("数据库表结构迁移完成") return nil } // Close 关闭数据库连接 func Close() error { if DB == nil { return nil } sqlDB, err := DB.DB() if err != nil { return err } return sqlDB.Close() } // CloseAll 关闭所有数据库连接(包括从库) func CloseAll() error { // 关闭主库 if err := Close(); err != nil { return err } // 关闭从库 for _, replica := range replicas { if replica == nil { continue } sqlDB, err := replica.DB() if err != nil { continue } sqlDB.Close() } replicas = nil DBRead = nil return nil } // Transaction 事务操作(自动使用主库) func Transaction(fn func(tx *gorm.DB) error) error { return DB.Transaction(fn) } // TransactionWithContext 带上下文的事务操作 func TransactionWithContext(ctx context.Context, fn func(tx *gorm.DB) error) error { return DB.WithContext(ctx).Transaction(fn) } // ReadQuery 读查询(自动路由到从库) func ReadQuery(ctx context.Context, model any, query string, args ...any) error { db := GetDBFromContext(ctx) return db.WithContext(ctx).Where(query, args...).Find(model).Error } // WriteQuery 写查询(强制使用主库) func WriteQuery(ctx context.Context, model any, query string, args ...any) error { return DB.WithContext(ctx).Where(query, args...).Find(model).Error } // HealthCheck 健康检查 func HealthCheck() map[string]bool { result := make(map[string]bool) // 检查主库 if DB != nil { sqlDB, err := DB.DB() if err == nil && sqlDB.Ping() == nil { result["master"] = true } else { result["master"] = false } } else { result["master"] = false } // 检查从库 for i, replica := range replicas { if replica != nil { sqlDB, err := replica.DB() if err == nil && sqlDB.Ping() == nil { result[fmt.Sprintf("replica_%d", i+1)] = true } else { result[fmt.Sprintf("replica_%d", i+1)] = false } } else { result[fmt.Sprintf("replica_%d", i+1)] = false } } return result }