sql: do not evaluate AOST timestamp in session migrations

Release note (bug fix): Fixed a bug where a session migration performed
by SHOW TRANSFER STATE would not handle prepared statements that used
the AS OF SYSTEM TIME clause. Users who encountered this bug would see
errors such as `expected 1 or 0 for number of format codes, got N`. This
bug was present since v22.2.0.

Author: rafiss

pkg/ccl/testccl/sqlccl/show_transfer_state_test.go

@@ -43,6 +43,12 @@ func TestShowTransferState(t *testing.T) {
 	require.NoError(t, err)
 	_, err = tenantDB.Exec("CREATE TYPE typ AS ENUM ('foo', 'bar')")
 	require.NoError(t, err)
+	_, err = tenantDB.Exec("CREATE TABLE tab (a INT4, b typ)")
+	require.NoError(t, err)
+	_, err = tenantDB.Exec("INSERT INTO tab VALUES (1, 'foo')")
+	require.NoError(t, err)
+	_, err = tenantDB.Exec("GRANT SELECT ON tab TO testuser")
+	require.NoError(t, err)
 
 	testUserConn := tenant.SQLConnForUser(t, username.TestUser, "")
 
@@ -90,12 +96,18 @@ func TestShowTransferState(t *testing.T) {
 		defer func() { _ = conn.Close(ctx) }()
 
 		// Add a prepared statement to make sure SHOW TRANSFER STATE handles it.
-		_, err = conn.Prepare(ctx, "prepared_stmt", "SELECT $1::INT4, 'foo'::typ WHERE 1 = 1")
+		_, err = conn.Prepare(ctx, "prepared_stmt_const", "SELECT $1::INT4, 'foo'::typ WHERE 1 = 1")
+		require.NoError(t, err)
+		_, err = conn.Prepare(ctx, "prepared_stmt_aost", "SELECT a, b FROM tab AS OF SYSTEM TIME '-1us'")
 		require.NoError(t, err)
 
 		var intResult int
 		var enumResult string
-		err = conn.QueryRow(ctx, "prepared_stmt", 1).Scan(&intResult, &enumResult)
+		err = conn.QueryRow(ctx, "prepared_stmt_const", 1).Scan(&intResult, &enumResult)
+		require.NoError(t, err)
+		require.Equal(t, 1, intResult)
+		require.Equal(t, "foo", enumResult)
+		err = conn.QueryRow(ctx, "prepared_stmt_aost").Scan(&intResult, &enumResult)
 		require.NoError(t, err)
 		require.Equal(t, 1, intResult)
 		require.Equal(t, "foo", enumResult)
@@ -171,14 +183,26 @@ func TestShowTransferState(t *testing.T) {
 		// session.
 		result := conn.PgConn().ExecPrepared(
 			ctx,
-			"prepared_stmt",
+			"prepared_stmt_const",
 			[][]byte{{0, 0, 0, 2}}, // binary representation of 2
 			[]int16{1},             // paramFormats - 1 means binary
-			[]int16{1},             // resultFormats - 1 means binary
+			[]int16{1, 1},          // resultFormats - 1 means binary
 		).Read()
+		require.NoError(t, result.Err)
 		require.Equal(t, [][][]byte{{
 			{0, 0, 0, 2}, {0x66, 0x6f, 0x6f}, // binary representation of 2, 'foo'
 		}}, result.Rows)
+		result = conn.PgConn().ExecPrepared(
+			ctx,
+			"prepared_stmt_aost",
+			[][]byte{},    // paramValues
+			[]int16{},     // paramFormats
+			[]int16{1, 1}, // resultFormats - 1 means binary
+		).Read()
+		require.NoError(t, result.Err)
+		require.Equal(t, [][][]byte{{
+			{0, 0, 0, 1}, {0x66, 0x6f, 0x6f}, // binary representation of 1, 'foo'
+		}}, result.Rows)
 	})
 
 	// Errors should be displayed as a SQL value.

pkg/sql/conn_executor_prepare.go

@@ -288,8 +288,12 @@ func (ex *connExecutor) prepare(
 	}
 
 	// Use the existing transaction.
-	if err := prepare(ctx, ex.state.mu.txn); err != nil && origin != PreparedStatementOriginSessionMigration {
-		return nil, err
+	if err := prepare(ctx, ex.state.mu.txn); err != nil {
+		if origin != PreparedStatementOriginSessionMigration {
+			return nil, err
+		} else {
+			log.Warningf(ctx, "could not prepare statement during session migration: %v", err)
+		}
 	}
 
 	// Account for the memory used by this prepared statement.
@@ -320,8 +324,15 @@ func (ex *connExecutor) populatePrepared(
 		return 0, err
 	}
 	p.extendedEvalCtx.PrepareOnly = true
-	if err := ex.handleAOST(ctx, p.stmt.AST); err != nil {
-		return 0, err
+	// If the statement is being prepared by a session migration, then we should
+	// not evaluate the AS OF SYSTEM TIME timestamp. During session migration,
+	// there is no way for the statement being prepared to be executed in this
+	// transaction, so there's no need to fix the timestamp, unlike how we must
+	// for pgwire- or SQL-level prepared statements.
+	if origin != PreparedStatementOriginSessionMigration {
+		if err := ex.handleAOST(ctx, p.stmt.AST); err != nil {
+			return 0, err
+		}
 	}
 
 	// PREPARE has a limited subset of statements it can be run with. Postgres

pkg/sql/testdata/session_migration/prepared_statements

@@ -33,6 +33,11 @@ wire_prepare s4
 INSERT INTO t2 VALUES($1, $2)
 ----
 
+# Regression test for transferring statements with AOST.
+wire_prepare s5
+SELECT a, b FROM t2 AS OF SYSTEM TIME '-2us'
+----
+
 wire_prepare s_empty
 ;
 ----
@@ -102,6 +107,15 @@ SELECT * FROM t2
 ----
 1 cat
 
+query
+SELECT pg_sleep(0.1)
+----
+true
+
+wire_query s5
+----
+1 cat
+
 reset
 ----