From 677aeafaed87b4ab3b5b67be2a551026cc6e4051 Mon Sep 17 00:00:00 2001
From: Yusuke Sakurai <kerokerokerop@gmail.com>
Date: Sun, 10 May 2020 14:30:18 +0900
Subject: [PATCH] feat: recover and rollback on panic during transaction (#13)

---
 tx.go      | 12 +++++++++++-
 tx_test.go | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/tx.go b/tx.go
index dee96e4..2d482d9 100644
--- a/tx.go
+++ b/tx.go
@@ -3,6 +3,7 @@ package exql
 import (
 	"context"
 	"database/sql"
+	"fmt"
 )
 
 type Tx interface {
@@ -51,7 +52,16 @@ func transaction(db *sql.DB, ctx context.Context, opts *sql.TxOptions, callback
 		return err
 	}
 	tx := &tx{tx: sqlTx, s: NewSaver(sqlTx), m: NewMapper()}
-	txErr := callback(tx)
+	var p interface{}
+	txErr := func() error {
+		defer func() {
+			p = recover()
+		}()
+		return callback(tx)
+	}()
+	if p != nil {
+		txErr = fmt.Errorf("recovered: %s", p)
+	}
 	if txErr != nil {
 		if err := sqlTx.Rollback(); err != nil {
 			return err
diff --git a/tx_test.go b/tx_test.go
index 36491ca..c648fbd 100644
--- a/tx_test.go
+++ b/tx_test.go
@@ -53,6 +53,23 @@ func TestTx_Transaction(t *testing.T) {
 		err = db.Map(rows, &dest)
 		assert.Error(t, err, ErrRecordNotFound.Error())
 	})
+	t.Run("should rollback if panic happened during transaction", func(t *testing.T) {
+		var user *model.Users
+		err := transaction(db.DB(), context.Background(), nil, func(tx Tx) error {
+			user = &model.Users{
+				FirstName: null.String{},
+				LastName:  null.String{},
+			}
+			_, err := tx.Insert(user)
+			assert.Nil(t, err)
+			panic("panic")
+		})
+		assert.EqualError(t, err, "recovered: panic")
+		rows, err := db.DB().Query(`select * from users where id = ?`, user.Id)
+		assert.Nil(t, err)
+		var dest model.Users
+		assert.Equal(t, db.Map(rows, &dest), ErrRecordNotFound)
+	})
 }
 func TestTx_Map(t *testing.T) {
 	db := testDb()