Snowykami
b0d224dc64
- Introduced `hertzx` package with `NewHertz` function for server initialization. - Implemented `logx` package with various log collectors: Console, Loki, Elasticsearch, and Prometheus. - Added `Logger` struct to manage logging levels and collectors. - Created environment variable loading functionality in `osx` package to support configuration. - Enhanced logging capabilities with structured log entries and asynchronous collection.
72 lines
2.2 KiB
Go
72 lines
2.2 KiB
Go
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 + "`"
|
||
}
|