Merge pull request #2591 from AlisinaDevelo/test/tx-rollback-conn-closed

tx: regression test for Rollback after mid-query connection close (#2557)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: jackc

conn.go

@@ -475,6 +475,9 @@ func (c *Conn) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.C
 	}
 
 	if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil {
+		if c.queryTracer != nil {
+			c.queryTracer.TraceQueryEnd(ctx, c, TraceQueryEndData{Err: err})
+		}
 		return pgconn.CommandTag{}, err
 	}
 

tx_test.go

@@ -644,6 +644,36 @@ func TestTxSendBatchClosed(t *testing.T) {
 	require.Error(t, err)
 }
 
+// When a query is cancelled mid-flight the underlying connection is closed, but
+// the Tx handle has not been finalized. The first Rollback then fails to send
+// ROLLBACK over the dead connection. That error must be matchable as
+// pgconn.ErrConnClosed (transport gone) rather than ErrTxClosed (handle already
+// finalized), so callers can tell the two apart. A second Rollback returns
+// ErrTxClosed. See https://github.com/jackc/pgx/issues/2557.
+func TestTxRollbackOnClosedConnReturnsErrConnClosed(t *testing.T) {
+	t.Parallel()
+
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	tx, err := conn.Begin(context.Background())
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	_, err = tx.Exec(ctx, "select pg_sleep(5)")
+	require.Error(t, err)
+	require.True(t, conn.IsClosed())
+
+	err = tx.Rollback(context.Background())
+	require.ErrorIs(t, err, pgconn.ErrConnClosed)
+	require.NotErrorIs(t, err, pgx.ErrTxClosed)
+
+	err = tx.Rollback(context.Background())
+	require.ErrorIs(t, err, pgx.ErrTxClosed)
+}
+
 func TestBeginTxNonFatalErrorKeepsConnAlive(t *testing.T) {
 	t.Parallel()