remove query process

Author: James Cor

engine.go

@@ -447,7 +447,7 @@ func (e *Engine) QueryWithBindings(ctx *sql.Context, query string, parsed sqlpar
 		return nil, nil, nil, err
 	}
 
-	iter = finalizeIters(analyzed, qFlags, iter)
+	iter = finalizeIters(ctx, analyzed, qFlags, iter)
 
 	return analyzed.Schema(), iter, qFlags, nil
 }
@@ -482,7 +482,7 @@ func (e *Engine) PrepQueryPlanForExecution(ctx *sql.Context, _ string, plan sql.
 		return nil, nil, nil, err
 	}
 
-	iter = finalizeIters(plan, nil, iter)
+	iter = finalizeIters(ctx, plan, nil, iter)
 
 	return plan.Schema(), iter, nil, nil
 }
@@ -830,7 +830,7 @@ func (e *Engine) executeEvent(ctx *sql.Context, dbName, createEventStatement, us
 		return err
 	}
 
-	iter = finalizeIters(definitionNode, nil, iter)
+	iter = finalizeIters(ctx, definitionNode, nil, iter)
 
 	// Drain the iterate to execute the event body/definition
 	// NOTE: No row data is returned for an event; we just need to execute the statements
@@ -865,8 +865,9 @@ func findCreateEventNode(planTree sql.Node) (*plan.CreateEvent, error) {
 }
 
 // finalizeIters applies the final transformations on sql.RowIter before execution.
-func finalizeIters(analyzed sql.Node, qFlags *sql.QueryFlags, iter sql.RowIter) sql.RowIter {
-	iter = rowexec.AddTransactionCommittingIter(iter, qFlags)
+func finalizeIters(ctx *sql.Context, analyzed sql.Node, qFlags *sql.QueryFlags, iter sql.RowIter) sql.RowIter {
+	iter = plan.AddTrackedRowIter(ctx, analyzed, iter)
+	iter = rowexec.AddTransactionCommittingIter(qFlags, iter)
 	iter = rowexec.AddExpressionCloser(analyzed, iter)
 	return iter
 }

enginetest/engine_only_test.go

@@ -293,74 +293,6 @@ b (2/6 partitions)
 	require.ElementsMatch(expected, rows)
 }
 
-// TODO: this was an analyzer test, but we don't have a mock process list for it to use, so it has to be here
-func TestTrackProcess(t *testing.T) {
-	require := require.New(t)
-	db := memory.NewDatabase("db")
-	provider := memory.NewDBProvider(db)
-	a := analyzer.NewDefault(provider)
-	sess := memory.NewSession(sql.NewBaseSession(), provider)
-
-	node := plan.NewInnerJoin(
-		plan.NewResolvedTable(&nonIndexableTable{memory.NewPartitionedTable(db.BaseDatabase, "foo", sql.PrimaryKeySchema{}, nil, 2)}, nil, nil),
-		plan.NewResolvedTable(memory.NewPartitionedTable(db.BaseDatabase, "bar", sql.PrimaryKeySchema{}, nil, 4), nil, nil),
-		expression.NewLiteral(int64(1), types.Int64),
-	)
-
-	pl := sqle.NewProcessList()
-
-	ctx := sql.NewContext(context.Background(), sql.WithPid(1), sql.WithProcessList(pl), sql.WithSession(sess))
-	pl.AddConnection(ctx.Session.ID(), "localhost")
-	pl.ConnectionReady(ctx.Session)
-	ctx, err := ctx.ProcessList.BeginQuery(ctx, "SELECT foo")
-	require.NoError(err)
-
-	rule := getRuleFrom(analyzer.OnceAfterAll, analyzer.TrackProcessId)
-	result, _, err := rule.Apply(ctx, a, node, nil, analyzer.DefaultRuleSelector, nil)
-	require.NoError(err)
-
-	processes := ctx.ProcessList.Processes()
-	require.Len(processes, 1)
-	require.Equal("SELECT foo", processes[0].Query)
-	require.Equal(
-		map[string]sql.TableProgress{
-			"foo": sql.TableProgress{
-				Progress:           sql.Progress{Name: "foo", Done: 0, Total: 2},
-				PartitionsProgress: map[string]sql.PartitionProgress{}},
-			"bar": sql.TableProgress{
-				Progress:           sql.Progress{Name: "bar", Done: 0, Total: 4},
-				PartitionsProgress: map[string]sql.PartitionProgress{}},
-		},
-		processes[0].Progress)
-
-	proc, ok := result.(*plan.QueryProcess)
-	require.True(ok)
-
-	join, ok := proc.Child().(*plan.JoinNode)
-	require.True(ok)
-	require.Equal(join.JoinType(), plan.JoinTypeInner)
-
-	lhs, ok := join.Left().(*plan.ResolvedTable)
-	require.True(ok)
-	_, ok = lhs.Table.(*plan.ProcessTable)
-	require.True(ok)
-
-	rhs, ok := join.Right().(*plan.ResolvedTable)
-	require.True(ok)
-	_, ok = rhs.Table.(*plan.ProcessTable)
-	require.True(ok)
-
-	iter, err := rowexec.DefaultBuilder.Build(ctx, proc, nil)
-	require.NoError(err)
-	_, err = sql.RowIterToRows(ctx, iter)
-	require.NoError(err)
-
-	procs := ctx.ProcessList.Processes()
-	require.Len(procs, 1)
-	require.Equal(procs[0].Command, sql.ProcessCommandSleep)
-	require.Error(ctx.Err())
-}
-
 func TestConcurrentProcessList(t *testing.T) {
 	enginetest.TestConcurrentProcessList(t, enginetest.NewDefaultMemoryHarness())
 }
