110 lines
2.1 KiB
Go
110 lines
2.1 KiB
Go
package model
|
||
|
||
import (
|
||
"database/sql"
|
||
"fmt"
|
||
"github.com/gofiber/fiber/v2/log"
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/driver/sqlite"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
|
||
"gitee.ltd/lxh/wechat-robot/internal/config"
|
||
)
|
||
|
||
var (
|
||
db *gorm.DB
|
||
)
|
||
|
||
// InitDB 初始化数据库连接
|
||
func InitDB() {
|
||
var err error
|
||
defer func() {
|
||
if err != nil {
|
||
log.Panicf("数据库初始化失败: %v", err)
|
||
}
|
||
}()
|
||
|
||
gormConfig := &gorm.Config{
|
||
Logger: logger.Default.LogMode(logger.Info),
|
||
}
|
||
|
||
dsn := config.Scd.Database.DSN()
|
||
|
||
// 根据数据库类型选择对应的驱动
|
||
switch config.Scd.Database.Type {
|
||
case config.PostgreSQL:
|
||
db, err = gorm.Open(postgres.Open(dsn), gormConfig)
|
||
case config.MySQL:
|
||
db, err = gorm.Open(mysql.Open(dsn), gormConfig)
|
||
case config.SQLite:
|
||
db, err = gorm.Open(sqlite.Open(dsn), gormConfig)
|
||
default:
|
||
err = fmt.Errorf("不支持的数据库类型: %s", config.Scd.Database.Type)
|
||
return
|
||
}
|
||
|
||
if err != nil {
|
||
log.Fatalf("无法连接到数据库: %v", err)
|
||
return
|
||
}
|
||
|
||
// 自动迁移数据库模型
|
||
err = migrateDB()
|
||
if err != nil {
|
||
log.Fatalf("迁移数据库失败: %v", err)
|
||
return
|
||
}
|
||
|
||
// 对于SQLite，执行一些特定的优化
|
||
if config.Scd.Database.Type == config.SQLite {
|
||
var sqlDB *sql.DB
|
||
sqlDB, err = db.DB()
|
||
if err != nil {
|
||
log.Errorf("无法获取基础SQL DB: %v", err)
|
||
return
|
||
}
|
||
|
||
// 启用外键约束
|
||
_, _ = sqlDB.Exec("PRAGMA foreign_keys = ON")
|
||
// 设置连接池大小
|
||
sqlDB.SetMaxOpenConns(1) // SQLite建议使用单连接
|
||
}
|
||
|
||
return
|
||
}
|
||
|
||
// GetDB 获取数据库连接实例
|
||
func GetDB() *gorm.DB {
|
||
if db == nil {
|
||
panic("数据库未初始化，先调用InitDB")
|
||
}
|
||
return db
|
||
}
|
||
|
||
// CloseDB 关闭数据库连接
|
||
func CloseDB() error {
|
||
if db == nil {
|
||
return nil
|
||
}
|
||
|
||
sqlDB, err := db.DB()
|
||
if err != nil {
|
||
return fmt.Errorf("获取sql.DB实例错误: %w", err)
|
||
}
|
||
|
||
return sqlDB.Close()
|
||
}
|
||
|
||
// migrateDB 进行数据库迁移
|
||
func migrateDB() error {
|
||
// 在这里添加需要自动迁移的模型
|
||
return db.AutoMigrate(
|
||
&Robot{},
|
||
&Contact{},
|
||
&GroupMember{},
|
||
&Message{},
|
||
)
|
||
}
|