From d84ce15b297a7098cec05c0ee6c82349cba4a271 Mon Sep 17 00:00:00 2001
From: chen quan <chenquan.dev@gmail.com>
Date: Wed, 18 Jan 2023 11:43:54 +0800
Subject: [PATCH] Revert "refactor: refactor hook (#16)"

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

diff --git a/hook.go b/hook.go
index e370e57..18cc6d4 100644
--- a/hook.go
+++ b/hook.go
@@ -24,7 +24,6 @@ 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)
 	}
@@ -35,14 +34,12 @@ 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 db8ebf0..be3788d 100644
--- a/hook_test.go
+++ b/hook_test.go
@@ -74,6 +74,10 @@ 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 {
@@ -85,6 +89,10 @@ 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 {
@@ -96,6 +104,10 @@ 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 {
@@ -107,6 +119,10 @@ 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 {
@@ -123,6 +139,11 @@ 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
 }
 
@@ -133,6 +154,11 @@ 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
 }
 
@@ -143,6 +169,11 @@ 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
 }
 
@@ -153,6 +184,11 @@ 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 fbd6818..ae6c91d 100644
--- a/tx.go
+++ b/tx.go
@@ -11,10 +11,22 @@ 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 := t.txContext
+	ctx := context.WithValue(context.Background(), txContextKey{}, t.txContext)
 	ctx, err = t.BeforeCommit(ctx, nil)
 	defer func() {
 		_, err = t.AfterCommit(ctx, err)
@@ -32,7 +44,7 @@ func (t *tx) Commit() (err error) {
 }
 
 func (t *tx) Rollback() (err error) {
-	ctx := t.txContext
+	ctx := context.WithValue(context.Background(), txContextKey{}, t.txContext)
 	ctx, err = t.BeforeRollback(ctx, nil)
 	defer func() {
 		_, err = t.AfterRollback(ctx, err)