Keep connections alive after BEGIN statement for non-FATAL errors

pgxpool/pool_test.go

@@ -1066,6 +1066,35 @@ func TestConnReleaseWhenBeginFail(t *testing.T) {
 		assert.NoError(t, err)
 	}
 
+	assert.EqualValues(t, 1, db.Stat().TotalConns())
+
+	var n int
+	require.NoError(t, db.QueryRow(ctx, "select 1").Scan(&n))
+	require.EqualValues(t, 1, n)
+}
+
+func TestConnDestroyedWhenBeginFailsFatally(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	controllerConn, err := pgx.Connect(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer controllerConn.Close(ctx)
+	pgxtest.SkipCockroachDB(t, controllerConn, "Server does not support pg_terminate_backend() (https://github.com/cockroachdb/cockroach/issues/35897)")
+
+	db, err := pgxpool.New(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer db.Close()
+
+	tx, err := db.BeginTx(ctx, pgx.TxOptions{BeginQuery: "select pg_terminate_backend(pg_backend_pid())"})
+	assert.Error(t, err)
+	if !assert.Zero(t, tx) {
+		err := tx.Rollback(ctx)
+		assert.NoError(t, err)
+	}
+
 	for range 1000 {
 		if db.Stat().TotalConns() == 0 {
 			break

tx.go

@@ -100,9 +100,12 @@ func (c *Conn) Begin(ctx context.Context) (Tx, error) {
 func (c *Conn) BeginTx(ctx context.Context, txOptions TxOptions) (Tx, error) {
 	_, err := c.Exec(ctx, txOptions.beginSQL())
 	if err != nil {
-		// begin should never fail unless there is an underlying connection issue or
-		// a context timeout. In either case, the connection is possibly broken.
-		c.die()
+		// begin kills the connection upon receiving fatal, panic or non PGError errors,
+		// but does not otherwise as the connection should be reusable.
+		var pgErr *pgconn.PgError
+		if !errors.As(err, &pgErr) || isConnectionFatal(pgErr) {
+			c.die()
+		}
 		return nil, err
 	}
 
@@ -112,6 +115,14 @@ func (c *Conn) BeginTx(ctx context.Context, txOptions TxOptions) (Tx, error) {
 	}, nil
 }
 
+func isConnectionFatal(pgErr *pgconn.PgError) bool {
+	severity := pgErr.SeverityUnlocalized
+	if severity == "" {
+		severity = pgErr.Severity
+	}
+	return severity == "FATAL" || severity == "PANIC"
+}
+
 // Tx represents a database transaction.
 //
 // Tx is an interface instead of a struct to enable connection pools to be implemented without relying on internal pgx

tx_test.go

@@ -643,3 +643,47 @@ func TestTxSendBatchClosed(t *testing.T) {
 	_, err = br.Query()
 	require.Error(t, err)
 }
+
+func TestBeginTxNonFatalErrorKeepsConnAlive(t *testing.T) {
+	t.Parallel()
+
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	pgxtest.SkipCockroachDB(t, conn, "Server returns a different severity/error for an invalid BEGIN")
+
+	ctx := context.Background()
+
+	_, err := conn.BeginTx(ctx, pgx.TxOptions{BeginQuery: "begin transaction isolation level nonsense"})
+	require.Error(t, err)
+
+	var pgErr *pgconn.PgError
+	require.True(t, errors.As(err, &pgErr))
+	require.NotEqual(t, "FATAL", pgErr.SeverityUnlocalized)
+
+	require.False(t, conn.IsClosed())
+
+	var n int
+	require.NoError(t, conn.QueryRow(ctx, "select 1").Scan(&n))
+	require.Equal(t, 1, n)
+}
+
+func TestBeginTxFatalErrorKillsConn(t *testing.T) {
+	t.Parallel()
+
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	pgxtest.SkipCockroachDB(t, conn, "Server does not support pg_terminate_backend()")
+
+	ctx := context.Background()
+
+	_, err := conn.BeginTx(ctx, pgx.TxOptions{BeginQuery: "select pg_terminate_backend(pg_backend_pid())"})
+	require.Error(t, err)
+
+	var pgErr *pgconn.PgError
+	require.True(t, errors.As(err, &pgErr))
+	require.Equal(t, "FATAL", pgErr.SeverityUnlocalized)
+
+	require.True(t, conn.IsClosed())
+}