core/database/new.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 = make([]any, 0)
)

// 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) {
	MigrateTables = append(MigrateTables, table)
}