Merge #153689

153689: sql, tree: add InjectHints function r=DrewKimball a=michae2

This PR adds `tree.(*HintInjectionDonor).InjectHints`, which rewrites a target statement with hints copied from a donor statement.

This is based on the last commit of [Drew's plan_hints prototype](#119120) but uses `tree.WalkStmt` instead of creating a new AST-walker.

See individual commits for details.

Informs: #153633

Release note: None

Co-authored-by: Michael Erickson <[email protected]>

Author: craig[bot]

pkg/internal/sqlsmith/relational.go

@@ -200,6 +200,17 @@ func makeJoinExpr(s *Smither, refs colRefs, forJoin bool) (tree.TableExpr, colRe
 		Left:     left,
 		Right:    right,
 	}
+	switch s.rnd.Intn(10) {
+	// INVERTED is unlikely to work unless we get very lucky, so skip it for now.
+	case 0:
+		joinExpr.Hint = "HASH"
+	case 1:
+		joinExpr.Hint = "LOOKUP"
+	case 2:
+		joinExpr.Hint = "MERGE"
+	case 3:
+		joinExpr.Hint = "STRAIGHT"
+	}
 
 	if s.disableCrossJoins {
 		if available, ok := getAvailablePairedColsForJoinPreds(s, leftRefs, rightRefs); ok {

pkg/sql/conn_executor.go

@@ -3643,23 +3643,6 @@ var allowBufferedWritesForWeakIsolation = settings.RegisterBoolSetting(
 
 var logIsolationLevelLimiter = log.Every(10 * time.Second)
 
-// Bitmask for enabling various query fingerprint formatting styles.
-// We don't publish this setting because most users should not need
-// to tweak the fingerprint generation.
-var queryFormattingForFingerprintsMask = settings.RegisterIntSetting(
-	settings.ApplicationLevel,
-	"sql.stats.statement_fingerprint.format_mask",
-	"enables setting additional fmt flags for statement fingerprint formatting. "+
-		"Flags set here will be applied in addition to FmtHideConstants",
-	int64(tree.FmtCollapseLists|tree.FmtConstantsAsUnderscores),
-	settings.WithValidateInt(func(i int64) error {
-		if i == 0 || int64(tree.FmtCollapseLists|tree.FmtConstantsAsUnderscores)&i == i {
-			return nil
-		}
-		return errors.Newf("invalid value %d", i)
-	}),
-)
-
 func (ex *connExecutor) txnIsolationLevelToKV(
 	ctx context.Context, level tree.IsolationLevel,
 ) isolation.Level {

pkg/sql/conn_executor_exec.go

@@ -380,7 +380,7 @@ func (ex *connExecutor) execStmtInOpenState(
 	os := ex.machine.CurState().(stateOpen)
 
 	isExtendedProtocol := prepared != nil
-	stmtFingerprintFmtMask := tree.FmtHideConstants | tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV))
+	stmtFingerprintFmtMask := tree.FmtHideConstants | tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV))
 
 	var stmt Statement
 	if isExtendedProtocol {
@@ -905,7 +905,7 @@ func (ex *connExecutor) execStmtInOpenState(
 				NumAnnotations:  stmt.NumAnnotations,
 			},
 			ex.server.cfg.GenerateID(),
-			tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV)),
+			tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV)),
 		)
 		var rawTypeHints []oid.Oid
 
@@ -1256,7 +1256,7 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 	os := ex.machine.CurState().(stateOpen)
 
 	isExtendedProtocol := portal != nil && portal.Stmt != nil
-	stmtFingerprintFmtMask := tree.FmtHideConstants | tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV))
+	stmtFingerprintFmtMask := tree.FmtHideConstants | tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV))
 
 	if isExtendedProtocol {
 		vars.stmt = makeStatementFromPrepared(portal.Stmt, queryID)
@@ -1867,7 +1867,7 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 				NumAnnotations:  vars.stmt.NumAnnotations,
 			},
 			ex.server.cfg.GenerateID(),
-			tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV)),
+			tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV)),
 		)
 		var rawTypeHints []oid.Oid
 
