core/database/new.go

// Package database provides database connection and management functionality
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 holds the tables that need to be auto-migrated on database initialization
	MigrateTables []any
	// Init is an optional initialization function that can be executed after database connection is established
	Init *func() = nil
)

// NewDatabase creates a new database connection based on the provided driver type
// It supports both MySQL and PostgreSQL databases
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
	}

	// Execute the Init function if it's not nil
	if Init != nil {
		(*Init)()
	}

	return db, nil
}

// NewMysql creates a MySQL database service
func NewMysql(dsn []string, options *types.SqlOptions) (gormDb *gorm.DB, err error) {
	//set connection default val.
	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)

	if len(MigrateTables) > 0 {
		err = gormDb.AutoMigrate(MigrateTables...)
		if err != nil {
			return nil, err
		}
	}

	return gormDb, nil
}

// NewPostgres creates a PostgreSQL database service
func NewPostgres(dsn []string, options *types.SqlOptions) (gormDb *gorm.DB, err error) {
	//set connection default val.
	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)

	if len(MigrateTables) > 0 {
		err = gormDb.AutoMigrate(MigrateTables...)
		if err != nil {
			return nil, err
		}
	}

	return
}