75 lines
1.6 KiB
Go
75 lines
1.6 KiB
Go
package dal
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
"gorm.io/plugin/dbresolver"
|
|
|
|
"gitea.micah.com/micah/standard/conf"
|
|
"gitea.micah.com/micah/standard/internal/dal/query"
|
|
)
|
|
|
|
var DB *gorm.DB
|
|
|
|
func Init() {
|
|
c := conf.GetConf()
|
|
var err error
|
|
loggerLevel := getLogLevel(c.MySQL.LogLevel)
|
|
if loggerLevel < logger.Silent || loggerLevel > logger.Info {
|
|
loggerLevel = logger.Info
|
|
}
|
|
|
|
DB, err = gorm.Open(mysql.Open(c.MySQL.DSN), &gorm.Config{
|
|
SkipDefaultTransaction: true,
|
|
PrepareStmt: true,
|
|
Logger: logger.Default.LogMode(loggerLevel),
|
|
})
|
|
if err != nil {
|
|
panic(fmt.Sprintf("注册 DB DSN 失败, %+v", err))
|
|
}
|
|
if len(c.MySQL.ReplicasDSN) > 0 {
|
|
dialects := make([]gorm.Dialector, len(c.MySQL.ReplicasDSN))
|
|
for i, s := range c.MySQL.ReplicasDSN {
|
|
dialects[i] = mysql.Open(s)
|
|
}
|
|
// 注册 DB Resolver 插件
|
|
err = DB.Use(dbresolver.Register(dbresolver.Config{
|
|
// 主库（写操作）
|
|
Sources: []gorm.Dialector{mysql.Open(c.MySQL.DSN)},
|
|
// 从库（读操作）
|
|
Replicas: dialects,
|
|
// 负载均衡策略（默认轮询）
|
|
Policy: dbresolver.RandomPolicy{}, // 可选：随机、轮询、加权等
|
|
// 自动检查连接健康
|
|
TraceResolverMode: true, // 记录解析器日志
|
|
}).SetMaxOpenConns(100))
|
|
|
|
if err != nil {
|
|
panic(fmt.Sprintf("注册 DB Resolver 失败, %+v", err))
|
|
}
|
|
}
|
|
|
|
query.SetDefault(DB)
|
|
}
|
|
|
|
func getLogLevel(level string) logger.LogLevel {
|
|
switch level {
|
|
case "info":
|
|
return logger.Info
|
|
case "warn":
|
|
return logger.Warn
|
|
case "error":
|
|
return logger.Error
|
|
default:
|
|
return logger.Silent
|
|
|
|
}
|
|
}
|
|
|
|
func GetDB() *gorm.DB {
|
|
return DB
|
|
}
|