@@ -3453,7 +3453,7 @@ func (ex *connExecutor) execStmtInNoTxnState(
 
 		p := &ex.planner
 		stmt := makeStatement(parserStmt, ex.server.cfg.GenerateID(),
-			tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV)))
+			tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV)))
 		p.stmt = stmt
 		p.semaCtx.Annotations = tree.MakeAnnotations(stmt.NumAnnotations)
 		p.extendedEvalCtx.Annotations = &p.semaCtx.Annotations

pkg/sql/conn_executor_prepare.go

@@ -73,7 +73,7 @@ func (ex *connExecutor) execPrepare(
 	}
 
 	stmt := makeStatement(parseCmd.Statement, ex.server.cfg.GenerateID(),
-		tree.FmtFlags(queryFormattingForFingerprintsMask.Get(ex.server.cfg.SV())))
+		tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(ex.server.cfg.SV())))
 	_, err := ex.addPreparedStmt(
 		ctx,
 		parseCmd.Name,

pkg/sql/distsql_plan_changefeed.go

@@ -82,7 +82,7 @@ func PlanCDCExpression(
 		AST: cdcExpr,
 		SQL: tree.AsString(cdcExpr),
 	}, clusterunique.ID{}, /* queryID */
-		tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&p.execCfg.Settings.SV)),
+		tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&p.execCfg.Settings.SV)),
 	)
 
 	p.curPlan.init(&p.stmt, &p.instrumentation)

pkg/sql/distsql_running_test.go

@@ -168,7 +168,7 @@ func TestDistSQLRunningInAbortedTxn(t *testing.T) {
 		// We need to re-plan every time, since the plan is closed automatically
 		// by PlanAndRun() below making it unusable across retries.
 		p.stmt = makeStatement(stmt, clusterunique.ID{},
-			tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&execCfg.Settings.SV)))
+			tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&execCfg.Settings.SV)))
 		if err := p.makeOptimizerPlan(ctx); err != nil {
 			t.Fatal(err)
 		}
@@ -286,7 +286,7 @@ func TestDistSQLRunningParallelFKChecksAfterAbort(t *testing.T) {
 		)
 
 		p.stmt = makeStatement(stmt, clusterunique.ID{},
-			tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&s.ClusterSettings().SV)))
+			tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&s.ClusterSettings().SV)))
 		if err := p.makeOptimizerPlan(ctx); err != nil {
 			t.Fatal(err)
 		}

pkg/sql/exec_log.go

@@ -353,7 +353,7 @@ func (p *planner) maybeLogStatementInternal(
 		if stmtFingerprintID == 0 {
 			repQuery := p.stmt.StmtNoConstants
 			if repQuery == "" {
-				flags := tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&p.execCfg.Settings.SV))
+				flags := tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&p.execCfg.Settings.SV))
 				f := tree.NewFmtCtx(flags)
 				f.FormatNode(p.stmt.AST)
 				repQuery = f.CloseAndGetString()

pkg/sql/explain_tree_test.go

@@ -72,7 +72,7 @@ func TestPlanToTreeAndPlanToString(t *testing.T) {
 			ih.collectBundle = true
 
 			p.stmt = makeStatement(stmt, clusterunique.ID{},
-				tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&execCfg.Settings.SV)))
+				tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&execCfg.Settings.SV)))
 			if err := p.makeOptimizerPlan(ctx); err != nil {
 				t.Fatal(err)
 			}

pkg/sql/plan_opt_test.go

@@ -646,7 +646,7 @@ func TestPlanGistControl(t *testing.T) {
 	)
 	defer cleanup()
 
-	fmtFingerprintMask := tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&s.ClusterSettings().SV))
+	fmtFingerprintMask := tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&s.ClusterSettings().SV))
 	p := internalPlanner.(*planner)
 	stmt, err := parser.ParseOne("SELECT 1")
 	if err != nil {

pkg/sql/schema_changer.go

@@ -504,7 +504,7 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 		localPlanner.MaybeReallocateAnnotations(stmt.NumAnnotations)
 		// Construct an optimized logical plan of the AS source stmt.
 		localPlanner.stmt = makeStatement(stmt, clusterunique.ID{}, /* queryID */
-			tree.FmtFlags(queryFormattingForFingerprintsMask.Get(&localPlanner.execCfg.Settings.SV)))
+			tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&localPlanner.execCfg.Settings.SV)))
 		localPlanner.optPlanningCtx.init(localPlanner)
 
 		localPlanner.runWithOptions(resolveFlags{skipCache: true}, func() {