core/database/new.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)
}