fix(tfstated): rollack transaction in case of panic

1 files changed, 13 insertions, 8 deletions

diff --git a/pkg/database/db.go b/pkg/database/db.go
index 5501a82..0777556 100644
--- a/pkg/database/db.go
+++ b/pkg/database/db.go
@@ -108,15 +108,20 @@ func (db *DB) SetVersionsHistoryLimit(n int) {
 func (db *DB) WithTransaction(f func(tx *sql.Tx) error) error {
 	tx, err := db.writeDB.Begin()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to begin transaction: %w", err)
 	}
-	err = f(tx)
-	if err == nil {
-		err = tx.Commit()
-	}
-	if err != nil {
-		if err2 := tx.Rollback(); err2 != nil {
-			panic(fmt.Sprintf("failed to rollback transaction: %+v. Reason for rollback: %+v", err2, err))
+	defer func() {
+		if err != nil {
+			if err2 := tx.Rollback(); err2 != nil {
+				panic(fmt.Sprintf("failed to rollback transaction: %+v. Reason for rollback: %+v", err2, err))
+			}
+		}
+	}()
+	if err = f(tx); err != nil {
+		return fmt.Errorf("failed to execute function inside transaction: %w", err)
+	} else {
+		if err = tx.Commit(); err != nil {
+			err = fmt.Errorf("failed to commit transaction: %w", err)
 		}
 	}
 	return err