go.mod: bump pgx to v5.9.2 for OAuthTokenProvider support

Bumps github.com/jackc/pgx/v5 from v5.7.2 to v5.9.2. The primary
motivation is that v5.9.2 adds OAuthTokenProvider to pgconn.Config,
enabling native OAUTHBEARER SASL client support.

Four compatibility fixes are required:

1. pgconn.PgConn.SecretKey() now returns []byte instead of uint32,
   supporting the new PostgreSQL protocol v3.2 variable-length cancel
   key. The sqlproxy cancel tests recover the original uint32 using
   binary.BigEndian.Uint32() since the proxy always sends a 4-byte
   key.

2. pgconn now returns a "conn closed" error from its internal state
   machine when operations are attempted on an already-closed
   connection, rather than propagating the OS-level "connection reset
   by peer" or "unexpected EOF". Two regex assertions in sqlproxy
   tests are updated to accept all three forms.

3. pgxpool.Config.BeforeAcquire is deprecated in v5.9.2 in favor of
   PrepareConn, which has the same semantics but also returns an
   error. The workload pool setup is migrated accordingly.

4. pgx v5.9.2 enforces the PostgreSQL pipeline protocol more strictly:
   ReadyForQuery must be sent after Sync, not after each Execute.
   TestPipelineMetric's fake server was sending ReadyForQuery after
   each Execute (via finishQuery(execPortal)), which worked with
   v5.7.2's lenient pipeline reader but fails with v5.9.2. Two new
   test helpers (expectExecPortalPipelined, expectSyncAndSendReadyForQuery)
   correctly simulate the pipeline response sequence.

Closes #171149
Closes #169280
Informs: CRDB-64341
Epic: CRDB-60766

Release note: None

Author: darshanjain10

DEPS.bzl

@@ -5134,10 +5134,10 @@ def go_deps():
         name = "com_github_jackc_pgx_v5",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/jackc/pgx/v5",
