draft

Author: ZhouXing19

pkg/BUILD.bazel

@@ -48,6 +48,7 @@ ALL_TESTS = [
     "//pkg/ccl/cloudccl/externalconn:externalconn_test",
     "//pkg/ccl/cloudccl/gcp:gcp_test",
     "//pkg/ccl/jwtauthccl:jwtauthccl_test",
+    "//pkg/ccl/kvccl/kvtenantccl/upgradeinterlockccl:upgradeinterlockccl_test",
     "//pkg/ccl/ldapccl:ldapccl_test",
     "//pkg/ccl/logictestccl/tests/3node-tenant-multiregion:3node-tenant-multiregion_test",
     "//pkg/ccl/logictestccl/tests/3node-tenant:3node-tenant_test",
@@ -1018,6 +1019,7 @@ GO_TARGETS = [
     "//pkg/ccl/gssapiccl:gssapiccl",
     "//pkg/ccl/jwtauthccl:jwtauthccl",
     "//pkg/ccl/jwtauthccl:jwtauthccl_test",
+    "//pkg/ccl/kvccl/kvtenantccl/upgradeinterlockccl:upgradeinterlockccl_test",
     "//pkg/ccl/ldapccl:ldapccl",
     "//pkg/ccl/ldapccl:ldapccl_test",
     "//pkg/ccl/logictestccl/tests/3node-tenant-multiregion:3node-tenant-multiregion_test",

pkg/sql/explain_bundle.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec/explain"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/parserutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -556,6 +557,52 @@ var stmtBundleIncludeAllFKReferences = settings.RegisterBoolSetting(
 	false,
 )
 
+func (c *stmtEnvCollector) getStatementHints(
+	ctx context.Context, stmt string, sv *settings.Values,
+) (donorSqls []string, err error) {
+	fingerprintFlags := tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(
+		sv,
+	))
+	fingerprint, err := parserutils.FingerprintStatement(stmt, fingerprintFlags)
+	if err != nil {
+		return donorSqls, err
+	}
+
+	query := `SELECT crdb_internal.pb_to_json('hintpb.StatementHintUnion', hint)->'injectHints'->>'donorSql' AS donor_sql FROM system.statement_hints WHERE fingerprint = $1 ORDER BY "created_at"`
+
+	rows, err := c.ie.QueryBufferedEx(ctx, "stmtEnvCollector", nil /* txn */, c.ieo, query, fingerprint)
+	if err != nil {
+		return donorSqls, err
+	}
+
+	donorSqls = make([]string, len(rows))
+	for i, row := range rows {
+		if len(row) == 0 {
+			return nil, errors.AssertionFailedf("expect not nil row")
+		}
+		strDatum, ok := row[0].(*tree.DString)
+		if !ok {
+			return nil, errors.AssertionFailedf("expect row to be a string")
+		}
+		donorSqls[i] = strDatum.String()
+	}
+	return donorSqls, nil
+}
+
+func (c *stmtEnvCollector) PrintStmtHints(
+	ctx context.Context, stmt string, sv *settings.Values, w io.Writer,
+) error {
+	donorSqls, err := c.getStatementHints(ctx, stmt, sv)
+	if err != nil {
+		return err
+	}
+	for _, donorSql := range donorSqls {
+		// donorSql has been quoted.
+		fmt.Fprintf(w, "SELECT information_schema.crdb_rewrite_inline_hints('%s', %s);\n", stmt, donorSql)
+	}
+	return err
+}
+
 // stmtBundleStatsFileRE is a regex that matches all "complex" characters that
 // might not be safe when used in file names on some systems.
 var stmtBundleStatsFileRE = regexp.MustCompile(`[^a-zA-Z0-9_.\-]`)
@@ -1007,6 +1054,11 @@ func (b *stmtBundleBuilder) addEnv(ctx context.Context) {
 			b.printError(fmt.Sprintf("-- error getting schema for view %s: %v", views[i].FQString(), err), &buf)
 		}
 	}
+
+	if err := c.PrintStmtHints(ctx, b.stmt, b.sv, &buf); err != nil {
+		b.printError(fmt.Sprintf("-- error getting statement hints: %v", err), &buf)
+	}
+
 	if buf.Len() == 0 {
 		buf.WriteString("-- there were no objects used in this query\n")
 	}

