raw draft

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,56 @@ 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`
+
+	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()
+	}
+	// 1. Fingerprint the statement
+	// 2. Query all the doner sql
+	// 3. Convert them into rewrite inline hints statements.
+	return donorSqls, nil
+}
+
+const printStmtHintPattern = `SELECT information_schema.crdb_rewrite_inline_hints('%s', '%s')\n`
+
+func (c *stmtEnvCollector) PrintStmtHints(
+	ctx context.Context, stmt string, sv *settings.Values, w io.Writer,
+) error {
+	donerSqls, err := c.getStatementHints(ctx, stmt, sv)
+	if err != nil {
+		return err
+	}
+	for _, donorSql := range donerSqls {
+		fmt.Fprintf(w, printStmtHintPattern, 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 +1058,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")
 	}