From 93f28bc116526ba4decdd969a7b2b0b245ad70f1 Mon Sep 17 00:00:00 2001
From: Joe <sunfuze.1989@gmail.com>
Date: Fri, 24 Jun 2022 10:33:39 +0800
Subject: [PATCH] use callback to handle transaction

- make transaction have before and after hooks, so plugin can have hack before
or after transaction
---
 callbacks.go    | 37 +++++++++++++++++++++++++++++++------
 finisher_api.go | 16 +---------------
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/callbacks.go b/callbacks.go
index c060ea70..1b4e58ea 100644
--- a/callbacks.go
+++ b/callbacks.go
@@ -2,6 +2,7 @@ package gorm
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 	"reflect"
@@ -15,12 +16,13 @@ import (
 func initializeCallbacks(db *DB) *callbacks {
 	return &callbacks{
 		processors: map[string]*processor{
-			"create": {db: db},
-			"query":  {db: db},
-			"update": {db: db},
-			"delete": {db: db},
-			"row":    {db: db},
-			"raw":    {db: db},
+			"create":      {db: db},
+			"query":       {db: db},
+			"update":      {db: db},
+			"delete":      {db: db},
+			"row":         {db: db},
+			"raw":         {db: db},
+			"transaction": {db: db},
 		},
 	}
 }
@@ -72,6 +74,29 @@ func (cs *callbacks) Raw() *processor {
 	return cs.processors["raw"]
 }
 
+func (cs *callbacks) Transaction() *processor {
+	return cs.processors["transaction"]
+}
+
+func (p *processor) Begin(tx *DB, opt *sql.TxOptions) *DB {
+	var err error
+
+	switch beginner := tx.Statement.ConnPool.(type) {
+	case TxBeginner:
+		tx.Statement.ConnPool, err = beginner.BeginTx(tx.Statement.Context, opt)
+	case ConnPoolBeginner:
+		tx.Statement.ConnPool, err = beginner.BeginTx(tx.Statement.Context, opt)
+	default:
+		err = ErrInvalidTransaction
+	}
+
+	if err != nil {
+		tx.AddError(err)
+	}
+
+	return tx
+}
+
 func (p *processor) Execute(db *DB) *DB {
 	// call scopes
 	for len(db.Statement.scopes) > 0 {
diff --git a/finisher_api.go b/finisher_api.go
index 7a3f27ba..3e406c1c 100644
--- a/finisher_api.go
+++ b/finisher_api.go
@@ -619,27 +619,13 @@ func (db *DB) Begin(opts ...*sql.TxOptions) *DB {
 		// clone statement
 		tx  = db.getInstance().Session(&Session{Context: db.Statement.Context, NewDB: db.clone == 1})
 		opt *sql.TxOptions
-		err error
 	)
 
 	if len(opts) > 0 {
 		opt = opts[0]
 	}
 
-	switch beginner := tx.Statement.ConnPool.(type) {
-	case TxBeginner:
-		tx.Statement.ConnPool, err = beginner.BeginTx(tx.Statement.Context, opt)
-	case ConnPoolBeginner:
-		tx.Statement.ConnPool, err = beginner.BeginTx(tx.Statement.Context, opt)
-	default:
-		err = ErrInvalidTransaction
-	}
-
-	if err != nil {
-		tx.AddError(err)
-	}
-
-	return tx
+	return tx.callbacks.Transaction().Begin(tx, opt)
 }
 
 // Commit commit a transaction