sql: use correct context in recursive CTE iterations

Previously, the `execbuilder.Builder` that we use in recursive CTE
iterations referenced the same context that the Builder captured on the
main query path. This is problematic since that context might have
a tracing span that's already been finished. This commit fixes this
issue by explicitly passing the context argument into the iteration
function. This was only exposed because we added some logging in the
following commit but has been present for a while now.

Release note: None

Author: yuzefovich

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

@@ -3418,7 +3418,6 @@ func (b *Builder) buildRecursiveCTE(
 	// To implement exec.RecursiveCTEIterationFn, we create a special Builder.
 
 	innerBldTemplate := &Builder{
-		ctx:     b.ctx,
 		mem:     b.mem,
 		catalog: b.catalog,
 		semaCtx: b.semaCtx,
@@ -3430,9 +3429,10 @@ func (b *Builder) buildRecursiveCTE(
 		withExprs: b.withExprs[:len(b.withExprs):len(b.withExprs)],
 	}
 
-	fn := func(ef exec.Factory, bufferRef exec.Node) (exec.Plan, error) {
+	fn := func(ctx context.Context, ef exec.Factory, bufferRef exec.Node) (exec.Plan, error) {
 		// Use a separate builder each time.
 		innerBld := *innerBldTemplate
+		innerBld.ctx = ctx
 		innerBld.factory = ef
 		innerBld.addBuiltWithExpr(rec.WithID, initialCols, bufferRef)
 		// TODO(mgartner): I think colOrdsAlloc can be reused for each recursive

pkg/sql/opt/exec/factory.go

@@ -285,7 +285,7 @@ type KVOption struct {
 // RecursiveCTEIterationFn creates a plan for an iteration of WITH RECURSIVE,
 // given the result of the last iteration (as a node created by
 // ConstructBuffer).
-type RecursiveCTEIterationFn func(ef Factory, bufferRef Node) (Plan, error)
+type RecursiveCTEIterationFn func(ctx context.Context, ef Factory, bufferRef Node) (Plan, error)
 
 // ApplyJoinPlanRightSideFn creates a plan for an iteration of ApplyJoin, given
 // a row produced from the left side. The plan is guaranteed to produce the

pkg/sql/recursive_cte.go

@@ -135,17 +135,21 @@ func (n *recursiveCTENode) Next(params runParams) (bool, error) {
 		rows:                lastWorkingRows,
 		label:               n.label,
 	}
-	newPlan, err := n.genIterationFn(newExecFactory(params.ctx, params.p), buf)
-	if err != nil {
-		return false, err
-	}
 
+	// Create a separate tracing span that will be used when planning and
+	// running this iteration. Note that we'll still use the "outer" params.ctx
+	// when accessing rows in the container.
 	n.iterationCount++
 	opName := "recursive-cte-iteration-" + strconv.Itoa(n.iterationCount)
-	ctx, sp := tracing.ChildSpan(params.ctx, opName)
+	planAndRunCtx, sp := tracing.ChildSpan(params.ctx, opName)
 	defer sp.Finish()
+
+	newPlan, err := n.genIterationFn(planAndRunCtx, newExecFactory(planAndRunCtx, params.p), buf)
+	if err != nil {
+		return false, err
+	}
 	if err := runPlanInsidePlan(
-		ctx, params, newPlan.(*planComponents), rowResultWriter(n),
+		planAndRunCtx, params, newPlan.(*planComponents), rowResultWriter(n),
 		nil /* deferredRoutineSender */, "", /* stmtForDistSQLDiagram */
 	); err != nil {
 		return false, err