sql: rerun execbuilder for gist matching

This commit re-runs the execbuilder with the explain factory whenever
plan-gist matching happens. This allows us to fully populate `plan.txt`
contents of the bundle. Some care had to be taken to ensure that the new
plan is used going forward (for correct annotation with exec stats). We
also make a copy of the memo to be safe (so that the original memo isn't
mutated by the execbuilder).

Release note: None

pkg/sql/conn_executor_exec.go

@@ -1863,7 +1863,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 		// the pausable portal, and we're not collecting a bundle yet, check
 		// whether we should get a bundle for this particular plan gist.
 		if ih := &planner.instrumentation; !ih.collectBundle && ih.outputMode == unmodifiedOutput {
-			ctx = ih.setupWithPlanGist(ctx, ex.server.cfg, stmt.StmtNoConstants, planGist, &planner.curPlan)
+			ctx = ih.setupWithPlanGist(ctx, planner, ex.server.cfg, planGist)
 		}
 	}
 

pkg/sql/explain_bundle_test.go

@@ -567,6 +567,7 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 
 	t.Run("plan-gist matching", func(t *testing.T) {
 		r.Exec(t, "CREATE TABLE gist (k INT PRIMARY KEY);")
+		r.Exec(t, "INSERT INTO gist SELECT generate_series(1, 10)")
 		const fprint = `SELECT * FROM gist`
 
 		// Come up with a target gist.
@@ -580,20 +581,20 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 				if name != "plan.txt" {
 					return nil
 				}
-				// Add a new line at the beginning for cleaner formatting in the
-				// test.
-				contents = "\n" + contents
-				// The gist appears to be somewhat non-deterministic (but its
-				// decoding stays the same), so we populate the expected
-				// contents based on the particular gist.
-				expected := fmt.Sprintf(`
--- plan is incomplete due to gist matching: %s
-
-• scan
-  table: gist@gist_pkey
-  spans: FULL SCAN`, gist)
-				if contents != expected {
-					return errors.Newf("unexpected contents of plan.txn\nexpected:\n%s\ngot:\n%s", expected, contents)
+				// We don't hard-code the full expected output here so that it
+				// doesn't need an update every time we change EXPLAIN ANALYZE
+				// output format. Instead, we only assert that a few lines are
+				// present in the output.
+				for _, expectedLine := range []string{
+					"• scan",
+					"  sql nodes: n1",
+					"  actual row count: 10",
+					"  table: gist@gist_pkey",
+					"  spans: FULL SCAN",
+				} {
+					if !strings.Contains(contents, expectedLine) {
+						return errors.Newf("didn't find %q in the output: %v", expectedLine, contents)
+					}
 				}
 				return nil
 			}, false, /* expectErrors */

pkg/sql/instrumentation.go

@@ -9,7 +9,6 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -29,7 +28,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec/execbuilder"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec/explain"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/indexrec"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
+	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
 	"github.com/cockroachdb/cockroach/pkg/sql/physicalplan"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -92,12 +94,6 @@ type instrumentationHelper struct {
 	// statement; it triggers saving of extra information like the plan string.
 	collectBundle bool
 
-	// planGistMatchingBundle is set when the bundle collection was enabled for
-	// a request with plan-gist matching enabled. In particular, such a bundle
-	// will be somewhat incomplete (it'll miss the plan string as well as the
-	// trace will miss all the events that happened in the optimizer).
-	planGistMatchingBundle bool
-
 	// collectExecStats is set when we are collecting execution statistics for a
 	// statement.
 	collectExecStats bool
@@ -534,16 +530,15 @@ func (ih *instrumentationHelper) Setup(
 // provided fingerprint and plan gist. It assumes that the bundle is not
 // currently being collected.
 func (ih *instrumentationHelper) setupWithPlanGist(
-	ctx context.Context, cfg *ExecutorConfig, fingerprint, planGist string, plan *planTop,
+	ctx context.Context, p *planner, cfg *ExecutorConfig, planGist string,
 ) context.Context {
 	ih.collectBundle, ih.diagRequestID, ih.diagRequest =
-		ih.stmtDiagnosticsRecorder.ShouldCollectDiagnostics(ctx, fingerprint, planGist)
+		ih.stmtDiagnosticsRecorder.ShouldCollectDiagnostics(ctx, p.stmt.StmtNoConstants, planGist)
 	// IsRedacted will be false when ih.collectBundle is false.
 	ih.explainFlags.RedactValues = ih.explainFlags.RedactValues || ih.diagRequest.IsRedacted()
 	if ih.collectBundle {
 		ih.needFinish = true
 		ih.collectExecStats = true
-		ih.planGistMatchingBundle = true
 		if ih.sp == nil || !ih.sp.IsVerbose() {
 			// We will create a verbose span
 			// - if we don't have a span yet, or
@@ -567,12 +562,52 @@ func (ih *instrumentationHelper) setupWithPlanGist(
 			)
 			ih.shouldFinishSpan = true
 			ih.finalizeSetup(ctx, cfg)
-			log.VEventf(ctx, 1, "plan-gist matching bundle collection began after the optimizer finished its part")
+		}
+		log.VEventf(ctx, 1, "plan-gist matching bundle collection began after the optimizer finished its part")
+		if cfg.TestingKnobs.DeterministicExplain {
+			ih.explainFlags.Deflake = explain.DeflakeAll
+		}
+		// Since we haven't enabled the bundle collection before the
+		// optimization, explain plan wasn't populated. We'll rerun the
+		// execbuilder with the explain factory to get that, on the copy of the
+		// memo (to be safe to not mutate the original memo).
+		copiedMemo := func(o *xform.Optimizer, mem *memo.Memo) *memo.Memo {
+			f := o.Factory()
+			f.CopyAndReplace(
+				mem.RootExpr().(memo.RelExpr),
+				mem.RootProps(),
+				f.CopyWithoutAssigningPlaceholders,
+			)
+			return f.DetachMemo()
+		}(&p.optPlanningCtx.optimizer, p.curPlan.mem)
+
+		explainFactory := explain.NewFactory(newExecFactory(ctx, p), p.SemaCtx(), p.EvalContext())
+		bld := execbuilder.New(
+			ctx, explainFactory, &p.optPlanningCtx.optimizer, copiedMemo,
+			p.curPlan.catalog, copiedMemo.RootExpr(),
+			p.SemaCtx(), p.EvalContext(), p.autoCommit, statements.IsANSIDML(p.stmt.AST),
+		)
+		// Disable telemetry in order to not double count things since we've
+		// already built the plan once.
+		bld.DisableTelemetry()
+		if plan, err := bld.Build(); err != nil {
+			log.VEventf(ctx, 1, "hit an error when using explain factory: %v", err)
+		} else {
+			ep := plan.(*explain.Plan)
+			ih.RecordExplainPlan(ep)
+			// We need to close the original plan since we're going to overwrite
+			// it. Note that the new plan will be closed correctly by the defer
+			// in dispatchToExecutionEngine.
+			p.curPlan.close(ctx)
+			// We need to use the newly-created plan going forward in order for
+			// execution stats to be attributed correctly (execNodeTraceMetadata
+			// is a map from pointers to plan nodes).
+			p.curPlan.planComponents = *ep.WrappedPlan.(*planComponents)
 		}
 	} else {
 		// We won't need the memo and the catalog, so free it up.
-		plan.mem = nil
-		plan.catalog = nil
+		p.curPlan.mem = nil
+		p.curPlan.catalog = nil
 	}
 	return ctx
 }
@@ -675,45 +710,10 @@ func (ih *instrumentationHelper) Finish(
 				}
 			}
 			planString := ob.BuildString()
-			if ih.planGistMatchingBundle {
-				// We don't have the plan string available since the stmt bundle
-				// collection was enabled _after_ the optimizer was done.
-				// Instead, we do have the gist available, so we'll decode it
-				// and use that as the plan string.
-				var sb strings.Builder
-				sb.WriteString("-- plan is incomplete due to gist matching: ")
-				sb.WriteString(ih.planGist.String())
-				// Perform best-effort decoding ignoring all errors.
-				if it, err := ie.QueryIterator(
-					bundleCtx, "plan-gist-decoding" /* opName */, nil, /* txn */
-					fmt.Sprintf("SELECT * FROM crdb_internal.decode_plan_gist('%s')", ih.planGist.String()),
-				); err == nil {
-					defer func() {
-						_ = it.Close()
-					}()
-					sb.WriteString("\n")
-					// Ignore the errors returned on Next call.
-					for ok, _ = it.Next(bundleCtx); ok; ok, _ = it.Next(bundleCtx) {
-						row := it.Cur()
-						var line string
-						// Be conservative in case the output format changes.
-						if len(row) == 1 {
-							var ds tree.DString
-							ds, ok = tree.AsDString(row[0])
-							line = string(ds)
-						} else {
-							ok = false
-						}
-						if !ok && buildutil.CrdbTestBuild {
-							return errors.AssertionFailedf("unexpected output format for decoding plan gist %s", ih.planGist.String())
-						}
-						if ok {
-							sb.WriteString("\n")
-							sb.WriteString(line)
-						}
-					}
-				}
-				planString = sb.String()
+			if planString == "" {
+				// This should only happen with plan-gist matching where we hit
+				// an error when using the explain factory.
+				planString = "-- plan is missing, probably hit an error with gist matching: " + ih.planGist.String()
 			}
 			bundle = buildStatementBundle(
 				bundleCtx, ih.explainFlags, cfg.DB, p, ie.(*InternalExecutor),

pkg/sql/opt/exec/execbuilder/builder.go

@@ -273,6 +273,11 @@ func New(
 	return b
 }
 
+// DisableTelemetry disables telemetry on this Builder.
+func (b *Builder) DisableTelemetry() {
+	b.disableTelemetry = true
+}
+
 // Build constructs the execution node tree and returns its root node if no
 // error occurred.
 func (b *Builder) Build() (_ exec.Plan, err error) {