sql: add SHOW TRANSFER STATE observer statement

Informs cockroachdb#76000.

This commit adds the SHOW TRANSFER STATE observer statement, as described in
the sqlproxy connection migration RFC. This observer statement will be used
whenever a connection is about to be migrated to retrieve the relevant values
needed for the transfer process. A unique aspect to this statement is that
serialization or token generation errors will be returned as a SQL value
instead of an ErrorResponse. This will allow the sqlproxy to react accordingly,
and to reduce ambiguity issues during transferring.

This observer statement will only work on tenants (due to the need of token
generation). Since this is meant to be used internally only, there is no
release note.

Release note: None

Co-authored-by: Jay <[email protected]>
Co-authored-by: Rafi Shamim <[email protected]>

docs/generated/sql/bnf/BUILD.bazel

@@ -207,6 +207,7 @@ FILES = [
     "show_tables",
     "show_trace",
     "show_transactions_stmt",
+    "show_transfer_stmt",
     "show_types_stmt",
     "show_users_stmt",
     "show_var",

docs/generated/sql/bnf/show_transfer_stmt.bnf

@@ -0,0 +1,3 @@
+show_transfer_stmt ::=
+	'SHOW' 'TRANSFER' 'STATE' 'WITH' 'SCONST'
+	| 'SHOW' 'TRANSFER' 'STATE'

docs/generated/sql/bnf/show_var.bnf

@@ -29,6 +29,7 @@ show_stmt ::=
 	| show_tables_stmt
 	| show_trace_stmt
 	| show_transactions_stmt
+	| show_transfer_stmt
 	| show_users_stmt
 	| show_zone_stmt
 	| show_full_scans_stmt

docs/generated/sql/bnf/stmt_block.bnf

@@ -245,6 +245,7 @@ show_stmt ::=
 	| show_tables_stmt
 	| show_trace_stmt
 	| show_transactions_stmt
+	| show_transfer_stmt
 	| show_users_stmt
 	| show_zone_stmt
 	| show_full_scans_stmt
@@ -782,6 +783,10 @@ show_transactions_stmt ::=
 	'SHOW' opt_cluster 'TRANSACTIONS'
 	| 'SHOW' 'ALL' opt_cluster 'TRANSACTIONS'
 
+show_transfer_stmt ::=
+	'SHOW' 'TRANSFER' 'STATE' 'WITH' 'SCONST'
+	| 'SHOW' 'TRANSFER' 'STATE'
+
 show_users_stmt ::=
 	'SHOW' 'USERS'
 
@@ -1162,6 +1167,7 @@ unreserved_keyword ::=
 	| 'SQL'
 	| 'SQLLOGIN'
 	| 'START'
+	| 'STATE'
 	| 'STATEMENTS'
 	| 'STATISTICS'
 	| 'STDIN'
@@ -1188,6 +1194,7 @@ unreserved_keyword ::=
 	| 'TRACE'
 	| 'TRANSACTION'
 	| 'TRANSACTIONS'
+	| 'TRANSFER'
 	| 'TRIGGER'
 	| 'TRUNCATE'
 	| 'TRUSTED'

pkg/ccl/testccl/sqlccl/BUILD.bazel

@@ -6,6 +6,7 @@ go_test(
         "main_test.go",
         "run_control_test.go",
         "session_revival_test.go",
+        "show_transfer_state_test.go",
         "temp_table_clean_test.go",
     ],
     deps = [
@@ -31,6 +32,7 @@ go_test(
         "//pkg/util/stop",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
+        "@com_github_cockroachdb_cockroach_go_v2//crdb",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_gogo_protobuf//types",
         "@com_github_jackc_pgx_v4//:pgx",

pkg/ccl/testccl/sqlccl/show_transfer_state_test.go

@@ -0,0 +1,256 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Licensed as a CockroachDB Enterprise file under the Cockroach Community
+// License (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
+
+package sqlccl
+
+import (
+	"context"
+	gosql "database/sql"
+	"net/url"
+	"testing"
+
+	"github.com/cockroachdb/cockroach-go/v2/crdb"
+	"github.com/cockroachdb/cockroach/pkg/security"
+	"github.com/cockroachdb/cockroach/pkg/sql/tests"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/stretchr/testify/require"
+)
+
+func TestShowTransferState(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	ctx := context.Background()
+
+	params, _ := tests.CreateTestServerParams()
+	s, _, _ := serverutils.StartServer(t, params)
+	defer s.Stopper().Stop(ctx)
+	tenant, mainDB := serverutils.StartTenant(t, s, tests.CreateTestTenantParams(serverutils.TestTenantID()))
+	defer tenant.Stopper().Stop(ctx)
+	defer mainDB.Close()
+
+	_, err := mainDB.Exec("CREATE USER testuser WITH PASSWORD 'hunter2'")
+	require.NoError(t, err)
+	_, err = mainDB.Exec("SET CLUSTER SETTING server.user_login.session_revival_token.enabled = true")
+	require.NoError(t, err)
+
+	t.Run("without_transfer_key", func(t *testing.T) {
+		pgURL, cleanup := sqlutils.PGUrl(
+			t,
+			tenant.SQLAddr(),
+			"TestShowTransferState-without_transfer_key",
+			url.UserPassword(security.TestUser, "hunter2"),
+		)
+		defer cleanup()
+
+		conn, err := gosql.Open("postgres", pgURL.String())
+		require.NoError(t, err)
+		defer conn.Close()
+
+		rows, err := conn.Query("SHOW TRANSFER STATE")
+		require.NoError(t, err, "show transfer state failed")
+		defer rows.Close()
+
+		resultColumns, err := rows.Columns()
+		require.NoError(t, err)
+
+		require.Equal(t, []string{
+			"error",
+			"session_state_base64",
+			"session_revival_token_base64",
+		}, resultColumns)
+
+		var errVal, sessionState, sessionRevivalToken gosql.NullString
+
+		rows.Next()
+		err = rows.Scan(&errVal, &sessionState, &sessionRevivalToken)
+		require.NoError(t, err, "unexpected error while reading transfer state")
+
+		require.False(t, errVal.Valid)
+		require.True(t, sessionState.Valid)
+		require.True(t, sessionRevivalToken.Valid)
+	})
+
+	var state, token string
+	t.Run("with_transfer_key", func(t *testing.T) {
+		pgURL, cleanup := sqlutils.PGUrl(
+			t,
+			tenant.SQLAddr(),
+			"TestShowTransferState-with_transfer_key",
+			url.UserPassword(security.TestUser, "hunter2"),
+		)
+		defer cleanup()
+
+		q := pgURL.Query()
+		q.Add("application_name", "carl")
+		pgURL.RawQuery = q.Encode()
+		conn, err := gosql.Open("postgres", pgURL.String())
+		require.NoError(t, err)
+		defer conn.Close()
+
+		rows, err := conn.Query(`SHOW TRANSFER STATE WITH 'foobar'`)
+		require.NoError(t, err, "show transfer state failed")
+		defer rows.Close()
+
+		resultColumns, err := rows.Columns()
+		require.NoError(t, err)
+
+		require.Equal(t, []string{
+			"error",
+			"session_state_base64",
+			"session_revival_token_base64",
+			"transfer_key",
+		}, resultColumns)
+
+		var key string
+		var errVal, sessionState, sessionRevivalToken gosql.NullString
+
+		rows.Next()
+		err = rows.Scan(&errVal, &sessionState, &sessionRevivalToken, &key)
+		require.NoError(t, err, "unexpected error while reading transfer state")
+
+		require.Equal(t, "foobar", key)
+		require.False(t, errVal.Valid)
+		require.True(t, sessionState.Valid)
+		require.True(t, sessionRevivalToken.Valid)
+		state = sessionState.String
+		token = sessionRevivalToken.String
+	})
+
+	t.Run("successful_transfer", func(t *testing.T) {
+		pgURL, cleanup := sqlutils.PGUrl(
+			t,
+			tenant.SQLAddr(),
+			"TestShowTransferState-successful_transfer",
+			url.User(security.TestUser), // Do not use a password here.
+		)
+		defer cleanup()
+
+		q := pgURL.Query()
+		q.Add("application_name", "someotherapp")
+		q.Add("crdb:session_revival_token_base64", token)
+		pgURL.RawQuery = q.Encode()
+		conn, err := gosql.Open("postgres", pgURL.String())
+		require.NoError(t, err)
+		defer conn.Close()
+
+		var appName string
+		err = conn.QueryRow("SHOW application_name").Scan(&appName)
+		require.NoError(t, err)
+		require.Equal(t, "someotherapp", appName)
+
+		var b bool
+		err = conn.QueryRow(
+			"SELECT crdb_internal.deserialize_session(decode($1, 'base64'))",
+			state,
+		).Scan(&b)
+		require.NoError(t, err)
+		require.True(t, b)
+
+		err = conn.QueryRow("SHOW application_name").Scan(&appName)
+		require.NoError(t, err)
+		require.Equal(t, "carl", appName)
+	})
+
+	// Errors should be displayed as a SQL value.
+	t.Run("errors", func(t *testing.T) {
+		t.Run("root_user", func(t *testing.T) {
+			var key string
+			var errVal, sessionState, sessionRevivalToken gosql.NullString
+			err := mainDB.QueryRow(`SHOW TRANSFER STATE WITH 'bar'`).Scan(&errVal, &sessionState, &sessionRevivalToken, &key)
+			require.NoError(t, err)
+
+			require.True(t, errVal.Valid)
+			require.Equal(t, "cannot create token for root user", errVal.String)
+			require.False(t, sessionState.Valid)
+			require.False(t, sessionRevivalToken.Valid)
+		})
+
+		t.Run("transaction", func(t *testing.T) {
+			pgURL, cleanup := sqlutils.PGUrl(
+				t,
+				tenant.SQLAddr(),
+				"TestShowTransferState-errors-transaction",
+				url.UserPassword(security.TestUser, "hunter2"),
+			)
+			defer cleanup()
+
+			conn, err := gosql.Open("postgres", pgURL.String())
+			require.NoError(t, err)
+			defer conn.Close()
+
+			var errVal, sessionState, sessionRevivalToken gosql.NullString
+			err = crdb.ExecuteTx(ctx, conn, nil /* txopts */, func(tx *gosql.Tx) error {
+				return tx.QueryRow("SHOW TRANSFER STATE").Scan(&errVal, &sessionState, &sessionRevivalToken)
+			})
+			require.NoError(t, err)
+
+			require.True(t, errVal.Valid)
+			require.Equal(t, "cannot serialize a session which is inside a transaction", errVal.String)
+			require.False(t, sessionState.Valid)
+			require.False(t, sessionRevivalToken.Valid)
+		})
+
+		t.Run("prepared_statements", func(t *testing.T) {
+			pgURL, cleanup := sqlutils.PGUrl(
+				t,
+				tenant.SQLAddr(),
+				"TestShowTransferState-errors-prepared_statements",
+				url.UserPassword(security.TestUser, "hunter2"),
+			)
+			defer cleanup()
+
+			conn, err := gosql.Open("postgres", pgURL.String())
+			require.NoError(t, err)
+			defer conn.Close()
+
+			// Use a dummy prepared statement.
+			stmt, err := conn.Prepare("SELECT 1 WHERE 1 = 1")
+			require.NoError(t, err)
+			defer stmt.Close()
+
+			var errVal, sessionState, sessionRevivalToken gosql.NullString
+			err = conn.QueryRow("SHOW TRANSFER STATE").Scan(&errVal, &sessionState, &sessionRevivalToken)
+			require.NoError(t, err)
+
+			require.True(t, errVal.Valid)
+			require.Equal(t, "cannot serialize a session which has portals or prepared statements", errVal.String)
+			require.False(t, sessionState.Valid)
+			require.False(t, sessionRevivalToken.Valid)
+		})
+
+		t.Run("temp_tables", func(t *testing.T) {
+			pgURL, cleanup := sqlutils.PGUrl(
+				t,
+				tenant.SQLAddr(),
+				"TestShowTransferState-errors-temp_tables",
+				url.UserPassword(security.TestUser, "hunter2"),
+			)
+			defer cleanup()
+
+			q := pgURL.Query()
+			q.Add("experimental_enable_temp_tables", "true")
+			pgURL.RawQuery = q.Encode()
+			conn, err := gosql.Open("postgres", pgURL.String())
+			require.NoError(t, err)
+			defer conn.Close()
+
+			_, err = conn.Exec("CREATE TEMP TABLE temp_tbl()")
+			require.NoError(t, err)
+
+			var errVal, sessionState, sessionRevivalToken gosql.NullString
+			err = conn.QueryRow("SHOW TRANSFER STATE").Scan(&errVal, &sessionState, &sessionRevivalToken)
+			require.NoError(t, err)
+
+			require.True(t, errVal.Valid)
+			require.Equal(t, "cannot serialize session with temporary schemas", errVal.String)
+			require.False(t, sessionState.Valid)
+			require.False(t, sessionRevivalToken.Valid)
+		})
+	})
+}

pkg/gen/docs.bzl

@@ -219,6 +219,7 @@ DOCS_SRCS = [
   "//docs/generated/sql/bnf:show_tables.bnf",
   "//docs/generated/sql/bnf:show_trace.bnf",
   "//docs/generated/sql/bnf:show_transactions_stmt.bnf",
+  "//docs/generated/sql/bnf:show_transfer_stmt.bnf",
   "//docs/generated/sql/bnf:show_types_stmt.bnf",
   "//docs/generated/sql/bnf:show_users_stmt.bnf",
   "//docs/generated/sql/bnf:show_var.bnf",

pkg/sql/BUILD.bazel

@@ -182,6 +182,7 @@ go_library(
         "sequence_select.go",
         "serial.go",
         "session_revival_token.go",
+        "session_state.go",
         "set_cluster_setting.go",
         "set_default_isolation.go",
         "set_schema.go",