-        sha256 = "c4b5a22a3f3db2e764f5b4df65ad1b2550d4587b3640f73ab27b868136ef9018",
-        strip_prefix = "github.com/jackc/pgx/v5@v5.7.2",
+        sha256 = "0e8e8456630e580729cd1b06bf402cc6b98b57ec3a867df0e89e6ff27aa361c9",
+        strip_prefix = "github.com/jackc/pgx/v5@v5.9.2",
         urls = [
-            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v5/com_github_jackc_pgx_v5-v5.7.2.zip",
+            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v5/com_github_jackc_pgx_v5-v5.9.2.zip",
         ],
     )
     go_repository(

build/bazelutil/distdir_files.bzl

@@ -672,7 +672,7 @@ DISTDIR_FILES = {
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgservicefile/com_github_jackc_pgservicefile-v0.0.0-20240606120523-5a60cdf6a761.zip": "c9e31c91aebf96eb246bd410d1849cc7666d955a1e20ca2eba1c30b4eb89335f",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgtype/com_github_jackc_pgtype-v1.14.3.zip": "24eb1112e74a18ba22ef7b47f59808060ddef00b98c2ade780d81b283e40f676",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v4/com_github_jackc_pgx_v4-v4.18.3.zip": "1d5955dc65b8de8f72f9856865b33dcd9a2238f7cf9b1f2a00f1558e7d4965da",
-    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v5/com_github_jackc_pgx_v5-v5.7.2.zip": "c4b5a22a3f3db2e764f5b4df65ad1b2550d4587b3640f73ab27b868136ef9018",
+    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/pgx/v5/com_github_jackc_pgx_v5-v5.9.2.zip": "0e8e8456630e580729cd1b06bf402cc6b98b57ec3a867df0e89e6ff27aa361c9",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/puddle/com_github_jackc_puddle-v1.3.0.zip": "b1eb42bb3cf9a430146af79cb183860b9dddfca51844c2d4b447dc2f43becc55",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jackc/puddle/v2/com_github_jackc_puddle_v2-v2.2.2.zip": "f1f0789098a0bcb5ff3c7024f9ee387a6748446e1bc6713b13a63d763a9fb11e",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/jaegertracing/jaeger/com_github_jaegertracing_jaeger-v1.18.1.zip": "256a95b2a52a66494aca6d354224bb450ff38ce3ea1890af46a7c8dc39203891",

go.mod

@@ -185,7 +185,7 @@ require (
 	github.com/iancoleman/strcase v0.2.0
 	github.com/influxdata/influxdb-client-go/v2 v2.3.1-0.20210518120617-5d1fff431040
 	github.com/irfansharif/recorder v0.0.0-20211218081646-a21b46510fd6
-	github.com/jackc/pgx/v5 v5.7.2
+	github.com/jackc/pgx/v5 v5.9.2
 	github.com/jaegertracing/jaeger v1.18.1
 	github.com/jordan-wright/email v4.0.1-0.20210109023952-943e75fe5223+incompatible
 	github.com/jordanlewis/gcassert v0.0.0-20240401195008-3141cbd028c0

go.sum

@@ -1519,8 +1519,8 @@ github.com/jackc/pgx/v4 v4.12.1-0.20210724153913-640aa07df17c/go.mod h1:1QD0+tgS
 github.com/jackc/pgx/v4 v4.18.2/go.mod h1:Ey4Oru5tH5sB6tV7hDmfWFahwF15Eb7DNXlRKx2CkVw=
 github.com/jackc/pgx/v4 v4.18.3 h1:dE2/TrEsGX3RBprb3qryqSV9Y60iZN1C6i8IrmW9/BA=
 github.com/jackc/pgx/v4 v4.18.3/go.mod h1:Ey4Oru5tH5sB6tV7hDmfWFahwF15Eb7DNXlRKx2CkVw=
-github.com/jackc/pgx/v5 v5.7.2 h1:mLoDLV6sonKlvjIEsV56SkWNCnuNv531l94GaIzO+XI=
-github.com/jackc/pgx/v5 v5.7.2/go.mod h1:ncY89UGWxg82EykZUwSpUKEfccBGGYq1xjrOpsbsfGQ=
+github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw=
+github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
 github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=

pkg/sql/pgwire/conn_test.go

@@ -289,7 +289,7 @@ func TestPipelineMetric(t *testing.T) {
 	expectPrepareStmt(ctx, t, "", "SELECT $1::INT8 + $2::INT8", &rd, serverSideConn)
 	expectBindStmt(ctx, t, "", &rd, serverSideConn)
 	expectDescribeStmt(ctx, t, "", pgwirebase.PreparePortal, &rd, serverSideConn)
-	expectExecPortal(ctx, t, "", &rd, serverSideConn)
+	expectExecPortalPipelined(ctx, t, "", &rd, serverSideConn)
 	require.EqualValues(t, 5, serverSideConn.stmtBuf.PipelineCount.Value())
 
 	// Send another query in the pipeline.
@@ -310,19 +310,21 @@ func TestPipelineMetric(t *testing.T) {
 		return nil
 	})
 
-	// Process all of the commands that are in the pipeline.
+	// Process all of the commands that are in the pipeline. In the PostgreSQL
+	// pipeline protocol, ReadyForQuery is sent once per Sync (not per Execute),
+	// so we use expectExecPortalPipelined and expectSyncAndSendReadyForQuery.
 	expectPrepareStmt(ctx, t, "", "SELECT ($1::INT8 + $2::INT8) + $3::INT8", &rd, serverSideConn)
 	expectBindStmt(ctx, t, "", &rd, serverSideConn)
 	expectDescribeStmt(ctx, t, "", pgwirebase.PreparePortal, &rd, serverSideConn)
-	expectExecPortal(ctx, t, "", &rd, serverSideConn)
-	expectSync(ctx, t, &rd)
+	expectExecPortalPipelined(ctx, t, "", &rd, serverSideConn)
+	expectSyncAndSendReadyForQuery(ctx, t, &rd, serverSideConn)
 	require.EqualValues(t, 5, serverSideConn.stmtBuf.PipelineCount.Value())
 
 	expectPrepareStmt(ctx, t, "", "SELECT $1::STRING", &rd, serverSideConn)
 	expectBindStmt(ctx, t, "", &rd, serverSideConn)
 	expectDescribeStmt(ctx, t, "", pgwirebase.PreparePortal, &rd, serverSideConn)
-	expectExecPortal(ctx, t, "", &rd, serverSideConn)
-	expectSync(ctx, t, &rd)
+	expectExecPortalPipelined(ctx, t, "", &rd, serverSideConn)
+	expectSyncAndSendReadyForQuery(ctx, t, &rd, serverSideConn)
 	require.EqualValues(t, 0, serverSideConn.stmtBuf.PipelineCount.Value())
 
 	err = pipeline.Close()
@@ -951,6 +953,46 @@ func expectSync(ctx context.Context, t *testing.T, rd *sql.StmtBufReader) {
 	}
 }
 
+// expectSyncAndSendReadyForQuery reads a Sync command from the stmtBuf and
+// sends ReadyForQuery back to the client. In the PostgreSQL pipeline protocol,
+// ReadyForQuery is sent after Sync, not after each Execute.
+func expectSyncAndSendReadyForQuery(
+	ctx context.Context, t *testing.T, rd *sql.StmtBufReader, c *conn,
+) {
+	t.Helper()
+	expectSync(ctx, t, rd)
+	c.msgBuilder.initMsg(pgwirebase.ServerMsgReady)
+	c.msgBuilder.writeByte('I') // transaction status: no txn
+	if err := c.msgBuilder.finishMsg(c.conn); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// expectExecPortalPipelined reads an ExecPortal command from the stmtBuf and
+// sends only CommandComplete (no ReadyForQuery) back to the client. This
+// matches the PostgreSQL pipeline protocol where ReadyForQuery is sent after
+// Sync, not after each Execute.
+func expectExecPortalPipelined(
+	ctx context.Context, t *testing.T, expName string, rd *sql.StmtBufReader, c *conn,
+) {
+	t.Helper()
+	cmd, err := rd.CurCmd()
+	if err != nil {
+		t.Fatal(err)
+	}
+	rd.AdvanceOne()
+	ep, ok := cmd.(sql.ExecPortal)
+	if !ok {
+		t.Fatalf("expected command ExecPortal, got: %T (%+v)", cmd, cmd)
+	}
+	if ep.Name != expName {
+		t.Fatalf("expected name %s, got %s", expName, ep.Name)
+	}
+	if err := finishQuery(cmdComplete, c); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func expectExecPortal(
 	ctx context.Context, t *testing.T, expName string, rd *sql.StmtBufReader, c *conn,
 ) {

pkg/sqlproxy/proxy_handler_test.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"crypto/tls"
 	gosql "database/sql"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"net"
@@ -438,7 +439,7 @@ func TestPrivateEndpointsACL(t *testing.T) {
 					"Expected the connection to eventually fail",
 				)
 				require.Error(t, err)
-				require.Regexp(t, "connection reset by peer|unexpected EOF", err.Error())
+				require.Regexp(t, "connection reset by peer|unexpected EOF|conn closed", err.Error())
 				require.Equal(t, int64(1), s.metrics.ExpiredClientConnCount.Count())
 			},
 		)
@@ -568,7 +569,7 @@ func TestAllowedCIDRRangesACL(t *testing.T) {
 				time.Second, 5*time.Millisecond,
 				"Expected the connection to eventually fail",
 			)
-			require.Regexp(t, "connection reset by peer|unexpected EOF", err.Error())
+			require.Regexp(t, "connection reset by peer|unexpected EOF|conn closed", err.Error())
 			require.Equal(t, int64(1), s.metrics.ExpiredClientConnCount.Count())
 		})
 	})
