Repository: jackc/pgx

Issue:
Occasional prepared statement already exist errors when using statement timeout
**Describe the bug**

We have a web service where we set a `statement_timeout` on our connections to ensure that queries do not run for an extended period of time. As part of this web service, we have a large, generated query. We occasionally see errors such as the following:

```
ERROR: prepared statement "stmtcache_2b7087dc2ab6296b87fa5baa7f329541c23d17376faa03b7" already exists (SQLSTATE 42P05)
```

My hypothesis is that there is a statement timeout in the `PREPARE` but after Postgres has actually created the prepared statement. This causes `pgx` to end up in a state where the prepared statement isn't added to the statement cache but does exist for the connection in Postgres. I was able to reproduce this with the code below.

**To Reproduce**

The following sample code reliably produces the error with a local Postgres server for me. This probably could be simplified a bit further.

```go
package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"os"

	"github.com/jackc/pgx/v5/pgconn"
	_ "github.com/jackc/pgx/v5/stdlib"
)

func main() {
	db, err := sql.Open("pgx", os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatal(err)
	}

	ctx := context.Background()

	conn, err := db.Conn(ctx)
	if err != nil {
		log.Fatal(err)
	}

	_, err = conn.ExecContext(ctx, "SET statement_timeout TO 1")
	if err != nil {
		log.Fatal(err)
	}

	p := "s"
	for {
		var v uint

		q := fmt.Sprintf("SELECT 1 AS %s", p)
		err := conn.QueryRowContext(
			ctx,
			q,
		).Scan(&v)
		if err != nil && !errors.Is(err, context.Canceled) {
			var pgErr *pgconn.PgError
			if errors.As(err, &pgErr) && pgErr.Code == "57014" {
				continue
			}
			log.Fatal(err)
		}
		p = randString(100000)
	}
}

var letterRunes = []rune("abAB")

func randString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}
```

**Expected behavior**

No prepared statement already exists error.

**Actual behavior**

```
ERROR: prepared statement "stmtcache_087f9f36620ff21723cd288c17d7806ba6fdd9b160f73186" already exists (SQLSTATE 42P05)
```

**Version**
 - Go: `go version go1.23.4 linux/amd64`
 - PostgreSQL: `PostgreSQL 16.6 (Ubuntu 16.6-1.pgdg24.10+1) on x86_64-pc-linux-gnu, compiled by gcc (Ubuntu 14.2.0-4ubuntu2) 14.2.0, 64-bit`
 - pgx: `v5.7.2`
 

Discussion:
@jackc: > My hypothesis is that there is a statement timeout in the PREPARE but after Postgres has actually created the prepared statement. This causes the pgx to end up in a state where the prepared statement isn't added to the statement cache but does exist for the connection in Postgres. I was able to reproduce this with the code below.

You are correct. Based on your example code I was able to build a reproduction case directly at the PostgreSQL protocol layer.

I send a `Parse`, sleep, and then send a `Describe`. The `Parse` succeeds and the `Describe` fails. Preparing and releasing statements are *not* transactional, so even though the extended protocol pseudo-transaction fails the prepare is not undone.

pgx and pgconn always send `Parse` and `Describe` together . pgx keeps a record of the statement description for each prepared statement. 

In theory, I know how to solve this, but in practice it's tricky / messy enough I'm not sure. `pgconn.PgConn.Prepare()` could return a special error that denoted that `Parse` succeeded and `Describe` failed. Callers would detect that error and deallocate the statement on the server. Similar logic would need to be implemented for batches and pipeline mode. However, there are many contexts in which the prepare path can be called, so I'm not confident in this solution's robustness.

