@andrew-farries PTAL

This looks good 👍

We have tests for this package in db_test.go; we should add tests for QueryContext just as we have for ExecContext:

pgroll/pkg/db/db_test.go

Lines 22 to 62 in d64ba26

	func TestExecContext(t *testing.T) {
	t.Parallel()
	testutils.WithConnectionToContainer(t, func(conn *sql.DB, connStr string) {
	ctx := context.Background()
	// create a table on which an exclusive lock is held for 2 seconds
	setupTableLock(t, connStr, 2*time.Second)
	// set the lock timeout to 100ms
	ensureLockTimeout(t, conn, 100)
	// execute a query that should retry until the lock is released
	rdb := &db.RDB{DB: conn}
	_, err := rdb.ExecContext(ctx, "INSERT INTO test(id) VALUES (1)")
	require.NoError(t, err)
	})
	}
	func TestExecContextWhenContextCancelled(t *testing.T) {
	t.Parallel()
	testutils.WithConnectionToContainer(t, func(conn *sql.DB, connStr string) {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	// create a table on which an exclusive lock is held for 2 seconds
	setupTableLock(t, connStr, 2*time.Second)
	// set the lock timeout to 100ms
	ensureLockTimeout(t, conn, 100)
	// execute a query that should retry until the lock is released
	rdb := &db.RDB{DB: conn}
	// Cancel the context before the lock times out
	go time.AfterFunc(500*time.Millisecond, cancel)
	_, err := rdb.ExecContext(ctx, "INSERT INTO test(id) VALUES (1)")
	require.Errorf(t, err, "context canceled")
	})
	}

Done