144 lines
3.6 KiB
Go
144 lines
3.6 KiB
Go
package db
|
||
|
||
import (
|
||
"fmt"
|
||
"log"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/glebarez/sqlite"
|
||
"github.com/ssdomei232/goodBaby/configs"
|
||
"github.com/ssdomei232/goodBaby/model"
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
)
|
||
|
||
var (
|
||
once sync.Once
|
||
instance *gorm.DB
|
||
initErr error
|
||
)
|
||
|
||
// GetGormDB 返回全局唯一的 gorm 连接。
|
||
//
|
||
// 之前每次调用都会新开一个 sqlite 连接，在并发执行规则时很容易触发
|
||
// "database is locked"，这里改为单例并开启 WAL。
|
||
func GetGormDB() (*gorm.DB, error) {
|
||
once.Do(func() {
|
||
instance, initErr = open()
|
||
})
|
||
return instance, initErr
|
||
}
|
||
|
||
// MustInit 在启动阶段初始化数据库并执行迁移，失败直接退出
|
||
func MustInit() *gorm.DB {
|
||
gormDB, err := GetGormDB()
|
||
if err != nil {
|
||
log.Fatalf("初始化数据库失败: %v", err)
|
||
}
|
||
if err := AutoMigrate(gormDB); err != nil {
|
||
log.Fatalf("数据库迁移失败: %v", err)
|
||
}
|
||
if err := ensureAdmin(gormDB); err != nil {
|
||
log.Fatalf("初始化管理员失败: %v", err)
|
||
}
|
||
return gormDB
|
||
}
|
||
|
||
// ensureAdmin 保证系统里至少有一个管理员。
|
||
//
|
||
// IsAdmin 是后加的字段，老数据库升级上来时所有人都不是管理员，
|
||
// 这里把最早注册的那个用户提升为管理员。
|
||
func ensureAdmin(gormDB *gorm.DB) error {
|
||
var adminCount int64
|
||
if err := gormDB.Model(&model.User{}).Where("is_admin = ?", true).Count(&adminCount).Error; err != nil {
|
||
return err
|
||
}
|
||
if adminCount > 0 {
|
||
return nil
|
||
}
|
||
|
||
var first model.User
|
||
if err := gormDB.Order("id ASC").First(&first).Error; err != nil {
|
||
// 一个用户都没有，等第一个注册的人来当管理员
|
||
return nil
|
||
}
|
||
|
||
if err := gormDB.Model(&model.User{}).Where("id = ?", first.ID).
|
||
Update("is_admin", true).Error; err != nil {
|
||
return err
|
||
}
|
||
log.Printf("已将用户 %s (ID: %d) 提升为管理员", first.Username, first.ID)
|
||
return nil
|
||
}
|
||
|
||
func open() (*gorm.DB, error) {
|
||
config, err := configs.GetConfig()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("读取配置失败: %w", err)
|
||
}
|
||
|
||
gormConfig := &gorm.Config{Logger: logger.Default.LogMode(logger.Warn)}
|
||
|
||
var (
|
||
gormDB *gorm.DB
|
||
driver = config.DatabaseDriver
|
||
)
|
||
|
||
switch driver {
|
||
case configs.DriverPostgres:
|
||
if config.DatabaseDSN == "" {
|
||
return nil, fmt.Errorf("使用 postgres 时必须配置 database_dsn")
|
||
}
|
||
gormDB, err = gorm.Open(postgres.Open(config.DatabaseDSN), gormConfig)
|
||
case configs.DriverSQLite, "":
|
||
driver = configs.DriverSQLite
|
||
dsn := fmt.Sprintf(
|
||
"file:%s?_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)",
|
||
config.DatabasePath,
|
||
)
|
||
gormDB, err = gorm.Open(sqlite.Open(dsn), gormConfig)
|
||
default:
|
||
return nil, fmt.Errorf("不支持的数据库驱动: %s（可选 sqlite / postgres）", config.DatabaseDriver)
|
||
}
|
||
if err != nil {
|
||
return nil, fmt.Errorf("打开数据库失败: %w", err)
|
||
}
|
||
|
||
sqlDB, err := gormDB.DB()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if driver == configs.DriverSQLite {
|
||
// sqlite 单写多读，限制连接数避免写冲突
|
||
sqlDB.SetMaxOpenConns(1)
|
||
sqlDB.SetMaxIdleConns(1)
|
||
} else {
|
||
sqlDB.SetMaxOpenConns(20)
|
||
sqlDB.SetMaxIdleConns(5)
|
||
sqlDB.SetConnMaxLifetime(time.Hour)
|
||
}
|
||
|
||
// 提前握手，配置写错时在启动阶段就报错，而不是等到第一次查询
|
||
if err := sqlDB.Ping(); err != nil {
|
||
return nil, fmt.Errorf("连接数据库失败: %w", err)
|
||
}
|
||
|
||
log.Printf("数据库已连接 (driver=%s)", driver)
|
||
return gormDB, nil
|
||
}
|
||
|
||
// AutoMigrate 建表 / 补齐新增字段
|
||
func AutoMigrate(gormDB *gorm.DB) error {
|
||
return gormDB.AutoMigrate(
|
||
&model.User{},
|
||
&model.Timer{},
|
||
&model.Rule{},
|
||
&model.Account{},
|
||
&model.ExecutionLog{},
|
||
&model.MessageGateway{},
|
||
)
|
||
}
|