hotime/db/transaction.go

package db

import (
	"fmt"
	"sync"
	"sync/atomic"
)

var savepointCounter uint64

// Action 事务操作
// 如果 action 返回 true 则提交事务；返回 false 则回滚
// 测试模式下（testTx != nil），使用 SAVEPOINT 代替真事务，确保嵌套事务不会绕过外层测试事务
func (that *HoTimeDB) Action(action func(db HoTimeDB) (isSuccess bool)) (isSuccess bool) {
	that.testLogBufMu.Lock()
	logBuf := that.testLogBuf
	that.testLogBufMu.Unlock()
	db := HoTimeDB{
		DB:          that.DB,
		ContextBase: that.ContextBase,
		DBName:      that.DBName,
		HoTimeCache: that.HoTimeCache,
		Log:         that.Log,
		Type:        that.Type,
		Prefix:      that.Prefix,
		ConnectFunc: that.ConnectFunc,
		limit:       that.limit,
		Tx:          that.Tx,
		SlaveDB:     that.SlaveDB,
		Dialect:     that.Dialect,
		limitMu:     sync.Mutex{},
		testTx:      that.testTx,
		testMu:      that.testMu,
		testLogBuf:  logBuf,
	}

	txFailed := false
	db.txFailed = &txFailed

	if that.testTx != nil {
		spName := fmt.Sprintf("sp_%d", atomic.AddUint64(&savepointCounter, 1))
		if that.testMu != nil {
			that.testMu.Lock()
		}
		_, _ = that.testTx.Exec("SAVEPOINT " + spName)
		if that.testMu != nil {
			that.testMu.Unlock()
		}
		db.Tx = that.testTx
		isSuccess = action(db)
		if txFailed || !isSuccess {
			if that.testMu != nil {
				that.testMu.Lock()
			}
			_, _ = that.testTx.Exec("ROLLBACK TO SAVEPOINT " + spName)
			if that.testMu != nil {
				that.testMu.Unlock()
			}
			return false
		}
		if that.testMu != nil {
			that.testMu.Lock()
		}
		_, _ = that.testTx.Exec("RELEASE SAVEPOINT " + spName)
		if that.testMu != nil {
			that.testMu.Unlock()
		}
		return isSuccess
	}

	tx, err := db.Begin()

	if err != nil {
		dbErr := &DBError{}
		dbErr.SetError(err)
		that.lastError.Store(dbErr)
		return isSuccess
	}

	db.Tx = tx

	isSuccess = action(db)

	if txFailed {
		_ = db.Tx.Rollback()
		return false
	}

	if !isSuccess {
		err = db.Tx.Rollback()
		if err != nil {
			dbErr := &DBError{}
			dbErr.SetError(err)
			that.lastError.Store(dbErr)
			return isSuccess
		}
		return isSuccess
	}

	err = db.Tx.Commit()
	if err != nil {
		dbErr := &DBError{}
		dbErr.SetError(err)
		that.lastError.Store(dbErr)
		return false
	}
	return true
}