杭州明婳科技
149 lines
4.5 KiB
Go
149 lines
4.5 KiB
Go
package database
|
||
|
||
import (
|
||
"fmt"
|
||
"strings"
|
||
"sync"
|
||
|
||
"github.com/EthanCodeCraft/xlgo-core/config"
|
||
"github.com/EthanCodeCraft/xlgo-core/logger"
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
// 内置驱动常量（更多驱动可通过 RegisterDialect 扩展）
|
||
const (
|
||
DriverMySQL = config.DriverMySQL
|
||
DriverPostgres = config.DriverPostgres
|
||
)
|
||
|
||
// DialectorFactory 根据 DSN 返回 GORM Dialector
|
||
type DialectorFactory func(dsn string) gorm.Dialector
|
||
|
||
// DialectSpec 描述一种数据库方言：如何建立连接 + 如何拼接 DSN
|
||
type DialectSpec struct {
|
||
// Name 驱动主名称（如 "mysql"、"postgres"、"sqlite"），大小写不敏感
|
||
Name string
|
||
// Aliases 驱动别名（如 postgres 的 "postgresql"、"pg"）
|
||
Aliases []string
|
||
// Dialector 由 DSN 构造 GORM Dialector
|
||
Dialector DialectorFactory
|
||
// DSN 由 DatabaseConfig 拼接连接字符串。可选——
|
||
// 不提供时使用 cfg.MySQLDSN() 兜底（适合自定义驱动通过 CustomDSN 指定连接串的场景）
|
||
DSN config.DSNBuilder
|
||
}
|
||
|
||
var (
|
||
dialectsMu sync.RWMutex
|
||
dialects = map[string]DialectorFactory{}
|
||
)
|
||
|
||
// RegisterDialect 注册一种数据库方言。
|
||
// 同时把 DSN 构建器登记到 config 包，使 cfg.Database.DSN() 也能识别新驱动。
|
||
// 已注册的同名驱动会被覆盖。
|
||
//
|
||
// 用法示例（接入 SQLite）：
|
||
//
|
||
// import "gorm.io/driver/sqlite"
|
||
//
|
||
// database.RegisterDialect(database.DialectSpec{
|
||
// Name: "sqlite",
|
||
// Dialector: func(dsn string) gorm.Dialector { return sqlite.Open(dsn) },
|
||
// DSN: func(c *config.DatabaseConfig) string { return c.Name }, // 文件路径
|
||
// })
|
||
func RegisterDialect(spec DialectSpec) {
|
||
if spec.Dialector == nil || strings.TrimSpace(spec.Name) == "" {
|
||
return
|
||
}
|
||
|
||
dialectsMu.Lock()
|
||
for _, n := range append([]string{spec.Name}, spec.Aliases...) {
|
||
key := normalizeDriver(n)
|
||
if key != "" {
|
||
dialects[key] = spec.Dialector
|
||
}
|
||
}
|
||
dialectsMu.Unlock()
|
||
|
||
if spec.DSN != nil {
|
||
config.RegisterDSNBuilder(spec.Name, spec.DSN, spec.Aliases...)
|
||
}
|
||
}
|
||
|
||
// LookupDialect 查找已注册的 Dialector 工厂
|
||
func LookupDialect(driver string) (DialectorFactory, bool) {
|
||
key := normalizeDriver(driver)
|
||
dialectsMu.RLock()
|
||
defer dialectsMu.RUnlock()
|
||
f, ok := dialects[key]
|
||
return f, ok
|
||
}
|
||
|
||
// RegisteredDialects 返回所有已注册的驱动名（用于诊断）
|
||
func RegisteredDialects() []string {
|
||
dialectsMu.RLock()
|
||
defer dialectsMu.RUnlock()
|
||
names := make([]string, 0, len(dialects))
|
||
for k := range dialects {
|
||
names = append(names, k)
|
||
}
|
||
return names
|
||
}
|
||
|
||
// Dialector 根据配置返回 GORM Dialector。
|
||
// 驱动由 cfg.Database.Driver 决定，未指定或未注册时按 MySQL 兜底（向后兼容）。
|
||
func Dialector(cfg *config.Config) gorm.Dialector {
|
||
if cfg == nil {
|
||
logger.Warn("database: 配置为空，回退到 MySQL 空 DSN")
|
||
return mysql.Open("")
|
||
}
|
||
return dialectorForDSN(cfg.Database.Driver, cfg.Database.DSN())
|
||
}
|
||
|
||
// dialectorForDSN 根据驱动名和 DSN 返回 Dialector
|
||
func dialectorForDSN(driver, dsn string) gorm.Dialector {
|
||
if f, ok := LookupDialect(driver); ok {
|
||
return f(dsn)
|
||
}
|
||
// 未注册时回退到 MySQL，与 config.DSN() 的回退保持一致。
|
||
// 拼写错误的驱动名（如 "mysq"/"postgrs"）会静默回退 MySQL，导致连接错误难排查（M10），
|
||
// 故在此告警一次，提示用户驱动名未注册。
|
||
logger.Warnf("database: 驱动 %q 未注册，回退到 MySQL（已注册: %s）；若是拼写错误请在配置中修正 driver",
|
||
normalizeDriver(driver), strings.Join(RegisteredDialects(), ", "))
|
||
return mysql.Open(dsn)
|
||
}
|
||
|
||
// normalizeDriver 规范化驱动名（小写、去空白）
|
||
func normalizeDriver(name string) string {
|
||
return strings.ToLower(strings.TrimSpace(name))
|
||
}
|
||
|
||
// driverDescription 返回带别名提示的驱动描述（用于错误信息和日志）
|
||
func driverDescription(driver string) string {
|
||
key := normalizeDriver(driver)
|
||
if key == "" {
|
||
return DriverMySQL + " (default)"
|
||
}
|
||
if _, ok := LookupDialect(key); ok {
|
||
return key
|
||
}
|
||
return fmt.Sprintf("%s (unregistered, fallback=%s)", key, DriverMySQL)
|
||
}
|
||
|
||
func init() {
|
||
// 内置 MySQL
|
||
RegisterDialect(DialectSpec{
|
||
Name: DriverMySQL,
|
||
Dialector: func(dsn string) gorm.Dialector { return mysql.Open(dsn) },
|
||
DSN: func(c *config.DatabaseConfig) string { return c.MySQLDSN() },
|
||
})
|
||
// 内置 PostgreSQL
|
||
RegisterDialect(DialectSpec{
|
||
Name: DriverPostgres,
|
||
Aliases: []string{"postgresql", "pg"},
|
||
Dialector: func(dsn string) gorm.Dialector { return postgres.Open(dsn) },
|
||
DSN: func(c *config.DatabaseConfig) string { return c.PostgresDSN() },
|
||
})
|
||
}
|