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

@@ -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,68 @@ var stmtBundleIncludeAllFKReferences = settings.RegisterBoolSetting(
 	false,
 )
 
+// getDonorSqls returns the string representation of all the donor sql statements
+// of the target statement. The returned donor sqls should have been sorted in
+// ascending order by their creation time.
+func (c *stmtEnvCollector) getDonorSqls(
+	ctx context.Context, stmt string, sv *settings.Values,
+) (donorSqls []string, err error) {
+	fingerprint, err := parserutils.FingerprintStatement(stmt,
+		tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(
+			sv,
+		)))
+	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, 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
+		}
+		strDatum, ok := row[0].(*tree.DString)
+		if !ok {
+			return nil, errors.AssertionFailedf("expect row to be a string, got %T", row[0])
+		}
+		donorSqls = append(donorSqls, strDatum.String())
+	}
+	return donorSqls, nil
+}
+
+func (c *stmtEnvCollector) PrintStmtHints(
+	ctx context.Context, stmt string, sv *settings.Values, w io.Writer,
+) error {
+	donorSqls, err := c.getDonorSqls(ctx, stmt, sv)
+	if err != nil {
+		return err
+	}
+	for _, donorSql := range donorSqls {
+		// donorSql has been quoted. Escape stmt as well to handle embedded quotes.
+		fmt.Fprintf(w, "SELECT information_schema.crdb_rewrite_inline_hints(%s, %s);\n", lexbase.EscapeSQLString(stmt), donorSql)
+	}
+	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 +1070,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

@@ -1178,6 +1178,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/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
+}