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{}, ) }