Merge pull request #18 from knocknote/feature/add_condition_before_setting_global_callback

Add condition before setting global callback

Author: goccy

database/sql/tx.go

@@ -122,14 +122,18 @@ func (proxy *Tx) begin(conn *connection.DBConnection) {
 		return
 	}
 	tx := conn.Begin(proxy.ctx, proxy.opts)
-	proxy.BeforeCommitCallback(func(writeQueries []*QueryLog) error {
-		return errors.WithStack(globalBeforeCommitCallback(proxy, writeQueries))
-	})
-	proxy.AfterCommitCallback(func() error {
-		return errors.WithStack(globalAfterCommitSuccessCallback(proxy))
-	}, func(isCritical bool, failureQueries []*QueryLog) error {
-		return errors.WithStack(globalAfterCommitFailureCallback(proxy, isCritical, failureQueries))
-	})
+	if proxy.beforeCommitCallback == nil {
+		proxy.BeforeCommitCallback(func(writeQueries []*QueryLog) error {
+			return errors.WithStack(globalBeforeCommitCallback(proxy, writeQueries))
+		})
+	}
+	if proxy.afterCommitSuccessCallback == nil && proxy.afterCommitFailureCallback == nil {
+		proxy.AfterCommitCallback(func() error {
+			return errors.WithStack(globalAfterCommitSuccessCallback(proxy))
+		}, func(isCritical bool, failureQueries []*QueryLog) error {
+			return errors.WithStack(globalAfterCommitFailureCallback(proxy, isCritical, failureQueries))
+		})
+	}
 	proxy.tx = tx
 }
 

transaction_test.go

@@ -320,6 +320,37 @@ func TestCommitCallbackForTx(t *testing.T) {
 	}
 }
 
+func TestCommitCallbackForTxSetCallbackBeforeQueryExec(t *testing.T) {
+	db, err := sql.Open("", "")
+	if err != nil {
+		t.Fatalf("%+v\n", err)
+	}
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatalf("%+v\n", err)
+	}
+	isInvokedBeforeCommitCallback := false
+	tx.BeforeCommitCallback(func(writeQueries []*sql.QueryLog) error {
+		isInvokedBeforeCommitCallback = true
+		return nil
+	})
+	isInvokedAfterCommitCallback := false
+	tx.AfterCommitCallback(func() error {
+		isInvokedAfterCommitCallback = true
+		return nil
+	}, func(isCriticalError bool, failureQueries []*sql.QueryLog) error {
+		return nil
+	})
+	insertRecords(tx, t)
+	checkErr(t, tx.Commit())
+	if !isInvokedBeforeCommitCallback {
+		t.Fatal("cannot invoke callback for before commit")
+	}
+	if !isInvokedAfterCommitCallback {
+		t.Fatal("cannot invoke callback for after commit")
+	}
+}
+
 func testIsAlreadyCommittedQueryLog(t *testing.T, queryLog *sql.QueryLog) {
 	initializeTables(t)
 	db, err := sql.Open("", "")