baize/infrastructure/database/operations/operations.go

package operations

import (
	dbSql "database/sql"
	"encoding/json"
	"git.sxidc.com/go-framework/baize/infrastructure/database/sql"
	"git.sxidc.com/go-tools/utils/template"
	"git.sxidc.com/service-supports/fslog"
	"gorm.io/gorm"
	"time"
)

type Config struct {
	UserName           string
	Password           string
	Address            string
	Port               string
	Database           string
	MaxConnections     int
	MaxIdleConnections int
}

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, err
	}

	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 {
					fslog.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 err
		}
	}

	tx.Commit()

	return nil
}

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

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

	// 简化一下类型体系，与DataService保持一致
	jsonTableRows, err := json.Marshal(tableRows)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(jsonTableRows, &tableRows)
	if err != nil {
		return nil, err
	}

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

	return results, nil
}

func (op *Operations) ExecuteSql(_ string, _ map[string]any) ([]sql.Result, error) {
	return make([]sql.Result, 0), nil
}