~bigbes/core-go: database: don't panic with ErrBadConn

a06a6247898ca95c63a7396423218c85f0e83ea3 — Simon Ser 3 years ago d83af99

database: don't panic with ErrBadConn

Sometimes we're panic'ing like so:

    panic: driver: bad connection

    goroutine 9333 [running]:
    git.sr.ht/~sircmpwn/core-go/database.WithTx({0x561c2e623478, 0xc0002b6ed0}, 0x50?, 0xc00025ff38)
    	65b1657b30a1">git.sr.ht/~sircmpwn/core-go@v0.0.0-20230816134313-65b1657b30a1/database/middleware.go:57 +0x1fc
    git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders.fetchUsersByID.func1({0xc0002d80f0, 0x1, 0x1})
    	git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders/middleware.go:34 +0x118
    git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders.(*usersByIDLoaderBatch).end(...)
    	git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders/usersbyidloader_gen.go:222
    git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders.(*usersByIDLoaderBatch).startTimer(0xc0004e5380, 0xc0000e5140)
    	git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders/usersbyidloader_gen.go:218 +0xe7
    created by git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders.(*usersByIDLoaderBatch).keyIndex
    	git.sr.ht/~sircmpwn/builds.sr.ht/api/loaders/usersbyidloader_gen.go:191 +0x119

I believe this is a pq bug, see the linked bug report. Stop
panic'ing in that case.

patch browse .tar.gz .zip

1 files changed, 17 insertions(+), 8 deletions(-)

M database/middleware.go

M database/middleware.go => database/middleware.go +17 -8

@@ 3,6 3,7 @@ package database
 import (
 	"context"
 	"database/sql"
+	sqldriver "database/sql/driver"
 	"errors"
 	"net/http"
 )


@@ 45,22 46,30 @@ func DBForContext(ctx context.Context) *sql.DB {
 
 func WithTx(ctx context.Context, opts *sql.TxOptions, fn func(tx *sql.Tx) error) error {
 	db := DBForContext(ctx)
+
 	tx, err := db.BeginTx(ctx, opts)
 	if err != nil {
 		return err
 	}
 	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
 }