goalfylearning-admin/internal/storage/database.go

package storage

import (
	"context"
	"fmt"
	"goalfymax-admin/internal/config"
	"goalfymax-admin/internal/models"
	"goalfymax-admin/pkg/utils"
	"time"

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

var DB *gorm.DB

// GormLoggerAdapter GORM logger适配器
type GormLoggerAdapter struct {
	logger   *utils.Logger
	logLevel logger.LogLevel
}

// LogMode 设置日志级别
func (g *GormLoggerAdapter) LogMode(level logger.LogLevel) logger.Interface {
	g.logLevel = level
	return g
}

// Info 记录信息日志
func (g *GormLoggerAdapter) Info(ctx context.Context, msg string, data ...interface{}) {
	g.logger.Info(msg, zap.Any("data", data))
}

// Warn 记录警告日志
func (g *GormLoggerAdapter) Warn(ctx context.Context, msg string, data ...interface{}) {
	g.logger.Warn(msg, zap.Any("data", data))
}

// Error 记录错误日志
func (g *GormLoggerAdapter) Error(ctx context.Context, msg string, data ...interface{}) {
	g.logger.Error(msg, zap.Any("data", data))
}

// Trace 记录跟踪日志
func (g *GormLoggerAdapter) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
	// 根据日志级别决定是否记录
	if g.logLevel < logger.Info {
		return
	}

	elapsed := time.Since(begin)
	sql, rows := fc()

	if err != nil {
		g.logger.Error("SQL执行失败",
			zap.String("sql", sql),
			zap.Int64("rows", rows),
			zap.Duration("elapsed", elapsed),
			zap.Error(err),
		)
	} else {
		g.logger.Debug("SQL执行成功",
			zap.String("sql", sql),
			zap.Int64("rows", rows),
			zap.Duration("elapsed", elapsed),
		)
	}
}

// InitDatabase 初始化数据库连接
func InitDatabase(appLogger *utils.Logger) error {
	cfg := config.GetConfig()

	// 配置数据库连接
	dbConfig := mysql.Config{
		DSN: cfg.Database.DSN,
	}

	// 创建GORM logger适配器
	gormLogger := &GormLoggerAdapter{logger: appLogger}

	// 根据配置设置日志级别
	switch cfg.Database.LogLevel {
	case "silent":
		gormLogger.LogMode(logger.Silent)
	case "error":
		gormLogger.LogMode(logger.Error)
	case "warn":
		gormLogger.LogMode(logger.Warn)
	case "info":
		gormLogger.LogMode(logger.Info)
	default:
		gormLogger.LogMode(logger.Info)
	}

	// 连接数据库
	db, err := gorm.Open(mysql.Open(dbConfig.DSN), &gorm.Config{
		Logger: gormLogger,
	})
	if err != nil {
		return fmt.Errorf("连接数据库失败: %w", err)
	}

	// 设置连接池
	sqlDB, err := db.DB()
	if err != nil {
		return fmt.Errorf("获取数据库实例失败: %w", err)
	}

	sqlDB.SetMaxIdleConns(cfg.Database.MaxIdleConns)
	sqlDB.SetMaxOpenConns(cfg.Database.MaxOpenConns)

	DB = db
	return nil
}

// AutoMigrate 自动迁移数据库表
func AutoMigrate() error {
	if DB == nil {
		return fmt.Errorf("数据库未初始化")
	}

	// 迁移所有模型
	err := DB.AutoMigrate(
		&models.UserLevelConfig{},
		&models.GoalfyMaxUser{},
		&models.BalanceOperationLog{},
		&models.InviteCodeApplication{},
		//&models.AuditLog{},
		//&models.User{},
		//&models.Role{},
		//&models.Permission{},
		//&models.Menu{},
		// &models.SystemConfig{},
		//&models.LoginLog{},
		//&models.OperationLog{},
		//&models.UserProjectQuota{},
		//&models.PKCEState{},
		//&models.LoginInfo{},
		//&models.User{},
		//&models.Role{},
		//&models.Permission{},
		//&models.Menu{},
		//&models.SystemConfig{},
		//&models.LoginLog{},
		//&models.OperationLog{},
		//&models.UserProjectQuota{},
		//&models.PKCEState{},
		//&models.LoginInfo{},
		// &models.InviteCode{},
	)

	if err != nil {
		return fmt.Errorf("数据库迁移失败: %w", err)
	}

	// 初始化默认用户等级配置
	if err := initDefaultUserLevelConfigs(); err != nil {
		return fmt.Errorf("初始化默认用户等级配置失败: %w", err)
	}

	return nil
}

// initDefaultUserLevelConfigs 初始化默认用户等级配置
func initDefaultUserLevelConfigs() error {
	// 检查是否已存在配置
	var count int64
	if err := DB.Model(&models.UserLevelConfig{}).Count(&count).Error; err != nil {
		return err
	}

	// 如果已存在配置，则跳过
	if count > 0 {
		return nil
	}

	// 创建默认配置
	defaultConfigs := []models.UserLevelConfig{
		{
			LevelName:    "普通",
			LevelCode:    "normal",
			ProjectLimit: 2,
			Description:  "普通用户等级，可创建2个项目",
			SortOrder:    1,
			Status:       1,
		},
		{
			LevelName:    "VIP",
			LevelCode:    "vip",
			ProjectLimit: 10,
			Description:  "VIP用户等级，可创建10个项目",
			SortOrder:    2,
			Status:       1,
		},
		{
			LevelName:    "内部",
			LevelCode:    "internal",
			ProjectLimit: 0,
			Description:  "内部用户等级，无项目数限制",
			SortOrder:    3,
			Status:       1,
		},
	}

	return DB.Create(&defaultConfigs).Error
}

// GetDB 获取数据库实例
func GetDB() *gorm.DB {
	return DB
}

// Close 关闭数据库连接
func Close() error {
	if DB != nil {
		sqlDB, err := DB.DB()
		if err != nil {
			return err
		}
		return sqlDB.Close()
	}
	return nil
}