From 6372611d647aad088b8298f8a72af0a65fb0b17e Mon Sep 17 00:00:00 2001
From: chenquan <chenquan.dev@gmail.com>
Date: Sun, 8 Jan 2023 19:04:55 +0800
Subject: [PATCH] refactor: refactor hook

---
 hook.go      |  3 +++
 hook_test.go | 36 ------------------------------------
 tx.go        | 16 ++--------------
 3 files changed, 5 insertions(+), 50 deletions(-)

diff --git a/hook.go b/hook.go
index 18cc6d4..e370e57 100644
--- a/hook.go
+++ b/hook.go
@@ -24,6 +24,7 @@ type (
 
 		BeforePrepareContext(ctx context.Context, query string, err error) (context.Context, string, error)
 		AfterPrepareContext(ctx context.Context, query string, ds driver.Stmt, err error) (context.Context, driver.Stmt, error)
+
 		BeforeClose(ctx context.Context, err error) (context.Context, error)
 		AfterClose(ctx context.Context, err error) (context.Context, error)
 	}
@@ -34,12 +35,14 @@ type (
 	TxHook interface {
 		BeforeCommit(ctx context.Context, err error) (context.Context, error)
 		AfterCommit(ctx context.Context, err error) (context.Context, error)
+
 		BeforeRollback(ctx context.Context, err error) (context.Context, error)
 		AfterRollback(ctx context.Context, err error) (context.Context, error)
 	}
 	StmtHook interface {
 		BeforeStmtQueryContext(ctx context.Context, query string, args []driver.NamedValue, err error) (context.Context, []driver.NamedValue, error)
 		AfterStmtQueryContext(ctx context.Context, query string, args []driver.NamedValue, rows driver.Rows, err error) (context.Context, driver.Rows, error)
+
 		BeforeStmtExecContext(ctx context.Context, query string, args []driver.NamedValue, err error) (context.Context, []driver.NamedValue, error)
 		AfterStmtExecContext(ctx context.Context, query string, args []driver.NamedValue, r driver.Result, err error) (context.Context, driver.Result, error)
 	}
diff --git a/hook_test.go b/hook_test.go
index be3788d..db8ebf0 100644
--- a/hook_test.go
+++ b/hook_test.go
@@ -74,10 +74,6 @@ func (m *mockHook) AfterPrepareContext(ctx context.Context, _ string, ds driver.
 
 func (m *mockHook) BeforeCommit(ctx context.Context, err error) (context.Context, error) {
 	m.Write("BeforeCommit")
-	txContext := TxContextFromContext(ctx)
-	if txContext == nil {
-		panic("txContext is nil")
-	}
 
 	prepareContext := PrepareContextFromContext(ctx)
 	if prepareContext != nil {
@@ -89,10 +85,6 @@ func (m *mockHook) BeforeCommit(ctx context.Context, err error) (context.Context
 
 func (m *mockHook) AfterCommit(ctx context.Context, err error) (context.Context, error) {
 	m.Write("AfterCommit")
-	txContext := TxContextFromContext(ctx)
-	if txContext == nil {
-		panic("txContext is nil")
-	}
 
 	prepareContext := PrepareContextFromContext(ctx)
 	if prepareContext != nil {
@@ -104,10 +96,6 @@ func (m *mockHook) AfterCommit(ctx context.Context, err error) (context.Context,
 
 func (m *mockHook) BeforeRollback(ctx context.Context, err error) (context.Context, error) {
 	m.Write("BeforeRollback")
-	txContext := TxContextFromContext(ctx)
-	if txContext == nil {
-		panic("txContext is nil")
-	}
 
 	prepareContext := PrepareContextFromContext(ctx)
 	if prepareContext != nil {
@@ -119,10 +107,6 @@ func (m *mockHook) BeforeRollback(ctx context.Context, err error) (context.Conte
 
 func (m *mockHook) AfterRollback(ctx context.Context, err error) (context.Context, error) {
 	m.Write("AfterRollback")
-	txContext := TxContextFromContext(ctx)
-	if txContext == nil {
-		panic("txContext is nil")
-	}
 
 	prepareContext := PrepareContextFromContext(ctx)
 	if prepareContext != nil {
@@ -139,11 +123,6 @@ func (m *mockHook) BeforeStmtQueryContext(ctx context.Context, _ string, args []
 		panic("prepareContext is nil")
 	}
 
-	txContext := TxContextFromContext(ctx)
-	if txContext != nil {
-		panic("txContext is not nil")
-	}
-
 	return ctx, args, err
 }
 
@@ -154,11 +133,6 @@ func (m *mockHook) AfterStmtQueryContext(ctx context.Context, _ string, _ []driv
 		panic("prepareContext is nil")
 	}
 
-	txContext := TxContextFromContext(ctx)
-	if txContext != nil {
-		panic("txContext is not nil")
-	}
-
 	return ctx, rows, err
 }
 
@@ -169,11 +143,6 @@ func (m *mockHook) BeforeStmtExecContext(ctx context.Context, _ string, args []d
 		panic("prepareContext is nil")
 	}
 
-	txContext := TxContextFromContext(ctx)
-	if txContext != nil {
-		panic("txContext is not nil")
-	}
-
 	return ctx, args, err
 }
 
@@ -184,11 +153,6 @@ func (m *mockHook) AfterStmtExecContext(ctx context.Context, _ string, _ []drive
 		panic("prepareContext is nil")
 	}
 
-	txContext := TxContextFromContext(ctx)
-	if txContext != nil {
-		panic("txContext is not nil")
-	}
-
 	return ctx, r, err
 }
 
diff --git a/tx.go b/tx.go
index ae6c91d..fbd6818 100644
--- a/tx.go
+++ b/tx.go
@@ -11,22 +11,10 @@ type (
 		TxHook
 		txContext context.Context
 	}
-	txContextKey struct{}
 )
 
-func TxContextFromContext(ctx context.Context) context.Context {
-	value := ctx.Value(txContextKey{})
-	if value != nil {
-		return value.(context.Context)
-	}
-
-	return nil
-}
-
-// -----------------
-
 func (t *tx) Commit() (err error) {
-	ctx := context.WithValue(context.Background(), txContextKey{}, t.txContext)
+	ctx := t.txContext
 	ctx, err = t.BeforeCommit(ctx, nil)
 	defer func() {
 		_, err = t.AfterCommit(ctx, err)
@@ -44,7 +32,7 @@ func (t *tx) Commit() (err error) {
 }
 
 func (t *tx) Rollback() (err error) {
-	ctx := context.WithValue(context.Background(), txContextKey{}, t.txContext)
+	ctx := t.txContext
 	ctx, err = t.BeforeRollback(ctx, nil)
 	defer func() {
 		_, err = t.AfterRollback(ctx, err)