```go
package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgproto3"
)

func main() {
	ctx := context.Background()

	conn, err := pgx.Connect(ctx, os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close(ctx)

	_, err = conn.Exec(ctx, "SET statement_timeout TO 1")
	if err != nil {
		log.Fatal(err)
	}

	conn.PgConn().Frontend().SendParse(&pgproto3.Parse{Name: "ps1", Query: "select 1", ParameterOIDs: nil})
	err = conn.PgConn().Frontend().Flush()
	if err != nil {
		log.Fatal(err)
	}
	time.Sleep(1 * time.Millisecond)
	conn.PgConn().Frontend().SendDescribe(&pgproto3.Describe{ObjectType: 'S', Name: "ps1"})
	time.Sleep(1 * time.Millisecond)
	conn.PgConn().Frontend().SendSync(&pgproto3.Sync{})
	err = conn.PgConn().Frontend().Flush()
	if err != nil {
		log.Fatal(err)
	}

	for range 3 {
		msg, err := conn.PgConn().Frontend().Receive()
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%#v\n", msg)
	}

	conn.PgConn().Frontend().SendParse(&pgproto3.Parse{Name: "ps1", Query: "select 1", ParameterOIDs: nil})
	err = conn.PgConn().Frontend().Flush()
	if err != nil {
		log.Fatal(err)
	}
	time.Sleep(1 * time.Millisecond)
	conn.PgConn().Frontend().SendDescribe(&pgproto3.Describe{ObjectType: 'S', Name: "ps1"})
	time.Sleep(1 * time.Millisecond)
	conn.PgConn().Frontend().SendSync(&pgproto3.Sync{})
	err = conn.PgConn().Frontend().Flush()
	if err != nil {
		log.Fatal(err)
	}

	for range 2 {
		msg, err := conn.PgConn().Frontend().Receive()
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%#v\n", msg)
	}
}
```

```
jack@glados ~/dev/pgx_issues/pgx-2223 ±master⚡ » gor .
&pgproto3.ParseComplete{}
&pgproto3.ErrorResponse{Severity:"ERROR", SeverityUnlocalized:"ERROR", Code:"57014", Message:"canceling statement due to statement timeout", Detail:"", Hint:"", Position:0, InternalPosition:0, InternalQuery:"", Where:"", SchemaName:"", TableName:"", ColumnName:"", DataTypeName:"", ConstraintName:"", File:"postgres.c", Line:3356, Routine:"ProcessInterrupts", UnknownFields:map[uint8]string(nil)}
&pgproto3.ReadyForQuery{TxStatus:0x49}
&pgproto3.ErrorResponse{Severity:"ERROR", SeverityUnlocalized:"ERROR", Code:"42P05", Message:"prepared statement \"ps1\" already exists", Detail:"", Hint:"", Position:0, InternalPosition:0, InternalQuery:"", Where:"", SchemaName:"", TableName:"", ColumnName:"", DataTypeName:"", ConstraintName:"", File:"prepare.c", Line:415, Routine:"StorePreparedStatement", UnknownFields:map[uint8]string(nil)}
&pgproto3.ReadyForQuery{TxStatus:0x49}
```


Solution patch:
diff --git a/conn.go b/conn.go
index 340bca5a..42e955fc 100644
--- a/conn.go
+++ b/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
 	}
 
diff --git a/internal/faultyconn/faultyconn.go b/internal/faultyconn/faultyconn.go
new file mode 100644
index 00000000..73d4a067
--- /dev/null
+++ b/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)
+}
diff --git a/pgconn/config.go b/pgconn/config.go
index 3937dc40..3cd5d877 100644
--- a/pgconn/config.go
+++ b/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.
diff --git a/pgconn/errors.go b/pgconn/errors.go
index d968d3f0..bc1e31e3 100644
--- a/pgconn/errors.go
+++ b/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
+}
diff --git a/pgconn/pgconn.go b/pgconn/pgconn.go
index 321656f9..5c77a836 100644
--- a/pgconn/pgconn.go
+++ b/pgconn/pgconn.go
@@ -343,6 +343,14 @@ func connectOne(ctx context.Context, config *Config, connectConfig *connectOneCo
 		pgConn.conn = tlsConn
 	}
 
+	if config.AfterNetConnect != nil {
+		pgConn.conn, err = config.AfterNetConnect(ctx, config, pgConn.conn)
+		if err != nil {
+			pgConn.conn.Close()
+			return nil, newPerDialConnectError("AfterNetConnect failed", err)
+		}
+	}
+
 	pgConn.contextWatcher = ctxwatch.NewContextWatcher(config.BuildContextWatcherHandler(pgConn))
 	pgConn.contextWatcher.Watch(ctx)
 	defer pgConn.contextWatcher.Unwatch()
