146 lines
3.8 KiB
Go
146 lines
3.8 KiB
Go
// Package database 提供数据库连接和管理功能
|
||
// 支持MySQL和PostgreSQL数据库，包含连接池管理和自动迁移
|
||
package database
|
||
|
||
import (
|
||
"fmt"
|
||
"strings"
|
||
|
||
"git.apinb.com/bsm-sdk/core/database/sql"
|
||
"git.apinb.com/bsm-sdk/core/types"
|
||
"git.apinb.com/bsm-sdk/core/vars"
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
var (
|
||
// MigrateTables 存储需要在数据库初始化时自动迁移的表
|
||
MigrateTables []any
|
||
)
|
||
|
||
// NewDatabase 根据提供的驱动类型创建新的数据库连接
|
||
// 支持MySQL和PostgreSQL数据库
|
||
// driver: 数据库驱动类型 ("mysql" 或 "postgres")
|
||
// dsn: 数据源名称数组
|
||
// options: 数据库连接选项
|
||
// 返回: GORM数据库实例
|
||
func NewDatabase(driver string, dsn []string, options *types.SqlOptions) (db *gorm.DB, err error) {
|
||
driver = strings.ToLower(driver)
|
||
switch driver {
|
||
case "mysql":
|
||
db, err = NewMysql(dsn, options)
|
||
case "postgres":
|
||
db, err = NewPostgres(dsn, options)
|
||
default:
|
||
return nil, fmt.Errorf("unsupported database driver: %s", driver)
|
||
}
|
||
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 自动迁移表结构
|
||
if len(MigrateTables) > 0 {
|
||
err = db.AutoMigrate(MigrateTables...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
|
||
return db, nil
|
||
}
|
||
|
||
// NewMysql 创建MySQL数据库服务
|
||
// dsn: 数据源名称数组
|
||
// options: 数据库连接选项
|
||
// 返回: GORM数据库实例
|
||
func NewMysql(dsn []string, options *types.SqlOptions) (gormDb *gorm.DB, err error) {
|
||
// 设置连接默认值
|
||
if options == nil {
|
||
options = &types.SqlOptions{
|
||
MaxIdleConns: vars.SqlOptionMaxIdleConns,
|
||
MaxOpenConns: vars.SqlOptionMaxOpenConns,
|
||
ConnMaxLifetime: vars.SqlOptionConnMaxLifetime,
|
||
LogStdout: false,
|
||
Debug: true,
|
||
}
|
||
}
|
||
|
||
gormDb, err = gorm.Open(mysql.Open(dsn[0]), &gorm.Config{
|
||
SkipDefaultTransaction: true,
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if options.Debug {
|
||
gormDb = gormDb.Debug()
|
||
}
|
||
|
||
// 获取通用数据库对象 sql.DB，然后使用其提供的功能
|
||
sqlDB, err := gormDb.DB()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// SetMaxIdleConns 用于设置连接池中空闲连接的最大数量
|
||
sqlDB.SetMaxIdleConns(options.MaxIdleConns)
|
||
// SetMaxOpenConns 设置打开数据库连接的最大数量
|
||
sqlDB.SetMaxOpenConns(options.MaxOpenConns)
|
||
// SetConnMaxLifetime 设置了连接可复用的最大时间
|
||
sqlDB.SetConnMaxLifetime(options.ConnMaxLifetime)
|
||
|
||
return gormDb, nil
|
||
}
|
||
|
||
// NewPostgres 创建PostgreSQL数据库服务
|
||
// dsn: 数据源名称数组
|
||
// options: 数据库连接选项
|
||
// 返回: GORM数据库实例
|
||
func NewPostgres(dsn []string, options *types.SqlOptions) (gormDb *gorm.DB, err error) {
|
||
// 设置连接默认值
|
||
if options == nil {
|
||
options = &types.SqlOptions{
|
||
MaxIdleConns: vars.SqlOptionMaxIdleConns,
|
||
MaxOpenConns: vars.SqlOptionMaxOpenConns,
|
||
ConnMaxLifetime: vars.SqlOptionConnMaxLifetime,
|
||
LogStdout: false,
|
||
Debug: true,
|
||
}
|
||
}
|
||
|
||
gormDb, err = sql.NewPostgreSql(dsn[0], options)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if options.Debug {
|
||
gormDb = gormDb.Debug()
|
||
}
|
||
|
||
// 获取通用数据库对象 sql.DB，然后使用其提供的功能
|
||
sqlDB, err := gormDb.DB()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// SetMaxIdleConns 用于设置连接池中空闲连接的最大数量
|
||
sqlDB.SetMaxIdleConns(options.MaxIdleConns)
|
||
// SetMaxOpenConns 设置打开数据库连接的最大数量
|
||
sqlDB.SetMaxOpenConns(options.MaxOpenConns)
|
||
// SetConnMaxLifetime 设置了连接可复用的最大时间
|
||
sqlDB.SetConnMaxLifetime(options.ConnMaxLifetime)
|
||
|
||
return gormDb, nil
|
||
}
|
||
|
||
// AppendMigrate 调用此函数后，会在数据库初始化时自动迁移表结构
|
||
//
|
||
// - table: 需要自动迁移的表
|
||
func AppendMigrate(table any) {
|
||
if MigrateTables == nil {
|
||
MigrateTables = make([]any, 0)
|
||
}
|
||
MigrateTables = append(MigrateTables, table)
|
||
}
|