120 lines
2.8 KiB
Go
120 lines
2.8 KiB
Go
// Package database 提供数据库连接和管理功能
|
||
// 支持MySQL和PostgreSQL数据库，包含连接池管理和自动迁移
|
||
package database
|
||
|
||
import (
|
||
"fmt"
|
||
"strings"
|
||
|
||
dbsql "git.apinb.com/bsm-sdk/core/database/sql"
|
||
"git.apinb.com/bsm-sdk/core/types"
|
||
"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
|
||
}
|
||
|
||
func firstDSN(dsn []string) (string, error) {
|
||
if len(dsn) == 0 || strings.TrimSpace(dsn[0]) == "" {
|
||
return "", fmt.Errorf("database dsn is empty")
|
||
}
|
||
return dsn[0], nil
|
||
}
|
||
|
||
func applySqlOptions(gormDb *gorm.DB, options *types.SqlOptions) (*gorm.DB, error) {
|
||
if options.Debug {
|
||
gormDb = gormDb.Debug()
|
||
}
|
||
|
||
sqlDB, err := gormDb.DB()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
sqlDB.SetMaxIdleConns(options.MaxIdleConns)
|
||
sqlDB.SetMaxOpenConns(options.MaxOpenConns)
|
||
sqlDB.SetConnMaxLifetime(options.ConnMaxLifetime)
|
||
|
||
return gormDb, nil
|
||
}
|
||
|
||
// NewMysql 创建MySQL数据库服务
|
||
// dsn: 数据源名称数组
|
||
// options: 数据库连接选项
|
||
// 返回: GORM数据库实例
|
||
func NewMysql(dsn []string, options *types.SqlOptions) (gormDb *gorm.DB, err error) {
|
||
options = dbsql.SetOptions(options)
|
||
|
||
dsn0, err := firstDSN(dsn)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
gormDb, err = gorm.Open(mysql.Open(dsn0), &gorm.Config{
|
||
SkipDefaultTransaction: true,
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return applySqlOptions(gormDb, options)
|
||
}
|
||
|
||
// NewPostgres 创建PostgreSQL数据库服务
|
||
// dsn: 数据源名称数组
|
||
// options: 数据库连接选项
|
||
// 返回: GORM数据库实例
|
||
func NewPostgres(dsn []string, options *types.SqlOptions) (gormDb *gorm.DB, err error) {
|
||
dsn0, err := firstDSN(dsn)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return dbsql.NewPostgreSql(dsn0, options)
|
||
}
|
||
|
||
// AppendMigrate 调用此函数后，会在数据库初始化时自动迁移表结构
|
||
//
|
||
// - table: 需要自动迁移的表
|
||
func AppendMigrate(table any) {
|
||
if MigrateTables == nil {
|
||
MigrateTables = make([]any, 0)
|
||
}
|
||
MigrateTables = append(MigrateTables, table)
|
||
}
|