package dbx import ( "context" "fmt" "strings" "time" "gorm.io/driver/mysql" "gorm.io/gorm" ) func GetMySQLDB(cfg *DBConfig, dbName string) *gorm.DB { db, err := getMySQLInstance(cfg, dbName) if err != nil { panic("failed to connect to MySQL database: " + err.Error()) } return db } func getMySQLInstance(cfg *DBConfig, dbName string) (*gorm.DB, error) { targetDSN := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local", cfg.User, cfg.Password, cfg.Host, cfg.Port, dbName) // 先尝试直接连接目标库 if db, err := gorm.Open(mysql.Open(targetDSN), &gorm.Config{}); err == nil { return db, nil } // 以不指定数据库的 DSN 连接(用于创建数据库) adminDSN := fmt.Sprintf("%s:%s@tcp(%s:%s)/?charset=utf8mb4&parseTime=True&loc=Local", cfg.User, cfg.Password, cfg.Host, cfg.Port) adminDB, err := gorm.Open(mysql.Open(adminDSN), &gorm.Config{}) if err != nil { return nil, fmt.Errorf("connect admin mysql failed: %w", err) } // 关闭底层连接 if sqlDB, e := adminDB.DB(); e == nil { defer sqlDB.Close() } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() // 检查是否存在(使用 INFORMATION_SCHEMA) var count int64 checkSQL := "SELECT count(*) FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = ?" if err := adminDB.WithContext(ctx).Raw(checkSQL, dbName).Scan(&count).Error; err != nil { return nil, fmt.Errorf("check mysql database existence failed: %w", err) } if count == 0 { ident := escapeMySQLIdentifier(dbName) createSQL := fmt.Sprintf("CREATE DATABASE %s DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci", ident) if err := adminDB.WithContext(ctx).Exec(createSQL).Error; err != nil { return nil, fmt.Errorf("create mysql database %s failed: %w", dbName, err) } } // 重试连接目标数据库 db2, err := gorm.Open(mysql.Open(targetDSN), &gorm.Config{}) if err != nil { return nil, fmt.Errorf("connect mysql target after create failed: %w", err) } return db2, nil } // 辅助:安全转义 MySQL 标识符(用反引号并把 ` 替换为 “) func escapeMySQLIdentifier(s string) string { s = strings.ReplaceAll(s, "`", "``") return "`" + s + "`" }