init

2026-04-30 23:26:14 +08:00
commit 2dfcbbbb22
76 changed files with 17478 additions and 0 deletions
@@ -0,0 +1,315 @@
+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
+}
@@ -0,0 +1,48 @@
+package database
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/EthanCodeCraft/xlgo-core/config"
+	"github.com/EthanCodeCraft/xlgo-core/logger"
+
+	"github.com/redis/go-redis/v9"
+	"go.uber.org/zap"
+)
+
+var (
+	// RedisClient 全局 Redis 客户端
+	RedisClient *redis.Client
+)
+
+// InitRedis 初始化 Redis 连接
+func InitRedis(cfg *config.Config) error {
+	RedisClient = redis.NewClient(&redis.Options{
+		Addr:     cfg.Redis.Addr(),
+		Password: cfg.Redis.Password,
+		DB:       cfg.Redis.DB,
+	})
+
+	// 测试连接
+	ctx := context.Background()
+	if err := RedisClient.Ping(ctx).Err(); err != nil {
+		return fmt.Errorf("Redis 连接失败: %w", err)
+	}
+
+	logger.Info("Redis 连接成功", zap.String("addr", cfg.Redis.Addr()))
+	return nil
+}
+
+// CloseRedis 关闭 Redis 连接
+func CloseRedis() error {
+	if RedisClient != nil {
+		return RedisClient.Close()
+	}
+	return nil
+}
+
+// GetRedis 获取 Redis 客户端
+func GetRedis() *redis.Client {
+	return RedisClient
+}