Prepare handles describe failure

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,73 @@ 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) {
+		pgxtest.SkipCockroachDB(t, conn, "Induced error does not occur on CockroachDB")
+
+		_, 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,69 @@ 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)
+
+	pgxtest.SkipCockroachDB(t, conn, "Induced error does not occur on CockroachDB")
+
+	_, 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/config.go

@@ -55,6 +55,13 @@ type Config struct {
 
 	SSLNegotiation string // sslnegotiation=postgres or sslnegotiation=direct
 
+	// AfterNetConnect is called after the network connection, including TLS if applicable, is established but before any
+	// PostgreSQL protocol communication. It takes the established net.Conn and returns a net.Conn that will be used in
+	// its place. It can be used to wrap the net.Conn (e.g. for logging, diagnostics, or testing). Its functionality has
+	// some overlap with DialFunc. However, DialFunc takes place before TLS is established and cannot be used to control
+	// the final net.Conn used for PostgreSQL protocol communication while AfterNetConnect can.
+	AfterNetConnect func(ctx context.Context, config *Config, conn net.Conn) (net.Conn, error)
+
 	// ValidateConnect is called during a connection attempt after a successful authentication with the PostgreSQL server.
 	// It can be used to validate that the server is acceptable. If this returns an error the connection is closed and the next
 	// fallback config is tried. This allows implementing high availability behavior such as libpq does with target_session_attrs.

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
+}