DeployHelper/internal/conf/db.go

package conf

import (
	"database/sql"
	"ego/pkg/logger"
	"fmt"
	"log"
	"os"
	"sync"
	"time"

	"go.uber.org/zap"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	glogger "gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
)

var (
	// Db 数据库链接单例
	Db        *gorm.DB
	once      sync.Once
	dbInitErr error
)

// NewDb 初始化数据库连接
func NewDb() *gorm.DB {
	return Db
}

// Database 初始化mysql连接
func Database(connString string) error {
	once.Do(func() {
		dbInitErr = initDatabase(connString)
	})
	return dbInitErr
}

// initDatabase 实际的数据库初始化逻辑
func initDatabase(connString string) error {
	if connString == "" {
		err := fmt.Errorf("数据库连接字符串不能为空")
		logger.Error(nil, "数据库初始化失败", zap.Error(err))
		return err
	}

	// 初始化GORM日志配置
	logLevel := getGormLogLevel()
	newLogger := glogger.New(
		log.New(os.Stdout, "\r\n", log.LstdFlags),
		glogger.Config{
			SlowThreshold:             time.Second,
			LogLevel:                  logLevel,
			IgnoreRecordNotFoundError: true,
			Colorful:                  false,
		},
	)

	// 配置 gorm
	db, err := gorm.Open(mysql.Open(connString), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
		Logger:                                   newLogger,
		NamingStrategy:                           schema.NamingStrategy{TablePrefix: ""},
	})

	if err != nil {
		logger.Error(nil, "数据库连接失败", zap.Error(err), zap.String("connString", maskConnectionString(connString)))
		return fmt.Errorf("数据库连接失败: %w", err)
	}

	// 获取底层的sql.DB对象
	sqlDB, err := db.DB()
	if err != nil {
		logger.Error(nil, "获取数据库实例失败", zap.Error(err))
		return fmt.Errorf("获取数据库实例失败: %w", err)
	}

	// 配置连接池
	if err := configureConnectionPool(sqlDB); err != nil {
		logger.Error(nil, "配置数据库连接池失败", zap.Error(err))
		return fmt.Errorf("配置数据库连接池失败: %w", err)
	}

	// 测试数据库连接
	if err := sqlDB.Ping(); err != nil {
		logger.Error(nil, "数据库连接测试失败", zap.Error(err))
		return fmt.Errorf("数据库连接测试失败: %w", err)
	}

	Db = db
	logger.Info(nil, "数据库连接成功",
		zap.String("connString", maskConnectionString(connString)),
		zap.Int("maxOpenConns", 20),
		zap.Int("maxIdleConns", 10))

	return nil
}

// configureConnectionPool 配置数据库连接池
func configureConnectionPool(sqlDB *sql.DB) error {
	// 设置空闲连接池中连接的最大数量
	sqlDB.SetMaxIdleConns(10)
	// 设置打开数据库连接的最大数量
	sqlDB.SetMaxOpenConns(20)
	// 设置连接可复用的最大时间
	sqlDB.SetConnMaxLifetime(time.Hour)
	// 设置连接空闲的最大时间
	sqlDB.SetConnMaxIdleTime(time.Minute * 30)

	return nil
}

// getGormLogLevel 根据环境变量获取GORM日志级别
func getGormLogLevel() glogger.LogLevel {
	switch os.Getenv("GORM_LOG_LEVEL") {
	case "silent":
		return glogger.Silent
	case "error":
		return glogger.Error
	case "warn":
		return glogger.Warn
	case "info":
		return glogger.Info
	default:
		// 根据应用环境决定默认级别
		if os.Getenv("GIN_MODE") == "release" {
			return glogger.Error
		}
		return glogger.Info
	}
}

// maskConnectionString 遮盖连接字符串中的敏感信息
func maskConnectionString(connString string) string {
	// 简单遮盖密码部分，实际项目中可以使用更复杂的正则表达式
	return "***已遮盖敏感信息***"
}

// Migration 执行数据迁移
func Migration() error {
	if Db == nil {
		return fmt.Errorf("数据库连接未初始化")
	}

	// 自动迁移模式
	models := []any{}

	for _, model := range models {
		if err := Db.AutoMigrate(model); err != nil {
			logger.Error(nil, "数据库迁移失败",
				zap.Error(err),
				zap.String("model", fmt.Sprintf("%T", model)))
			return fmt.Errorf("迁移模型 %T 失败: %w", model, err)
		}
	}

	logger.Info(nil, "数据库迁移完成", zap.Int("models", len(models)))
	return nil
}

// GetDB 获取数据库连接（用于外部调用）
func GetDB() (*gorm.DB, error) {
	if Db == nil {
		return nil, fmt.Errorf("数据库连接未初始化")
	}
	return Db, nil
}

// closeDatabase 安全关闭数据库连接
func closeDatabase() error {
	if Db == nil {
		return nil
	}

	sqlDB, err := Db.DB()
	if err != nil {
		return fmt.Errorf("获取数据库连接失败: %w", err)
	}

	return sqlDB.Close()
}