baize/framework/core/infrastructure/database/operations/operations.go

package operations

import (
	dbSql "database/sql"
	"git.sxidc.com/go-framework/baize/framework/core/infrastructure/database/sql"
	"git.sxidc.com/go-framework/baize/framework/core/infrastructure/logger"
	"git.sxidc.com/go-tools/utils/template"
	"github.com/pkg/errors"
	"gorm.io/gorm"
	"time"
)

type Config struct {
	UserName           string `json:"user_name" yaml:"user_name"`
	Password           string `json:"password" yaml:"password"`
	Address            string `json:"address" yaml:"address"`
	Port               string `json:"port" yaml:"port"`
	Database           string `json:"database" yaml:"database"`
	MaxConnections     int    `json:"max_connections" yaml:"max_connections"`
	MaxIdleConnections int    `json:"max_idle_connections" yaml:"max_idle_connections"`
	LogLevel           string `json:"log_level" yaml:"log_level"`
}

type Operations struct {
	db           *gorm.DB
	stopPingChan chan any
}

func NewOperations(dbConfig *Config) (*Operations, error) {
	gormDB, err := newGormDB(dbConfig)
	if err != nil {
		return nil, err
	}

	sqlDB, err := gormDB.DB()
	if err != nil {
		return nil, errors.New(err.Error())
	}

	if dbConfig.MaxConnections == 0 {
		dbConfig.MaxConnections = 50
	}

	if dbConfig.MaxIdleConnections == 0 {
		dbConfig.MaxIdleConnections = 10
	}

	sqlDB.SetMaxOpenConns(dbConfig.MaxConnections)
	sqlDB.SetMaxIdleConns(dbConfig.MaxIdleConnections)

	op := &Operations{
		db:           gormDB,
		stopPingChan: make(chan any),
	}

	op.startBeatHeart(sqlDB)

	return op, nil
}

func DestroyOperation(op *Operations) error {
	if op == nil {
		return nil
	}

	if op.db == nil {
		return nil
	}

	op.stopBeatHeart()

	err := destroyGormDB(op.db)
	if err != nil {
		return err
	}

	op = nil

	return nil
}

func (op *Operations) startBeatHeart(sqlDB *dbSql.DB) {
	go func() {
		pingTicker := time.NewTicker(time.Minute * 1)
		defer pingTicker.Stop()

		for {
			select {
			case <-op.stopPingChan:
				return
			case <-pingTicker.C:
				err := sqlDB.Ping()
				if err != nil {
					logger.GetInstance().Error(err)
				}
			}
		}
	}()
}

func (op *Operations) stopBeatHeart() {
	if op.stopPingChan != nil {
		op.stopPingChan <- nil
		close(op.stopPingChan)
		op.stopPingChan = nil
	}
}

func (op *Operations) BeginTransaction() *Operations {
	tx := op.db.Begin()
	return &Operations{
		db: tx,
	}
}

func (op *Operations) RollbackTransaction() {
	op.db.Rollback()
}

func (op *Operations) CommitTransaction() {
	op.db.Commit()
}

func (op *Operations) AutoMigrate(tables ...Table) error {
	tx := op.db.Begin()

	for _, table := range tables {
		dbModel, err := table.ToDBModel()
		if err != nil {
			tx.Rollback()
			return err
		}

		err = tx.Table(table.TableName).AutoMigrate(dbModel)
		if err != nil {
			tx.Rollback()
			return errors.New(err.Error())
		}
	}

	tx.Commit()

	return nil
}

func (op *Operations) ExecuteRawSql(sqlStr string, args ...any) ([]sql.Result, error) {
	return op.ExecuteRawSqlTemplate(sqlStr, nil, args...)
}

func (op *Operations) ExecuteRawSqlTemplate(sqlStr string, executeParams map[string]any, args ...any) ([]sql.Result, error) {
	parsedSql, err := template.ParseTemplateStringToString(sqlStr, executeParams)
	if err != nil {
		return nil, errors.New(err.Error())
	}

	tableRows := make([]map[string]any, 0)
	err = op.db.Raw(parsedSql, args...).Scan(&tableRows).Error
	if err != nil {
		return nil, errors.New(err.Error())
	}

	results := make([]sql.Result, len(tableRows))
	for i, row := range tableRows {
		results[i] = row
	}

	return results, nil
}