From 0135852a0eaaca6c516c784b4f365393aeac5b6b Mon Sep 17 00:00:00 2001
From: paul cannon <paul@thepaul.org>
Date: Thu, 19 Dec 2019 03:27:27 -0600
Subject: [PATCH] storage/postgreskv: use transactional helper

We may never need this code to work with CockroachDB, but I'm on a
mission to avoid problematic uses of Begin() and BeginTx(), and anywhere
they appear is a possible place for someone to copy-and-paste and do
something wrong. dbutil.WithTx makes this code a little bit simpler too,
so it seems worthwhile.

Change-Id: I9b4ab484db4590cad5ab07de515bbf5d9708daed
---
 storage/postgreskv/client_test.go | 51 +++++++++++++------------------
 1 file changed, 21 insertions(+), 30 deletions(-)

diff --git a/storage/postgreskv/client_test.go b/storage/postgreskv/client_test.go
index 7ecd904b8..18b3d9863 100644
--- a/storage/postgreskv/client_test.go
+++ b/storage/postgreskv/client_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/zeebo/errs"
 
 	"storj.io/storj/private/dbutil/pgutil/pgtest"
+	"storj.io/storj/private/dbutil/txutil"
 	"storj.io/storj/storage"
 	"storj.io/storj/storage/testsuite"
 )
@@ -90,40 +91,30 @@ func BenchmarkSuite(b *testing.B) {
 	testsuite.RunBenchmarks(b, store)
 }
 
-func bulkImport(db *sql.DB, iter storage.Iterator) (err error) {
-	txn, err2 := db.Begin()
-	if err2 != nil {
-		return errs.New("Failed to start transaction: %v", err2)
-	}
-	defer func() {
-		if err == nil {
-			err = errs.Combine(err, txn.Commit())
-		} else {
-			err = errs.Combine(err, txn.Rollback())
+func bulkImport(db *sql.DB, iter storage.Iterator) error {
+	return txutil.WithTx(ctx, db, nil, func(ctx context.Context, txn *sql.Tx) (err error) {
+		stmt, err := txn.Prepare(pq.CopyIn("pathdata", "bucket", "fullpath", "metadata"))
+		if err != nil {
+			return errs.New("Failed to initialize COPY FROM: %v", err)
 		}
-	}()
+		defer func() {
+			err2 := stmt.Close()
+			if err2 != nil {
+				err = errs.Combine(err, errs.New("Failed to close COPY FROM statement: %v", err2))
+			}
+		}()
 
-	stmt, err2 := txn.Prepare(pq.CopyIn("pathdata", "bucket", "fullpath", "metadata"))
-	if err2 != nil {
-		return errs.New("Failed to initialize COPY FROM: %v", err)
-	}
-	defer func() {
-		err2 := stmt.Close()
-		if err2 != nil {
-			err = errs.Combine(err, errs.New("Failed to close COPY FROM statement: %v", err2))
+		var item storage.ListItem
+		for iter.Next(ctx, &item) {
+			if _, err := stmt.Exec([]byte(""), []byte(item.Key), []byte(item.Value)); err != nil {
+				return err
+			}
 		}
-	}()
-
-	var item storage.ListItem
-	for iter.Next(ctx, &item) {
-		if _, err := stmt.Exec([]byte(""), []byte(item.Key), []byte(item.Value)); err != nil {
-			return err
+		if _, err = stmt.Exec(); err != nil {
+			return errs.New("Failed to complete COPY FROM: %v", err)
 		}
-	}
-	if _, err = stmt.Exec(); err != nil {
-		return errs.New("Failed to complete COPY FROM: %v", err)
-	}
-	return nil
+		return nil
+	})
 }
 
 func bulkDeleteAll(db *sql.DB) error {