Handle for preparing statements that fail during the Describe phase

A failure that occurred after a successful Parse could leave a prepared
statement on the server that must be deallocated before re-preparing.

Also, introduce a testing utility faultyconn.Conn that can be used to
intercept and modify frontend messages. This is used to induce difficult
difficult to reproduce timing and protocol errors in tests.

Author: jackc

batch_test.go

@@ -4,12 +4,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
+	"net"
 	"os"
 	"testing"
 	"time"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/internal/faultyconn"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgproto3"
 	"github.com/jackc/pgx/v5/pgxtest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -1060,6 +1064,71 @@ func TestSendBatchStatementTimeout(t *testing.T) {
 
 }
 
+func TestSendBatchHandlesTimeoutBetweenParseAndDescribe(t *testing.T) {
+	// Not parallel because it is a timing sensitive test.
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	var faultyConn *faultyconn.Conn
+	faultyConnTestRunner := pgxtest.DefaultConnTestRunner()
+	faultyConnTestRunner.CreateConfig = func(ctx context.Context, t testing.TB) *pgx.ConnConfig {
+		config, err := pgx.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+		require.NoError(t, err)
+		config.AfterNetConnect = func(ctx context.Context, config *pgconn.Config, conn net.Conn) (net.Conn, error) {
+			faultyConn = faultyconn.New(conn)
+			return faultyConn, nil
+		}
+		return config
+	}
+
+	// Only need to test modes that use Parse/Describe.
+	extendedQueryModes := []pgx.QueryExecMode{
+		pgx.QueryExecModeCacheStatement,
+		pgx.QueryExecModeCacheDescribe,
+		pgx.QueryExecModeDescribeExec,
+		pgx.QueryExecModeExec,
+	}
+
+	pgxtest.RunWithQueryExecModes(ctx, t, faultyConnTestRunner, extendedQueryModes, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		_, err := conn.Exec(ctx, "set statement_timeout = '100ms'")
+		require.NoError(t, err)
+
+		batch := &pgx.Batch{}
+		batch.Queue("select 1")
+		batch.Queue("select 2")
+
+		faultyConn.HandleFrontendMessage = func(backendWriter io.Writer, msg pgproto3.FrontendMessage) error {
+			if _, ok := msg.(*pgproto3.Describe); ok {
+				time.Sleep(200 * time.Millisecond)
+			}
+			buf, err := msg.Encode(nil)
+			if err != nil {
+				return err
+			}
+			_, err = backendWriter.Write(buf)
+			return err
+		}
+
+		err = conn.SendBatch(ctx, batch).Close()
+		require.Error(t, err)
+		var pgErr *pgconn.PgError
+		require.True(t, errors.As(err, &pgErr))
+		require.Equal(t, "57014", pgErr.Code)
+
+		faultyConn.HandleFrontendMessage = nil
+
+		_, err = conn.Exec(ctx, "set statement_timeout = default")
+		require.NoError(t, err)
+
+		batch = &pgx.Batch{}
+		batch.Queue("select 1")
+		batch.Queue("select 2")
+		err = conn.SendBatch(ctx, batch).Close()
+		require.NoError(t, err)
+	})
+}
+
 func ExampleConn_SendBatch() {
 	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
 	defer cancel()

conn.go

@@ -65,11 +65,12 @@ func (cc *ConnConfig) ConnString() string { return cc.connString }
 // Conn is a PostgreSQL connection handle. It is not safe for concurrent usage. Use a connection pool to manage access
 // to multiple database connections from multiple goroutines.
 type Conn struct {
-	pgConn             *pgconn.PgConn
-	config             *ConnConfig // config used when establishing this connection
-	preparedStatements map[string]*pgconn.StatementDescription
-	statementCache     stmtcache.Cache
-	descriptionCache   stmtcache.Cache
+	pgConn                  *pgconn.PgConn
+	config                  *ConnConfig // config used when establishing this connection
+	preparedStatements      map[string]*pgconn.StatementDescription
+	failedDescribeStatement string
+	statementCache          stmtcache.Cache
+	descriptionCache        stmtcache.Cache
 
 	queryTracer    QueryTracer
 	batchTracer    BatchTracer
@@ -314,6 +315,14 @@ func (c *Conn) Close(ctx context.Context) error {
 // Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same name and sql arguments. This
 // allows a code path to Prepare and Query/Exec without concern for if the statement has already been prepared.
 func (c *Conn) Prepare(ctx context.Context, name, sql string) (sd *pgconn.StatementDescription, err error) {
+	if c.failedDescribeStatement != "" {
+		err = c.Deallocate(ctx, c.failedDescribeStatement)
+		if err != nil {
+			return nil, fmt.Errorf("failed to deallocate previously failed statement %q: %w", c.failedDescribeStatement, err)
+		}
+		c.failedDescribeStatement = ""
+	}
+
 	if c.prepareTracer != nil {
 		ctx = c.prepareTracer.TracePrepareStart(ctx, c, TracePrepareStartData{Name: name, SQL: sql})
 	}
@@ -346,6 +355,10 @@ func (c *Conn) Prepare(ctx context.Context, name, sql string) (sd *pgconn.Statem
 
 	sd, err = c.pgConn.Prepare(ctx, psName, sql, nil)
 	if err != nil {
+		var pErr *pgconn.PrepareError
+		if errors.As(err, &pErr) {
+			c.failedDescribeStatement = psKey
+		}
 		return nil, err
 	}
 

conn_test.go

@@ -4,14 +4,18 @@ import (
 	"bytes"
 	"context"
 	"database/sql"
+	"io"
+	"net"
 	"os"
 	"strings"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/internal/faultyconn"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgproto3"
 	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/jackc/pgx/v5/pgxtest"
 	"github.com/stretchr/testify/assert"
@@ -466,6 +470,67 @@ func TestPrepare(t *testing.T) {
 	}
 }
 
+// https://github.com/jackc/pgx/issues/2223
+func TestPrepareHandlesTimeoutBetweenParseAndDescribe(t *testing.T) {
+	// Not parallel because it is a timing sensitive test.
+
+	config, err := pgx.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	var faultyConn *faultyconn.Conn
+	config.AfterNetConnect = func(ctx context.Context, config *pgconn.Config, conn net.Conn) (net.Conn, error) {
+		faultyConn = faultyconn.New(conn)
+		return faultyConn, nil
+	}
+
+	ctx := context.Background()
+	conn, err := pgx.ConnectConfig(ctx, config)
+	require.NoError(t, err)
+	defer closeConn(t, conn)
+	require.NotNil(t, faultyConn)
+
+	_, err = conn.Exec(ctx, "set statement_timeout = '100ms'")
+	require.NoError(t, err)
+
+	faultyConn.HandleFrontendMessage = func(backendWriter io.Writer, msg pgproto3.FrontendMessage) error {
+		if _, ok := msg.(*pgproto3.Describe); ok {
+			time.Sleep(200 * time.Millisecond)
+		}
+		buf, err := msg.Encode(nil)
+		if err != nil {
+			return err
+		}
+		_, err = backendWriter.Write(buf)
+		return err
+	}
+
+	psd, err := conn.Prepare(ctx, "test", "select $1::varchar")
+	var pgErr *pgconn.PgError
+	require.ErrorAs(t, err, &pgErr)
+	require.Equal(t, "57014", pgErr.Code)
+	require.Nil(t, psd)
+
+	faultyConn.HandleFrontendMessage = nil
+
+	_, err = conn.Exec(ctx, "set statement_timeout = default")
+	require.NoError(t, err)
+
+	var existsOnServer bool
+	err = conn.QueryRow(
+		ctx,
+		"select exists(select 1 from pg_prepared_statements where name = 'test')",
+		// Avoid using the prepared statement cache or it will clear the broken statement before we can check for its
+		// existence.
+		pgx.QueryExecModeExec,
+	).Scan(&existsOnServer)
+	require.NoError(t, err)
+	require.True(t, existsOnServer)
+
+	psd, err = conn.Prepare(ctx, "test", "select $1::varchar")
+	require.NoError(t, err)
+	require.NotNil(t, psd)
+}
+
 func TestPrepareBadSQLFailure(t *testing.T) {
 	t.Parallel()
 

internal/faultyconn/faultyconn.go

@@ -0,0 +1,99 @@
+package faultyconn
+
+import (
+	"bytes"
+	"io"
+	"net"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgproto3"
+)
+
+// Conn is a wrapper for a net.Conn that allows inspection and modification of messages between a PostgreSQL client and
+// server. It is designed to be used in tests that use a *pgx.Conn or *pgconn.PgConn connected to a real PostgreSQL
+// server. Instead of mocking an entire server connection, this is used to specify and modify only the particular
+// aspects of a connection that are necessary. This can be easier to setup and is more true to real world conditions.
+//
+// It currently only supports handling frontend messages.
+type Conn struct {
+	// HandleFrontendMessage is called for each frontend message received. It should use backendWriter to write to the
+	// backend.
+	HandleFrontendMessage func(backendWriter io.Writer, msg pgproto3.FrontendMessage) error
+
+	// TODO: Implement this if we need to handle backend messages.
+	// HandleBackendMessage  func(w io.Writer, msg pgproto3.BackendMessage) error
+
+	conn            net.Conn
+	fromFrontendBuf *bytes.Buffer
+	fromBackendBuf  *bytes.Buffer
+	backend         *pgproto3.Backend
+	frontend        *pgproto3.Frontend
+}
+
+// NewConn creates a new Conn that proxies the messages sent to and from the given net.Conn. New can be used with
+// pgconn.Config.AfterNetConnect to wrap a net.Conn for testing purposes.
+func New(c net.Conn) *Conn {
+	fromFrontendBuf := &bytes.Buffer{}
+	fromBackendBuf := &bytes.Buffer{}
+
+	return &Conn{
+		conn:            c,
+		fromFrontendBuf: fromFrontendBuf,
+		fromBackendBuf:  fromBackendBuf,
+		backend:         pgproto3.NewBackend(fromFrontendBuf, c),
+		frontend:        pgproto3.NewFrontend(fromBackendBuf, c),
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	return c.conn.Read(b)
+}
+
+func (c *Conn) Write(b []byte) (n int, err error) {
+	if c.HandleFrontendMessage == nil {
+		return c.conn.Write(b)
+	}
+
+	c.fromFrontendBuf.Write(b)
+
+	for {
+		msg, err := c.backend.Receive()
+		if err != nil {
+			if err == io.ErrUnexpectedEOF {
+				break
+			}
+			return len(b), err
+		}
+
+		err = c.HandleFrontendMessage(c.conn, msg)
+		if err != nil {
+			return len(b), err
+		}
+	}
+
+	return len(b), nil
+}
+
+func (c *Conn) Close() error {
+	return c.conn.Close()
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	return c.conn.LocalAddr()
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.conn.RemoteAddr()
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	return c.conn.SetDeadline(t)
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	return c.conn.SetReadDeadline(t)
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	return c.conn.SetWriteDeadline(t)
+}

pgconn/errors.go

@@ -254,3 +254,20 @@ func (e *NotPreferredError) SafeToRetry() bool {
 func (e *NotPreferredError) Unwrap() error {
 	return e.err
 }
+
+type PrepareError struct {
+	err error
+
+	ParseComplete bool // Indicates whether the error occurred after a ParseComplete message was received.
+}
+
+func (e *PrepareError) Error() string {
+	if e.ParseComplete {
+		return fmt.Sprintf("prepare failed after ParseComplete: %s", e.err.Error())
+	}
+	return e.err.Error()
+}
+
+func (e *PrepareError) Unwrap() error {
+	return e.err
+}

pgconn/pgconn.go

@@ -866,6 +866,10 @@ type StatementDescription struct {
 //
 // Prepare does not send a PREPARE statement to the server. It uses the PostgreSQL Parse and Describe protocol messages
 // directly.
+//
+// In extremely rare cases, Prepare may fail after the Parse is successful, but before the Describe is complete. In this
+// case, the returned error will be an error where errors.As with a *PrepareError succeeds and the *PrepareError has
+// ParseComplete set to true.
 func (pgConn *PgConn) Prepare(ctx context.Context, name, sql string, paramOIDs []uint32) (*StatementDescription, error) {
 	if err := pgConn.lock(); err != nil {
 		return nil, err
@@ -893,7 +897,8 @@ func (pgConn *PgConn) Prepare(ctx context.Context, name, sql string, paramOIDs [
 
 	psd := &StatementDescription{Name: name, SQL: sql}
 
-	var parseErr error
+	var ParseComplete bool
+	var pgErr *PgError
 
 readloop:
 	for {
@@ -904,20 +909,22 @@ readloop:
 		}
 
 		switch msg := msg.(type) {
+		case *pgproto3.ParseComplete:
+			ParseComplete = true
 		case *pgproto3.ParameterDescription:
 			psd.ParamOIDs = make([]uint32, len(msg.ParameterOIDs))
 			copy(psd.ParamOIDs, msg.ParameterOIDs)
 		case *pgproto3.RowDescription:
 			psd.Fields = pgConn.convertRowDescription(nil, msg)
 		case *pgproto3.ErrorResponse:
-			parseErr = ErrorResponseToPgError(msg)
+			pgErr = ErrorResponseToPgError(msg)
 		case *pgproto3.ReadyForQuery:
 			break readloop
 		}
 	}
 
-	if parseErr != nil {
-		return nil, parseErr
+	if pgErr != nil {
+		return nil, &PrepareError{err: pgErr, ParseComplete: ParseComplete}
 	}
 	return psd, nil
 }