package database import ( "fmt" "log" "anticaptcha/internal/config" "anticaptcha/internal/model" "gorm.io/driver/mysql" "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/logger" ) var DB *gorm.DB func Init(cfg *config.DatabaseConfig) error { var err error var gormConfig *gorm.Config // 生产环境禁用日志 gormConfig = &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), } switch cfg.Type { case "sqlite": DB, err = gorm.Open(sqlite.Open(cfg.SQLite.Path), gormConfig) case "mysql": dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local", cfg.User, cfg.Password, cfg.Host, cfg.Port, cfg.Database) DB, err = gorm.Open(mysql.Open(dsn), gormConfig) default: return fmt.Errorf("不支持的数据库类型: %s", cfg.Type) } if err != nil { return fmt.Errorf("数据库连接失败: %v", err) } // 自动迁移 if err = DB.AutoMigrate( &model.User{}, &model.RegistrationCode{}, &model.EndpointCost{}, &model.Config{}, ); err != nil { return fmt.Errorf("数据库迁移失败: %v", err) } return nil } func Close() error { if DB != nil { sqlDB, err := DB.DB() if err != nil { return err } return sqlDB.Close() } return nil } // 确保 admin 用户存在 func EnsureAdminUser() error { var count int64 DB.Model(&model.User{}).Where("role = ?", model.RoleAdmin).Count(&count) if count > 0 { return nil } admin := &model.User{ Username: "admin", HashedPassword: HashPassword("admin"), Role: model.RoleAdmin, Balance: 1000000, } return DB.Create(admin).Error } func HashPassword(password string) string { // 简化处理,实际应使用 bcrypt // 后续补充完整实现 return password }