From c7f668094ad2cdb0d3588901228ea235aeafeb5f Mon Sep 17 00:00:00 2001
From: anton98i <12586459+anton98i@users.noreply.github.com>
Date: Tue, 23 Aug 2022 16:07:44 +0200
Subject: [PATCH] Improved WrapInTx error handling

---
 main.go | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/main.go b/main.go
index de243c26..d07b7978 100644
--- a/main.go
+++ b/main.go
@@ -526,19 +526,35 @@ func (s *DB) Rollback() *DB {
 }
 
 // WrapInTx wraps a method in a transaction
-func (s *DB) WrapInTx(f func(tx *DB) error) error {
+func (s *DB) WrapInTx(f func(tx *DB) error) (err error) {
 	if _, ok := s.db.(*sql.Tx); ok {
 		// Already in a transaction
 		return f(s)
 	} else {
 		// Lets start a new transaction
 		tx := s.Begin()
-		if err := f(tx); err != nil {
-			tx.Rollback()
-			return err
+		if err = tx.Error; err != nil {
+			return
 		}
-		tx.Commit()
-		return nil
+		panicked := true
+		defer func() {
+			if panicked || err != nil {
+				rollbackErr := tx.Rollback().Error
+				if rollbackErr != nil {
+					if err == nil {
+						err = rollbackErr
+					} else {
+						err = fmt.Errorf("Transacton code and rollback failed: %s; %s", err, rollbackErr)
+					}
+				}
+			}
+		}()
+		err = f(tx)
+		if err == nil {
+			err = tx.Commit().Error
+		}
+		panicked = false
+		return
 	}
 }