sql/analyzer: remove QueryProcess nodes from subqueries

Closes src-d/gitbase#679

Because subqueries are analyzed as if they were a whole independent
query QueryProcess nodes were added to them. That caused that as
soon as a subquery finished, the process were marked as done. As a
result, some queries could error with "context canceled", as the
context of the query had already been canceled.

This fix removes the QueryProcess nodes from any subquery and makes
sure only the root node is wrapped in a QueryProcess.

Signed-off-by: Miguel Molina <[email protected]>

Author: erizocosmico

sql/analyzer/process.go

@@ -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
 }

sql/analyzer/process_test.go

@@ -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:]