Files
ops2/backend/ai_work/internal/database/connection.go
T
2026-04-01 12:09:02 +08:00

81 lines
1.6 KiB
Go

package database
import (
"fmt"
"log"
"ops/internal/config"
"time"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
// DB 全局数据库实例
var DB *gorm.DB
// Init 初始化数据库连接
func Init() error {
cfg := config.Current.Database
var dialector gorm.Dialector
switch cfg.Type {
case "sqlite":
dialector = sqlite.Open(cfg.Path)
case "mysql":
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
cfg.User, cfg.Pass, cfg.Host, cfg.Port, cfg.Name)
dialector = mysql.Open(dsn)
case "postgres", "pg":
dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s sslmode=disable TimeZone=Asia/Shanghai",
cfg.Host, cfg.User, cfg.Pass, cfg.Name, cfg.Port)
dialector = postgres.Open(dsn)
default:
return fmt.Errorf("不支持的数据库类型: %s", cfg.Type)
}
// 配置GORM
gormConfig := &gorm.Config{
Logger: logger.Default.LogMode(logger.Info),
}
// 连接数据库
var err error
DB, err = gorm.Open(dialector, gormConfig)
if err != nil {
return fmt.Errorf("数据库连接失败: %v", err)
}
// 配置连接池
sqlDB, err := DB.DB()
if err != nil {
return err
}
sqlDB.SetMaxIdleConns(10)
sqlDB.SetMaxOpenConns(100)
sqlDB.SetConnMaxLifetime(time.Hour)
log.Println("数据库连接成功")
return nil
}
// GetDB 获取数据库实例
func GetDB() *gorm.DB {
return DB
}
// Close 关闭数据库连接
func Close() error {
if DB != nil {
sqlDB, err := DB.DB()
if err != nil {
return err
}
return sqlDB.Close()
}
return nil
}