sql: use the new EXPLAIN infrastructure for UI plans

This change reworks how the `roachpb.ExplainPlanTreeNode` tree (used by the UI
to show statement plans) is populated. We use the new EXPLAIN infrastructure.
The main change is that we need to decide upfront if we will need the plan or
not.

Release note (admin ui change): various improvements to the statement plans in
the UI.

pkg/sql/app_stats.go

@@ -292,13 +292,14 @@ func (a *appStats) recordTransaction(txnTimeSec float64, ev txnEvent, implicit b
 // sample logical plan for its corresponding fingerprint. We use
 // `logicalPlanCollectionPeriod` to assess how frequently to sample logical
 // plans.
-func (a *appStats) shouldSaveLogicalPlanDescription(
-	stmt *Statement, useDistSQL bool, vectorized bool, implicitTxn bool, err error,
-) bool {
+func (a *appStats) shouldSaveLogicalPlanDescription(stmt *Statement, implicitTxn bool) bool {
 	if !sampleLogicalPlans.Get(&a.st.SV) {
 		return false
 	}
-	stats := a.getStatsForStmt(stmt, implicitTxn, err, false /* createIfNonexistent */)
+	// We don't know yet if we will hit an error, so we assume we don't. The worst
+	// that can happen is that for statements that always error out, we will
+	// always save the tree plan.
+	stats := a.getStatsForStmt(stmt, implicitTxn, nil /* error */, false /* createIfNonexistent */)
 	if stats == nil {
 		// Save logical plan the first time we see new statement fingerprint.
 		return true

pkg/sql/conn_executor_exec.go

@@ -909,10 +909,12 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 // makeExecPlan creates an execution plan and populates planner.curPlan using
 // the cost-based optimizer.
 func (ex *connExecutor) makeExecPlan(ctx context.Context, planner *planner) error {
-	planner.curPlan.init(planner.stmt, ex.appStats)
-	if planner.collectBundle {
-		planner.curPlan.savePlanString = true
-	}
+	savePlanString := planner.collectBundle
+	planner.curPlan.init(
+		planner.stmt,
+		ex.appStats,
+		savePlanString,
+	)
 
 	if err := planner.makeOptimizerPlan(ctx); err != nil {
 		log.VEventf(ctx, 1, "optimizer plan failed: %v", err)

pkg/sql/executor_statement_metrics.go

@@ -159,7 +159,7 @@ func (ex *connExecutor) recordStatementSummary(
 	}
 
 	ex.statsCollector.recordStatement(
-		stmt, planner.curPlan.savedPlanForStats,
+		stmt, planner.curPlan.planForStats,
 		flags.IsDistributed(), flags.IsSet(planFlagVectorized),
 		flags.IsSet(planFlagImplicitTxn), automaticRetryCount, rowsAffected, err,
 		parseLat, planLat, runLat, svcLat, execOverhead, stats,

pkg/sql/explain_plan.go

@@ -30,10 +30,6 @@ const (
 	// explainSubqueryFmtFlags is the format for subqueries within `EXPLAIN SQL` statements.
 	// Since these are individually run, we don't need to scrub any data from subqueries.
 	explainSubqueryFmtFlags = tree.FmtSimple
-
-	// sampledLogicalPlanFmtFlags is the format for sampled logical plans. Because these exposed
-	// in the Admin UI, sampled plans should be scrubbed of sensitive information.
-	sampledLogicalPlanFmtFlags = tree.FmtHideConstants
 )
 
 // explainPlanNode wraps the logic for EXPLAIN as a planNode.

pkg/sql/explain_tree_test.go

@@ -66,16 +66,16 @@ func TestPlanToTreeAndPlanToString(t *testing.T) {
 
 			p.stmt = &Statement{Statement: stmt}
 			p.curPlan.savePlanString = true
+			p.curPlan.savePlanForStats = true
 			if err := p.makeOptimizerPlan(ctx); err != nil {
 				t.Fatal(err)
 			}
+			p.curPlan.flags.Set(planFlagExecDone)
+			p.curPlan.close(ctx)
 			if d.Cmd == "plan-string" {
-				p.curPlan.flags.Set(planFlagExecDone)
-				p.curPlan.close(ctx)
 				return p.curPlan.planString
 			}
-			tree := planToTree(ctx, &p.curPlan)
-			treeYaml, err := yaml.Marshal(tree)
+			treeYaml, err := yaml.Marshal(p.curPlan.planForStats)
 			if err != nil {
 				t.Fatal(err)
 			}

pkg/sql/logictest/testdata/logic_test/statement_statistics

@@ -120,7 +120,6 @@ SELECT sqrt(_)                                                    !
 SELECT x FROM (VALUES (_, _, __more1__), (__more1__)) AS t (x)    ·
 SELECT x FROM test WHERE y = (_ / z)                              !+
 SELECT x FROM test WHERE y IN (_, _, _ + x, _, _)                 ·
-SELECT x FROM test WHERE y IN (_, _, __more3__)                   ·
 SELECT x FROM test WHERE y IN (_, _, __more3__)                   +
 SELECT x FROM test WHERE y NOT IN (_, _, __more3__)               ·
 SET CLUSTER SETTING "debug.panic_on_failed_assertions" = DEFAULT  ·
@@ -147,7 +146,6 @@ SELECT _ FROM (VALUES (_, _, __more1__), (__more1__)) AS _ (_)
 SELECT _ FROM _ WHERE _ = (_ / _)
 SELECT _ FROM _ WHERE _ IN (_, _, _ + _, _, _)
 SELECT _ FROM _ WHERE _ IN (_, _, __more3__)
-SELECT _ FROM _ WHERE _ IN (_, _, __more3__)
 SELECT _ FROM _ WHERE _ NOT IN (_, _, __more3__)
 SET CLUSTER SETTING "debug.panic_on_failed_assertions" = DEFAULT
 SET CLUSTER SETTING "debug.panic_on_failed_assertions" = _

pkg/sql/opt/exec/explain/emit.go

@@ -641,7 +641,15 @@ func (e *emitter) emitSpans(
 			e.ob.Attr(field, "FULL SCAN")
 		}
 	} else {
-		e.ob.Attr(field, e.spanFormatFn(table, index, scanParams))
+		if e.ob.flags.HideValues {
+			n := len(scanParams.InvertedConstraint)
+			if scanParams.IndexConstraint != nil {
+				n = scanParams.IndexConstraint.Spans.Count()
+			}
+			e.ob.Attrf(field, "%d span%s", n, util.Pluralize(int64(n)))
+		} else {
+			e.ob.Attr(field, e.spanFormatFn(table, index, scanParams))
+		}
 	}
 }
 

pkg/sql/opt/exec/explain/flags.go

@@ -20,6 +20,10 @@ type Flags struct {
 	// ShowTypes indicates that the types of columns are shown.
 	// If ShowTypes is true, then Verbose is also true.
 	ShowTypes bool
+	// If HideValues is true, we hide fields that may contain values from the
+	// query (e.g. spans). Used internally for the plan visible in the UI.
+	// If HideValues is true, then Verbose must be false.
+	HideValues bool
 }
 
 // MakeFlags crates Flags from ExplainOptions.

pkg/sql/opt/exec/explain/output.go

@@ -116,6 +116,9 @@ func (ob *OutputBuilder) Expr(key string, expr tree.TypedExpr, varColumns sqlbas
 	if ob.flags.ShowTypes {
 		flags |= tree.FmtShowTypes
 	}
+	if ob.flags.HideValues {
+		flags |= tree.FmtHideConstants
+	}
 	f := tree.NewFmtCtx(flags)
 	f.SetIndexedVarFormat(func(ctx *tree.FmtCtx, idx int) {
 		// Ensure proper quoting.

pkg/sql/opt/exec/explain/result_columns.go

@@ -196,7 +196,7 @@ func getResultColumns(
 		}
 		return sqlbase.ShowTraceColumns, nil
 
-	case createTableOp, createTableAsOp, createViewOp, controlJobsOp,
+	case createTableOp, createTableAsOp, createViewOp, controlJobsOp, controlSchedulesOp,
 		cancelQueriesOp, cancelSessionsOp, errorIfRowsOp, deleteRangeOp:
 		// These operations produce no columns.
 		return nil, nil

pkg/sql/plan.go

@@ -12,6 +12,7 @@ package sql
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
@@ -25,7 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
-	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 )
 
 // runParams is a struct containing all parameters passed to planNode.Next() and
@@ -301,14 +302,18 @@ type planTop struct {
 	// If we are collecting query diagnostics, flow diagrams are saved here.
 	distSQLDiagrams []execinfrapb.FlowDiagram
 
-	appStats          *appStats
-	savedPlanForStats *roachpb.ExplainTreePlanNode
+	// If savePlanForStats is true, an ExplainTreePlanNode tree will be saved in
+	// planForStats when the plan is closed.
+	savePlanForStats bool
+	// appStats is used to populate savePlanForStats.
+	appStats     *appStats
+	planForStats *roachpb.ExplainTreePlanNode
 
 	// If savePlanString is set to true, an EXPLAIN (VERBOSE)-style plan string
-	// will be saved in planString.
+	// will be saved in planString when the plan is closed.
 	savePlanString bool
-	explainPlan    *explain.Plan
 	planString     string
+	explainPlan    *explain.Plan
 }
 
 // physicalPlanTop is a utility wrapper around PhysicalPlan that allows for
@@ -419,40 +424,24 @@ func (p *planComponents) close(ctx context.Context) {
 
 // init resets planTop to point to a given statement; used at the start of the
 // planning process.
-func (p *planTop) init(stmt *Statement, appStats *appStats) {
-	*p = planTop{stmt: stmt, appStats: appStats}
+func (p *planTop) init(stmt *Statement, appStats *appStats, savePlanString bool) {
+	*p = planTop{
+		stmt:           stmt,
+		appStats:       appStats,
+		savePlanString: savePlanString,
+	}
 }
 
 // close ensures that the plan's resources have been deallocated.
 func (p *planTop) close(ctx context.Context) {
-	if p.main.planNode != nil && p.flags.IsSet(planFlagExecDone) {
-		// TODO(yuzefovich): update this once we support creating table reader
-		// specs directly in the optimizer (see #47474).
-		if p.appStats != nil && p.appStats.shouldSaveLogicalPlanDescription(
-			p.stmt,
-			p.flags.IsDistributed(),
-			p.flags.IsSet(planFlagVectorized),
-			p.flags.IsSet(planFlagImplicitTxn),
-			p.execErr,
-		) {
-			p.savedPlanForStats = planToTree(ctx, p)
-		}
-
-		if p.savePlanString && p.explainPlan != nil {
-			var err error
-			p.planString, err = p.formatExplain()
-			if err != nil {
-				p.planString = err.Error()
-			}
-		}
+	if p.explainPlan != nil && p.flags.IsSet(planFlagExecDone) {
+		p.savePlanInfo(ctx)
 	}
 	p.planComponents.close(ctx)
 }
 
-func (p *planTop) formatExplain() (string, error) {
-	if p.explainPlan == nil {
-		return "", errors.AssertionFailedf("no plan")
-	}
+// savePlanInfo uses p.explainPlan to populate the plan string and/or tree.
+func (p *planTop) savePlanInfo(ctx context.Context) {
 	vectorized := p.flags.IsSet(planFlagVectorized)
 	distribution := physicalplan.LocalPlan
 	if p.flags.IsSet(planFlagFullyDistributed) {
@@ -461,14 +450,28 @@ func (p *planTop) formatExplain() (string, error) {
 		distribution = physicalplan.PartiallyDistributedPlan
 	}
 
-	ob := explain.NewOutputBuilder(explain.Flags{
-		Verbose:   true,
-		ShowTypes: true,
-	})
-	if err := emitExplain(ob, p.codec, p.explainPlan, distribution, vectorized); err != nil {
-		return "", err
+	if p.savePlanForStats {
+		ob := explain.NewOutputBuilder(explain.Flags{
+			HideValues: true,
+		})
+		if err := emitExplain(ob, p.codec, p.explainPlan, distribution, vectorized); err != nil {
+			log.Warningf(ctx, "unable to emit explain plan tree: %v", err)
+		} else {
+			p.planForStats = ob.BuildProtoTree()
+		}
+	}
+
+	if p.savePlanString {
+		ob := explain.NewOutputBuilder(explain.Flags{
+			Verbose:   true,
+			ShowTypes: true,
+		})
+		if err := emitExplain(ob, p.codec, p.explainPlan, distribution, vectorized); err != nil {
+			p.planString = fmt.Sprintf("error emitting plan: %v", err)
+		} else {
+			p.planString = ob.BuildString()
+		}
 	}
-	return ob.BuildString(), nil
 }
 
 // formatOptPlan returns a visual representation of the optimizer plan that was

pkg/sql/plan_opt.go

@@ -552,7 +552,16 @@ func (opc *optPlanningCtx) runExecBuilder(
 	var result *planComponents
 	var explainPlan *explain.Plan
 	var isDDL bool
-	if !planTop.savePlanString {
+	if planTop.appStats != nil {
+		// We do not set this flag upfront when initializing planTop because the
+		// planning process could in principle modify the AST, resulting in a
+		// different statement signature.
+		planTop.savePlanForStats = planTop.appStats.shouldSaveLogicalPlanDescription(
+			planTop.stmt,
+			allowAutoCommit,
+		)
+	}
+	if !planTop.savePlanString && !planTop.savePlanForStats {
 		// No instrumentation.
 		bld := execbuilder.New(f, mem, &opc.catalog, mem.RootExpr(), evalCtx, allowAutoCommit)
 		plan, err := bld.Build()

pkg/sql/sem/tree/hide_constants.go

@@ -118,9 +118,9 @@ func (node *Tuple) formatHideConstants(ctx *FmtCtx) {
 		v2.Exprs = append(make(Exprs, 0, 3), v2.Exprs[:2]...)
 		if len(node.Exprs) > 2 {
 			v2.Exprs = append(v2.Exprs, arityIndicator(len(node.Exprs)-2))
-		}
-		if node.Labels != nil {
-			v2.Labels = node.Labels[:2]
+			if node.Labels != nil {
+				v2.Labels = node.Labels[:2]
+			}
 		}
 		v2.Format(ctx)
 		return