@@ -515,7 +447,6 @@ func TestAnalyzer_Exp(t *testing.T) {
 			require.NoError(t, err)
 
 			analyzed, err := e.EngineAnalyzer().Analyze(ctx, parsed, nil, nil)
-			analyzed = analyzer.StripPassthroughNodes(analyzed)
 			if tt.err != nil {
 				require.Error(t, err)
 				assert.True(t, tt.err.Is(err))
@@ -527,10 +458,6 @@ func TestAnalyzer_Exp(t *testing.T) {
 }
 
 func assertNodesEqualWithDiff(t *testing.T, expected, actual sql.Node) {
-	if x, ok := actual.(*plan.QueryProcess); ok {
-		actual = x.Child()
-	}
-
 	if !assert.Equal(t, expected, actual) {
 		expectedStr := sql.DebugString(expected)
 		actualStr := sql.DebugString(actual)

sql/analyzer/describe.go

@@ -32,5 +32,5 @@ func resolveDescribeQuery(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan
 		return nil, transform.SameTree, err
 	}
 
-	return d.WithQuery(StripPassthroughNodes(q)), transform.NewTree, nil
+	return d.WithQuery(q), transform.NewTree, nil
 }

sql/analyzer/inserts.go

@@ -62,8 +62,6 @@ func resolveInsertRows(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Sc
 			if err != nil {
 				return nil, transform.SameTree, err
 			}
-
-			source = StripPassthroughNodes(source)
 		}
 
 		dstSchema := insertable.Schema()

sql/analyzer/parallelize.go

@@ -63,8 +63,7 @@ func parallelize(ctx *sql.Context, a *Analyzer, node sql.Node, scope *plan.Scope
 		return node, transform.SameTree, nil
 	}
 
-	proc, ok := node.(*plan.QueryProcess)
-	if (ok && !shouldParallelize(proc.Child(), nil)) || !shouldParallelize(node, scope) {
+	if !shouldParallelize(node, scope) {
 		return node, transform.SameTree, nil
 	}
 

sql/analyzer/process.go

@@ -39,15 +39,10 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope,
 	if !n.Resolved() {
 		return n, transform.SameTree, nil
 	}
-
-	if _, ok := n.(*plan.QueryProcess); ok {
-		return n, transform.SameTree, nil
-	}
-
 	processList := ctx.ProcessList
 
 	var seen = make(map[string]struct{})
-	n, _, err := transform.Node(n, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
+	n, same, err := transform.Node(n, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
 		switch n := n.(type) {
 		case *plan.ResolvedTable:
 			switch n.Table.(type) {
@@ -106,41 +101,9 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope,
 			return n, transform.SameTree, nil
 		}
 	})
-	if err != nil {
-		return nil, transform.SameTree, err
-	}
-
-	// Don't wrap CreateIndex in a QueryProcess, as it is a CreateIndexProcess.
-	// CreateIndex will take care of marking the process as done on its own.
-	if _, ok := n.(*plan.CreateIndex); ok {
-		return n, transform.SameTree, nil
-	}
 
-	// Remove QueryProcess nodes from the subqueries and trigger bodies. Otherwise, the process
-	// will be marked as done as soon as a subquery / trigger finishes.
-	node, _, err := transform.Node(n, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
-		if sq, ok := n.(*plan.SubqueryAlias); ok {
-			if qp, ok := sq.Child.(*plan.QueryProcess); ok {
-				n, err := sq.WithChildren(qp.Child())
-				return n, transform.NewTree, err
-			}
-		}
-		if t, ok := n.(*plan.TriggerExecutor); ok {
-			if qp, ok := t.Right().(*plan.QueryProcess); ok {
-				n, err := t.WithChildren(t.Left(), qp.Child())
-				return n, transform.NewTree, err
-			}
-		}
-		return n, transform.SameTree, nil
-	})
 	if err != nil {
 		return nil, transform.SameTree, err
 	}
-
-	return plan.NewQueryProcess(node, func() {
-		processList.EndQuery(ctx)
-		if span := ctx.RootSpan(); span != nil {
-			span.End()
-		}
-	}), transform.NewTree, nil
+	return n, same, nil
 }

sql/analyzer/process_test.go

@@ -55,7 +55,7 @@ func resolveCreateSelect(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.
 		return nil, transform.SameTree, err
 	}
 
-	return plan.NewTableCopier(ct.Database(), StripPassthroughNodes(analyzedCreate), StripPassthroughNodes(analyzedSelect), plan.CopierProps{}), transform.NewTree, nil
+	return plan.NewTableCopier(ct.Database(), analyzedCreate, analyzedSelect, plan.CopierProps{}), transform.NewTree, nil
 }
 
 // stripSchema removes all non-type information from a schema, such as the key info, default value, etc.

sql/analyzer/resolve_create_select.go

@@ -252,7 +252,7 @@ func analyzeSubqueryExpression(ctx *sql.Context, a *Analyzer, n sql.Node, sq *pl
 	// to the expense of positive errors, where a rule reports a change when the plan
 	// is the same before/after.
 	// .Resolved() might be useful for fixing these bugs.
-	return sq.WithQuery(StripPassthroughNodes(analyzed)).WithExecBuilder(a.ExecBuilder), transform.NewTree, nil
+	return sq.WithQuery(analyzed).WithExecBuilder(a.ExecBuilder), transform.NewTree, nil
 }
 
 // analyzeSubqueryAlias runs analysis on the specified subquery alias, |sqa|. The |finalize| parameter indicates if this is
@@ -282,28 +282,10 @@ func analyzeSubqueryAlias(ctx *sql.Context, a *Analyzer, sqa *plan.SubqueryAlias
 	if same {
 		return sqa, transform.SameTree, nil
 	}
-	newn, err := sqa.WithChildren(StripPassthroughNodes(child))
+	newn, err := sqa.WithChildren(child)
 	return newn, transform.NewTree, err
 }
 
-// StripPassthroughNodes strips all top-level passthrough nodes meant to apply only to top-level queries (query
-// tracking, transaction logic, etc) from the node tree given and return the first non-passthrough child element. This
-// is useful for when we invoke the analyzer recursively when e.g. analyzing subqueries or triggers
-// TODO: instead of stripping this node off after analysis, it would be better to just not add it in the first place.
-func StripPassthroughNodes(n sql.Node) sql.Node {
-	nodeIsPassthrough := true
-	for nodeIsPassthrough {
-		switch tn := n.(type) {
-		case *plan.QueryProcess:
-			n = tn.Child()
-		default:
-			nodeIsPassthrough = false
-		}
-	}
-
-	return n
-}
-
 // cacheSubqueryAlisesInJoins will look for joins against subquery aliases that
 // will repeatedly execute the subquery, and will insert a *plan.CachedResults
 // node on top of those nodes. The left-most child of a join root is an exception

sql/analyzer/resolve_subqueries.go

@@ -51,7 +51,7 @@ func resolveUnions(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope,
 			return nil, transform.SameTree, err
 		}
 
-		ret, err := n.WithChildren(StripPassthroughNodes(left), StripPassthroughNodes(right))
+		ret, err := n.WithChildren(left, right)
 		if err != nil {
 			return nil, transform.SameTree, err
 		}
@@ -95,7 +95,7 @@ func finalizeUnions(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope
 
 		scope.SetJoin(false)
 
-		newn, err := n.WithChildren(StripPassthroughNodes(left), StripPassthroughNodes(right))
+		newn, err := n.WithChildren(left, right)
 		if err != nil {
 			return nil, transform.SameTree, err
 		}