From 2d790cea0345ffbbbae68fb5845e8d852127b54d Mon Sep 17 00:00:00 2001
From: nabarun <nabarun@deepstacksoft.com>
Date: Thu, 10 Mar 2022 18:37:10 +0530
Subject: [PATCH] Refactor database transaction defer to match pattern
 elsewhere

---
 statediff/indexer/indexer.go | 60 +++++++++++++++++++++---------------
 1 file changed, 36 insertions(+), 24 deletions(-)

diff --git a/statediff/indexer/indexer.go b/statediff/indexer/indexer.go
index 58e0bafc5..65a5e950f 100644
--- a/statediff/indexer/indexer.go
+++ b/statediff/indexer/indexer.go
@@ -558,11 +558,20 @@ func (sdi *StateDiffIndexer) PushCodeAndCodeHash(tx *BlockTx, codeAndCodeHash sd
 
 // InsertWatchedAddresses inserts the given addresses in the database
 func (sdi *StateDiffIndexer) InsertWatchedAddresses(args []sdtypes.WatchAddressArg, currentBlockNumber *big.Int) error {
-	tx, err := sdi.dbWriter.db.Begin()
+	tx, err := sdi.dbWriter.db.Beginx()
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback()
+	defer func() {
+		if p := recover(); p != nil {
+			shared.Rollback(tx)
+			panic(p)
+		} else if err != nil {
+			shared.Rollback(tx)
+		} else {
+			err = tx.Commit()
+		}
+	}()
 
 	for _, arg := range args {
 		_, err = tx.Exec(`INSERT INTO eth_meta.watched_addresses (address, created_at, watched_at) VALUES ($1, $2, $3) ON CONFLICT (address) DO NOTHING`,
@@ -572,21 +581,25 @@ func (sdi *StateDiffIndexer) InsertWatchedAddresses(args []sdtypes.WatchAddressA
 		}
 	}
 
-	err = tx.Commit()
-	if err != nil {
-		return err
-	}
-
-	return nil
+	return err
 }
 
 // RemoveWatchedAddresses removes the given watched addresses from the database
 func (sdi *StateDiffIndexer) RemoveWatchedAddresses(args []sdtypes.WatchAddressArg) error {
-	tx, err := sdi.dbWriter.db.Begin()
+	tx, err := sdi.dbWriter.db.Beginx()
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback()
+	defer func() {
+		if p := recover(); p != nil {
+			shared.Rollback(tx)
+			panic(p)
+		} else if err != nil {
+			shared.Rollback(tx)
+		} else {
+			err = tx.Commit()
+		}
+	}()
 
 	for _, arg := range args {
 		_, err = tx.Exec(`DELETE FROM eth_meta.watched_addresses WHERE address = $1`, arg.Address)
@@ -595,21 +608,25 @@ func (sdi *StateDiffIndexer) RemoveWatchedAddresses(args []sdtypes.WatchAddressA
 		}
 	}
 
-	err = tx.Commit()
-	if err != nil {
-		return err
-	}
-
-	return nil
+	return err
 }
 
 // SetWatchedAddresses clears and inserts the given addresses in the database
 func (sdi *StateDiffIndexer) SetWatchedAddresses(args []sdtypes.WatchAddressArg, currentBlockNumber *big.Int) error {
-	tx, err := sdi.dbWriter.db.Begin()
+	tx, err := sdi.dbWriter.db.Beginx()
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback()
+	defer func() {
+		if p := recover(); p != nil {
+			shared.Rollback(tx)
+			panic(p)
+		} else if err != nil {
+			shared.Rollback(tx)
+		} else {
+			err = tx.Commit()
+		}
+	}()
 
 	_, err = tx.Exec(`DELETE FROM eth_meta.watched_addresses`)
 	if err != nil {
@@ -624,12 +641,7 @@ func (sdi *StateDiffIndexer) SetWatchedAddresses(args []sdtypes.WatchAddressArg,
 		}
 	}
 
-	err = tx.Commit()
-	if err != nil {
-		return err
-	}
-
-	return nil
+	return err
 }
 
 // ClearWatchedAddresses clears all the watched addresses from the database