From 86bdaae472e404d14877786f69a23392a7a572aa Mon Sep 17 00:00:00 2001
From: liuyuecai <liuyuecai@360.cn>
Date: Sat, 23 Mar 2024 18:21:26 +0800
Subject: [PATCH] fix trx log

---
 pkg/runtime/transaction/hook.go | 25 +++++++++++++------------
 pkg/runtime/tx_test.go          |  5 +++--
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/pkg/runtime/transaction/hook.go b/pkg/runtime/transaction/hook.go
index 4ced1554..c7a3f0f0 100644
--- a/pkg/runtime/transaction/hook.go
+++ b/pkg/runtime/transaction/hook.go
@@ -19,6 +19,7 @@ package transaction
 
 import (
 	"context"
+	rcontext "github.com/arana-db/arana/pkg/runtime/context"
 )
 
 import (
@@ -40,15 +41,15 @@ func NewXAHook(tenant string, enable bool) (*xaHook, error) {
 		enable: enable,
 	}
 
-	trxStateChangeFunc := map[runtime.TxState]handleFunc{
-		runtime.TrxStarted:       xh.onStarted,
-		runtime.TrxPreparing:     xh.onPreparing,
-		runtime.TrxPrepared:      xh.onPrepared,
-		runtime.TrxCommitting:    xh.onCommitting,
-		runtime.TrxCommitted:     xh.onCommitted,
-		runtime.TrxAborted:       xh.onAborting,
-		runtime.TrxRolledBacking: xh.onRollbackOnly,
-		runtime.TrxRolledBacked:  xh.onRolledBack,
+	trxStateChangeFunc := map[rcontext.TxState]handleFunc{
+		rcontext.TrxStarted:       xh.onStarted,
+		rcontext.TrxPreparing:     xh.onPreparing,
+		rcontext.TrxPrepared:      xh.onPrepared,
+		rcontext.TrxCommitting:    xh.onCommitting,
+		rcontext.TrxCommitted:     xh.onCommitted,
+		rcontext.TrxAborted:       xh.onAborting,
+		rcontext.TrxRolledBacking: xh.onRollbackOnly,
+		rcontext.TrxRolledBacked:  xh.onRolledBack,
 	}
 
 	xh.trxMgr = trxMgr
@@ -64,10 +65,10 @@ type xaHook struct {
 	enable             bool
 	trxMgr             *TrxManager
 	trxLog             *GlobalTrxLog
-	trxStateChangeFunc map[runtime.TxState]handleFunc
+	trxStateChangeFunc map[rcontext.TxState]handleFunc
 }
 
-func (xh *xaHook) OnTxStateChange(ctx context.Context, state runtime.TxState, tx runtime.CompositeTx) error {
+func (xh *xaHook) OnTxStateChange(ctx context.Context, state rcontext.TxState, tx runtime.CompositeTx) error {
 	if !xh.enable {
 		return nil
 	}
@@ -162,7 +163,7 @@ func (xh *xaHook) onRollbackOnly(ctx context.Context, tx runtime.CompositeTx) er
 }
 
 func (xh *xaHook) onRolledBack(ctx context.Context, tx runtime.CompositeTx) error {
-	xh.trxLog.Status = runtime.TrxRolledBacking
+	xh.trxLog.Status = rcontext.TrxRolledBacking
 	if err := xh.trxMgr.trxLog.AddOrUpdateGlobalTxLog(*xh.trxLog); err != nil {
 		return err
 	}
diff --git a/pkg/runtime/tx_test.go b/pkg/runtime/tx_test.go
index 19832e1e..0050cde2 100644
--- a/pkg/runtime/tx_test.go
+++ b/pkg/runtime/tx_test.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	rcontext "github.com/arana-db/arana/pkg/runtime/context"
 	"testing"
 	"time"
 )
@@ -41,7 +42,7 @@ func Test_branchTx_CallFieldList(t *testing.T) {
 	type fields struct {
 		closed   atomic.Bool
 		parent   *AtomDB
-		state    TxState
+		state    rcontext.TxState
 		prepare  dbFunc
 		commit   dbFunc
 		rollback dbFunc
@@ -89,7 +90,7 @@ func Test_compositeTx_Rollback(t *testing.T) {
 		beginTime time.Time
 		endTime   time.Time
 		isoLevel  sql.IsolationLevel
-		txState   TxState
+		txState   rcontext.TxState
 		beginFunc dbFunc
 		rt        *defaultRuntime
 		txs       map[string]*branchTx