pkg/sql/explain_bundle_test.go

@@ -1178,6 +1178,44 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 		)
 	})
 
+	t.Run("statement hints", func(t *testing.T) {
+		r.Exec(t, "CREATE TABLE table161829(x INT PRIMARY KEY, y INT)")
+		r.Exec(t, "CREATE INDEX xy161829 ON table161829 (x, y)")
+		r.Exec(t, "CREATE INDEX y161829 ON table161829 (y)")
+		r.Exec(t, `SELECT information_schema.crdb_rewrite_inline_hints('SELECT * FROM table161829 WHERE y = 10', 'SELECT * FROM table161829@primary WHERE y = 10')`)
+		r.Exec(t, `SELECT information_schema.crdb_rewrite_inline_hints('SELECT * FROM table161829 WHERE y = 10', 'SELECT * FROM table161829@xy161829 WHERE y = 10')`)
+		r.Exec(t, `SELECT information_schema.crdb_rewrite_inline_hints('SELECT * FROM table161829 WHERE y = 10', 'SELECT * FROM table161829@y161829 WHERE y = 10')`)
+
+		rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM table161829 WHERE y = 10")
+		checkBundle(
+			t, fmt.Sprint(rows), "table161829", func(name, contents string) error {
+				if name == "schema.sql" {
+					reg := regexp.MustCompile(`crdb_rewrite_inline_hints\(.*\@primary.*\n.*crdb_rewrite_inline_hints.*\@xy161829.*\n.*crdb_rewrite_inline_hints.*\@y161829`)
+					if reg.FindString(contents) == "" {
+						return errors.Errorf("could not find full crdb_rewrite_inline_hints in schema.sql")
+					}
+				}
+				return nil
+			}, false, /* expectErrors */
+			base, plans, "distsql.html vec.txt vec-v.txt stats-defaultdb.public.table161829.sql",
+		)
+
+		// Bundle for statements that share the same fingerprint should all see the same hints.
+		rows = r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM table161829 WHERE y = 100")
+		checkBundle(
+			t, fmt.Sprint(rows), "table161829", func(name, contents string) error {
+				if name == "schema.sql" {
+					reg := regexp.MustCompile(`crdb_rewrite_inline_hints\(.*\@primary.*\n.*crdb_rewrite_inline_hints.*\@xy161829.*\n.*crdb_rewrite_inline_hints.*\@y161829`)
+					if reg.FindString(contents) == "" {
+						return errors.Errorf("could not find full crdb_rewrite_inline_hints in schema.sql")
+					}
+				}
+				return nil
+			}, false, /* expectErrors */
+			base, plans, "distsql.html vec.txt vec-v.txt stats-defaultdb.public.table161829.sql",
+		)
+	})
+
 	// TODO(yuzefovich): figure out why this test occasionally fails under
 	// stress (i.e. it seems that no bundle is collected altogether).
 	//t.Run("under different users", func(t *testing.T) {

pkg/sql/parserutils/utils.go

@@ -105,3 +105,22 @@ var ParseQualifiedTableName = func(sql string) (*tree.TableName, error) {
 var PLpgSQLParse = func(sql string) (statements.PLpgStatement, error) {
 	return statements.PLpgStatement{}, errors.AssertionFailedf("sql.DoParserInjection hasn't been called")
 }
+
+// FingerprintStatement parses a SQL string and converts it to a statement
+// fingerprint. If the input is already a fingerprint, it is returned as-is.
+// The optional extraFlags are OR'd into the default FmtHideConstants flags used
+// for fingerprinting.
+func FingerprintStatement(sql string, extraFlags ...tree.FmtFlags) (string, error) {
+	stmts, err := Parse(sql)
+	if err != nil {
+		return "", pgerror.Wrapf(
+			err, pgcode.InvalidParameterValue, "could not parse statement",
+		)
+	}
+	if len(stmts) != 1 {
+		return "", pgerror.Newf(
+			pgcode.InvalidParameterValue, "could not parse as a single SQL statement",
+		)
+	}
+	return tree.FormatStatementHideConstants(stmts[0].AST, extraFlags...), nil
+}