@@ -1700,7 +1701,7 @@ func TestCancelQuery(t *testing.T) {
 		cancelFn = func() {
 			cancelRequest := proxyCancelRequest{
 				ProxyIP:   net.IP{},
-				SecretKey: conn.PgConn().SecretKey(),
+				SecretKey: binary.BigEndian.Uint32(conn.PgConn().SecretKey()),
 				ClientIP:  net.IP{127, 0, 0, 1},
 			}
 			u := "http://" + addrs.httpAddr + "/_status/cancel/"
@@ -1730,7 +1731,7 @@ func TestCancelQuery(t *testing.T) {
 			_ = conn.PgConn().CancelRequest(ctx)
 		}
 		defer testutils.TestingHook(&defaultTransferTimeout, 3*time.Minute)()
-		origCancelInfo, found := proxy.handler.cancelInfoMap.getCancelInfo(conn.PgConn().SecretKey())
+		origCancelInfo, found := proxy.handler.cancelInfoMap.getCancelInfo(binary.BigEndian.Uint32(conn.PgConn().SecretKey()))
 		require.True(t, found)
 		b := tds.DrainPod(tenantID, tenants[0].SQLAddr())
 		require.True(t, b)
@@ -1753,7 +1754,7 @@ func TestCancelQuery(t *testing.T) {
 		timeSource.Advance(2 * time.Minute)
 		proxy.handler.balancer.RebalanceTenant(ctx, tenantID)
 		testutils.SucceedsSoon(t, func() error {
-			newCancelInfo, found := proxy.handler.cancelInfoMap.getCancelInfo(conn.PgConn().SecretKey())
+			newCancelInfo, found := proxy.handler.cancelInfoMap.getCancelInfo(binary.BigEndian.Uint32(conn.PgConn().SecretKey()))
 			if !found {
 				return errors.New("expected to find cancel info")
 			}
@@ -1775,7 +1776,7 @@ func TestCancelQuery(t *testing.T) {
 		snapshot := snapshotMetrics()
 		cancelRequest := proxyCancelRequest{
 			ProxyIP:   net.IP{},
-			SecretKey: conn.PgConn().SecretKey(),
+			SecretKey: binary.BigEndian.Uint32(conn.PgConn().SecretKey()),
 			ClientIP:  net.IP{210, 1, 2, 3},
 		}
 		u := "http://" + addrs.httpAddr + "/_status/cancel/"
@@ -1812,7 +1813,7 @@ func TestCancelQuery(t *testing.T) {
 		})()
 		crdbRequest := &pgproto3.CancelRequest{
 			ProcessID: 1,
-			SecretKey: conn.PgConn().SecretKey() + 1,
+			SecretKey: binary.BigEndian.Uint32(conn.PgConn().SecretKey()) + 1,
 		}
 		buf, err := crdbRequest.Encode(nil /* buf */)
 		require.NoError(t, err)
@@ -1829,7 +1830,7 @@ func TestCancelQuery(t *testing.T) {
 		require.Equal(t, "http://0.0.0.1:8080/_status/cancel/", forwardedTo)
 		expectedReq := proxyCancelRequest{
 			ProxyIP:   net.IP{0, 0, 0, 1},
-			SecretKey: conn.PgConn().SecretKey() + 1,
+			SecretKey: binary.BigEndian.Uint32(conn.PgConn().SecretKey()) + 1,
 			ClientIP:  net.IP{127, 0, 0, 1},
 		}
 		require.Equal(t, expectedReq, forwardedReq)
@@ -1840,7 +1841,7 @@ func TestCancelQuery(t *testing.T) {
 		snapshot := snapshotMetrics()
 		cancelRequest := proxyCancelRequest{
 			ProxyIP:   net.IP{},
-			SecretKey: conn.PgConn().SecretKey() + 1,
+			SecretKey: binary.BigEndian.Uint32(conn.PgConn().SecretKey()) + 1,
 			ClientIP:  net.IP{127, 0, 0, 1},
 		}
 		u := "http://" + addrs.httpAddr + "/_status/cancel/"

pkg/workload/pgx_helpers.go

@@ -224,7 +224,7 @@ func NewMultiConnPool(
 				minConns = numConns
 			}
 			poolCfg.MinConns = int32(minConns)
-			poolCfg.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+			poolCfg.PrepareConn = func(ctx context.Context, conn *pgx.Conn) (bool, error) {
 				m.mu.RLock()
 				defer m.mu.RUnlock()
 				for name, sql := range m.mu.preparedStatements {
@@ -233,10 +233,10 @@ func NewMultiConnPool(
 					// communication to the server.
 					if _, err := conn.Prepare(ctx, name, sql); err != nil {
 						log.Dev.Warningf(ctx, "error preparing statement. name=%s sql=%s %v", name, sql, err)
-						return false
+						return false, nil
 					}
 				}
-				return true
+				return true, nil
 			}
 
 			// Attach the supplied tracer to the ConnConfig.