@@ -858,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
@@ -885,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 {
@@ -896,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
 }



Test patch:
diff --git a/batch_test.go b/batch_test.go
index b4c421e5..d45ff526 100644
--- a/batch_test.go
+++ b/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()
diff --git a/conn_test.go b/conn_test.go
index 10959f06..d570ad8b 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"context"
 	"database/sql"
+	"io"
+	"net"
 	"os"
 	"strings"
 	"sync"
@@ -11,7 +13,9 @@ import (
 	"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()
 
diff --git a/pgconn/pgconn_test.go b/pgconn/pgconn_test.go
index 4d6770f8..9d0ff8c3 100644
--- a/pgconn/pgconn_test.go
+++ b/pgconn/pgconn_test.go
@@ -544,6 +544,63 @@ func TestConnectFailsWithResolveFailureAndFailedConnectionAttempts(t *testing.T)
 	require.ErrorContains(t, err, ":2 (127.0.0.1): dial error:")
 }
 
+type testConnWrapper struct {
+	conn net.Conn
+}
+
+func (w *testConnWrapper) Read(b []byte) (n int, err error) {
+	return w.conn.Read(b)
+}
+
+func (w *testConnWrapper) Write(b []byte) (n int, err error) {
+	return w.conn.Write(b)
+}
+
+func (w *testConnWrapper) Close() error {
+	return w.conn.Close()
+}
+
+func (w *testConnWrapper) LocalAddr() net.Addr {
+	return w.conn.LocalAddr()
+}
+
+func (w *testConnWrapper) RemoteAddr() net.Addr {
+	return w.conn.RemoteAddr()
+}
+
+func (w *testConnWrapper) SetDeadline(t time.Time) error {
+	return w.conn.SetDeadline(t)
+}
+
+func (w *testConnWrapper) SetReadDeadline(t time.Time) error {
+	return w.conn.SetReadDeadline(t)
+}
+
+func (w *testConnWrapper) SetWriteDeadline(t time.Time) error {
+	return w.conn.SetWriteDeadline(t)
+}
+
+func TestConnectWithAfterNetConnect(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	var wrappedConn net.Conn
+	config.AfterNetConnect = func(ctx context.Context, config *pgconn.Config, conn net.Conn) (net.Conn, error) {
+		wrappedConn = &testConnWrapper{conn: conn}
+		return wrappedConn, nil
+	}
+	conn, err := pgconn.ConnectConfig(ctx, config)
+	require.NoError(t, err)
+
+	require.Equal(t, wrappedConn, conn.Conn())
+	closeConn(t, conn)
+}
+
 func TestConnectWithValidateConnect(t *testing.T) {
 	t.Parallel()
 
diff --git a/stdlib/sql_test.go b/stdlib/sql_test.go
index 84f48f6c..f9d441fe 100644
--- a/stdlib/sql_test.go
+++ b/stdlib/sql_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math"
 	"os"
@@ -511,7 +512,7 @@ func TestConnQueryFailure(t *testing.T) {
 	testWithAllQueryExecModes(t, func(t *testing.T, db *sql.DB) {
 		_, err := db.Query("select 'foo")
 		require.Error(t, err)
-		require.IsType(t, new(pgconn.PgError), err)
+		require.ErrorAs(t, err, new(*pgconn.PgError))
 	})
 }
 
@@ -755,7 +756,8 @@ func TestBeginTxContextCancel(t *testing.T) {
 
 		var n int
 		err = db.QueryRow("select count(*) from t").Scan(&n)
-		if pgErr, ok := err.(*pgconn.PgError); !ok || pgErr.Code != "42P01" {
+		var pgErr *pgconn.PgError
+		if !errors.As(err, &pgErr) || pgErr.Code != "42P01" {
 			t.Fatalf(`err => %v, want PgError{Code: "42P01"}`, err)
 		}
 	})