From 918e04d24e53da6eb6d9ae2e891d05a20b608887 Mon Sep 17 00:00:00 2001
From: Peter Sanchez <peter@netlandish.com>
Date: Wed, 16 Aug 2023 12:45:23 -0600
Subject: [PATCH] When a query fails because the context has been canceled, pq
 will return "driver: bad connection" from tx.Rollback. Do not panic in this
 case. See https://github.com/lib/pq/issues/1137

Taken from this sourcehut patch submission:

https://lists.sr.ht/~sircmpwn/sr.ht-dev/patches/43630
---
 database/sql.go | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/database/sql.go b/database/sql.go
index 1344638..ce4e5ab 100644
--- a/database/sql.go
+++ b/database/sql.go
@@ -3,6 +3,7 @@ package database
 import (
 	"context"
 	"database/sql"
+	sqldriver "database/sql/driver"
 	"errors"
 	"time"
 
@@ -18,7 +19,7 @@ type contextKey struct {
 }
 
 // TxOptionsRO read only transaction rules
-var TxOptionsRO *sql.TxOptions = &sql.TxOptions{Isolation: 0, ReadOnly: true}
+var TxOptionsRO = &sql.TxOptions{Isolation: 0, ReadOnly: true}
 
 // ContextWithTimeout returns a context with query timeout
 func ContextWithTimeout(
@@ -49,17 +50,24 @@ func WithTx(ctx context.Context, opts *sql.TxOptions, fn func(tx *sql.Tx) error)
 	}
 	defer tx.Rollback()
 	err = fn(tx)
+
+	var txErr error
 	if err != nil {
-		err := tx.Rollback()
-		if err != nil && err != sql.ErrTxDone {
-			panic(err)
-		}
+		txErr = tx.Rollback()
 	} else {
-		err := tx.Commit()
-		if err != nil && err != sql.ErrTxDone {
-			panic(err)
-		}
+		txErr = tx.Commit()
+	}
+
+	if errors.Is(err, context.Canceled) && errors.Is(txErr, sqldriver.ErrBadConn) {
+		// When a query fails because the context has been canceled, pq will
+		// return "driver: bad connection" from tx.Rollback. Do not panic in
+		// this case. See https://github.com/lib/pq/issues/1137
+		return err
 	}
+	if txErr != nil && txErr != sql.ErrTxDone {
+		panic(err)
+	}
+
 	return err
 }
 
-- 
2.45.2