src-d · erizocosmico · Jan 23, 2019 · Jan 23, 2019
@@ -71,5 +71,19 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 		return n, nil
 	}
 
-	return plan.NewQueryProcess(n, func() { processList.Done(ctx.Pid()) }), nil
+	// Remove QueryProcess nodes from the subqueries. Otherwise, the process
+	// will be marked as done as soon as a subquery finishes.
+	node, err := n.TransformUp(func(n sql.Node) (sql.Node, error) {
+		if sq, ok := n.(*plan.SubqueryAlias); ok {
+			if qp, ok := sq.Child.(*plan.QueryProcess); ok {
+				return plan.NewSubqueryAlias(sq.Name(), qp.Child), nil
+			}
+		}
+		return n, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return plan.NewQueryProcess(node, func() { processList.Done(ctx.Pid()) }), nil
 }
@@ -70,6 +70,37 @@ func TestTrackProcess(t *testing.T) {
 	}
 }
 
+func TestTrackProcessSubquery(t *testing.T) {
+	require := require.New(t)
+	rule := getRuleFrom(OnceAfterAll, "track_process")
+	catalog := sql.NewCatalog()
+	a := NewDefault(catalog)
+
+	node := plan.NewProject(
+		nil,
+		plan.NewSubqueryAlias("f",
+			plan.NewQueryProcess(
+				plan.NewResolvedTable(mem.NewTable("foo", nil)),
+				nil,
+			),
+		),
+	)
+
+	result, err := rule.Apply(sql.NewEmptyContext(), a, node)
+	require.NoError(err)
+
+	expectedChild := plan.NewProject(
+		nil,
+		plan.NewSubqueryAlias("f",
+			plan.NewResolvedTable(mem.NewTable("foo", nil)),
+		),
+	)
+
+	proc, ok := result.(*plan.QueryProcess)
+	require.True(ok)
+	require.Equal(expectedChild, proc.Child)
+}
+
 func withoutProcessTracking(a *Analyzer) *Analyzer {
 	afterAll := a.Batches[len(a.Batches)-1]
 	afterAll.Rules = afterAll.Rules[1:]