sql: add statement hints to statement diagnostics bundles

Fixes cockroachdb#161829

Release note (sql change): We now include `information_schema.crdb_rewrite_inline_hints`
statements in the `schema.sql` file of a statement diagnostics bundle
for re-creating all the statement hints bound to the statement. The hints
recreation statement have been sorted in ascending order of the original
hints' creation time.

Author: ZhouXing19

pkg/sql/explain_bundle.go

@@ -23,11 +23,13 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/lexbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"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 +558,74 @@ var stmtBundleIncludeAllFKReferences = settings.RegisterBoolSetting(
 	false,
 )
 
+// getStmtHintRecreateStmts returns the SQL statements that can be used to
+// recreate all the statement hints bound to the target statement. The returned
+// statements are sorted in ascending order by the original hints' creation time.
+func (c *stmtEnvCollector) getStmtHintRecreateStmts(
+	ctx context.Context, stmt string, sv *settings.Values,
+) (recreateStmts []string, err error) {
+	fingerprint, err := parserutils.FingerprintStatement(stmt,
+		tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(
+			sv,
+		)))
+	if err != nil {
+		return nil, err
+	}
+
+	query := `SELECT hint 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 nil, err
+	}
+
+	recreateStmts = make([]string, 0, len(rows))
+	for _, row := range rows {
+		if len(row) == 0 {
+			return nil, errors.AssertionFailedf("expect not nil row")
+		}
+		if row[0] == tree.DNull {
+			continue
+		}
+		hintBytes, ok := row[0].(*tree.DBytes)
+		if !ok {
+			return nil, errors.AssertionFailedf("expected hint to be bytes, got %T", row[0])
+		}
+		// Deserialize from the protobuf.
+		hint, err := hintpb.FromBytes([]byte(*hintBytes))
+		if err != nil {
+			return nil, errors.Wrap(err, "deserializing statement hint")
+		}
+		recreateSQL, ok := hint.RecreateStmt(stmt)
+		if !ok {
+			continue
+		}
+		recreateStmts = append(recreateStmts, recreateSQL)
+	}
+	return recreateStmts, nil
+}
+
+// PrintStmtHints writes the SQL statements needed to recreate the statement
+// hints for the given statement to w.
+func (c *stmtEnvCollector) PrintStmtHints(
+	ctx context.Context, stmt string, sv *settings.Values, w io.Writer,
+) error {
+	recreateStmts, err := c.getStmtHintRecreateStmts(ctx, stmt, sv)
+	if err != nil {
+		return err
+	}
+	for _, s := range recreateStmts {
+		fmt.Fprintf(w, "%s\n", s)
+	}
+	return nil
+}
+
 // 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 +1077,13 @@ func (b *stmtBundleBuilder) addEnv(ctx context.Context) {
 			b.printError(fmt.Sprintf("-- error getting schema for view %s: %v", views[i].FQString(), err), &buf)
 		}
 	}
+
+	if !b.flags.RedactValues {
+		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

@@ -1145,6 +1145,7 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 		r.Exec(t, "CREATE POLICY policy1 ON rls1 USING (u = 'hal')")
 		r.Exec(t, "CREATE USER rls_user")
 		r.Exec(t, "GRANT SYSTEM VIEWCLUSTERSETTING TO rls_user")
+		r.Exec(t, "GRANT SYSTEM VIEWSYSTEMTABLE TO rls_user")
 		r.Exec(t, "ALTER TABLE rls1 OWNER TO rls_user")
 		r.Exec(t, "SET ROLE rls_user")
 		defer r.Exec(t, "SET ROLE root")
@@ -1178,6 +1179,61 @@ 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, z STRING)")
+		r.Exec(t, "CREATE INDEX xy161829 ON table161829 (x, y)")
+		r.Exec(t, "CREATE INDEX y161829 ON table161829 (y)")
+		r.Exec(t, "CREATE INDEX xz161829 ON table161829 (x, z)")
+		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",
+		)
+
+		// Test that the stmt argument must have its single quotes escaped.
+		r.Exec(t, `SELECT information_schema.crdb_rewrite_inline_hints('SELECT * FROM table161829 WHERE z = ''hello''', 'SELECT * FROM table161829@xz161829 WHERE z = ''hello''')`)
+		rows = r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM table161829 WHERE z = 'hello'")
+		checkBundle(
+			t, fmt.Sprint(rows), "table161829", func(name, contents string) error {
+				if name == "schema.sql" {
+					reg := regexp.MustCompile(`crdb_rewrite_inline_hints\(e\'SELECT \* FROM table161829 WHERE z = \\'hello\\'\'`)
+					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/hintpb/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/hintpb",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/sql/lexbase",
         "//pkg/sql/protoreflect",
         "//pkg/util/json",
         "//pkg/util/protoutil",

pkg/sql/hintpb/statement_hint.go

@@ -6,6 +6,9 @@
 package hintpb
 
 import (
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/lexbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/protoreflect"
 	"github.com/cockroachdb/cockroach/pkg/util/json"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
@@ -44,6 +47,21 @@ func (h StatementHintUnion) HintType() string {
 	}
 }
 
+// RecreateStmt returns the SQL statement that can be used to recreate the hint.
+// Returns the empty string and false if the hint type is not supported.
+func (h StatementHintUnion) RecreateStmt(stmt string) (string, bool) {
+	switch t := h.GetValue().(type) {
+	case *InjectHints:
+		return fmt.Sprintf(
+			"SELECT information_schema.crdb_rewrite_inline_hints(%s, %s);",
+			lexbase.EscapeSQLString(stmt),
+			lexbase.EscapeSQLString(t.DonorSQL),
+		), true
+	default:
+		return "", false
+	}
+}
+
 // Details returns a JSON representation of the hint details. This is used for
 // displaying hint information in SHOW STATEMENT HINTS WITH DETAILS.
 func (h *StatementHintUnion) Details() (json.JSON, error) {

pkg/sql/opt/exec/execbuilder/testdata/system

@@ -173,3 +173,52 @@ distribution: local
       table: table_statistics_locks@primary
       spans: [/0/1 - /0/1]
       locking strength: for update
+
+# Query used by the stmt bundle to fetch hints for recreation.
+query T
+EXPLAIN (OPT, VERBOSE) SELECT hint FROM system.statement_hints WHERE fingerprint = 'SELECT * FROM t WHERE x = _' ORDER BY created_at
+----
+distribute
+ ├── columns: hint:4  [hidden: created_at:5]
+ ├── stats: [rows=10]
+ ├── cost: 291.694386
+ ├── cost-flags: unbounded-cardinality
+ ├── ordering: +5
+ ├── distribution: test
+ ├── input distribution:
+ ├── prune: (4,5)
+ └── sort
+      ├── columns: hint:4 created_at:5
+      ├── stats: [rows=10]
+      ├── cost: 91.6743864
+      ├── cost-flags: unbounded-cardinality
+      ├── ordering: +5
+      ├── prune: (4,5)
+      └── project
+           ├── columns: hint:4 created_at:5
+           ├── stats: [rows=10]
+           ├── cost: 90.5900007
+           ├── cost-flags: unbounded-cardinality
+           ├── prune: (4,5)
+           └── select
+                ├── columns: fingerprint:3 hint:4 created_at:5
+                ├── stats: [rows=10, distinct(3)=1, null(3)=0]
+                ├── cost: 90.4700007
+                ├── cost-flags: unbounded-cardinality
+                ├── fd: ()-->(3)
+                ├── index-join statement_hints
+                │    ├── columns: fingerprint:3 hint:4 created_at:5
+                │    ├── stats: [rows=10]
+                │    ├── cost: 90.3400007
+                │    ├── cost-flags: unbounded-cardinality
+                │    ├── interesting orderings: (+1)
+                │    └── scan statement_hints@hash_idx
+                │         ├── columns: row_id:1
+                │         ├── constraint: /2/1: [/6457443302007961978 - /6457443302007961978]
+                │         ├── stats: [rows=10, distinct(2)=1, null(2)=0]
+                │         ├── cost: 28.6200001
+                │         ├── cost-flags: unbounded-cardinality
+                │         ├── key: (1)
+                │         └── interesting orderings: (+1)
+                └── filters
+                     └── fingerprint:3 = 'SELECT * FROM t WHERE x = _' [outer=(3), constraints=(/3: [/'SELECT * FROM t WHERE x = _' - /'SELECT * FROM t WHERE x = _']; tight), fd=()-->(3)]

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
+}