Gen4: Handling subquery in query graph

go/vt/vtgate/planbuilder/plan_test.go

@@ -149,6 +149,8 @@ func init() {
 	vindexes.Register("costly", newCostlyIndex)
 }
 
+const samePlanMarker = "Gen4 plan same as above\n"
+
 func TestPlan(t *testing.T) {
 	vschemaWrapper := &vschemaWrapper{
 		v:             loadSchema(t, "schema_test.json"),
@@ -418,9 +420,6 @@ func testFile(t *testing.T, filename, tempDir string, vschema *vschemaWrapper, c
 			if tcase.output2ndPlanner == "" {
 				empty = true
 			}
-			if tcase.output2ndPlanner == "{\n}\n" {
-				tcase.output2ndPlanner = tcase.output
-			}
 
 			vschema.version = V4
 			out, err := getPlanOutput(tcase, vschema)
@@ -444,7 +443,7 @@ func testFile(t *testing.T, filename, tempDir string, vschema *vschemaWrapper, c
 					}
 
 					if tcase.output == out {
-						expected.WriteString("{\n}\n")
+						expected.WriteString(samePlanMarker)
 					} else {
 						expected.WriteString(out)
 					}
@@ -555,7 +554,9 @@ func iterateExecFile(name string) (testCaseIterator chan testCase) {
 			if err != nil && err != io.EOF {
 				panic(fmt.Sprintf("error reading file %s line# %d: %s", name, lineno, err.Error()))
 			}
-			if len(binput) > 0 && (binput[0] == '"' || binput[0] == '{') {
+			if len(binput) > 0 && string(binput) == samePlanMarker {
+				output2Planner = output
+			} else if len(binput) > 0 && (binput[0] == '"' || binput[0] == '{') {
 				output2Planner = append(output2Planner, binput...)
 				for {
 					l, err := r.ReadBytes('\n')
@@ -573,9 +574,6 @@ func iterateExecFile(name string) (testCaseIterator chan testCase) {
 						break
 					}
 				}
-				if string(output2Planner) == "{\n}" {
-					output2Planner = output
-				}
 			}
 
 			testCaseIterator <- testCase{

go/vt/vtgate/planbuilder/querygraph.go

@@ -41,6 +41,9 @@ type (
 
 		// noDeps contains the predicates that can be evaluated anywhere.
 		noDeps sqlparser.Expr
+
+		// subqueries contains the subqueries that depend on this query graph
+		subqueries map[*sqlparser.Subquery][]*queryGraph
 	}
 
 	// queryTable is a single FROM table, including all predicates particular to this table
@@ -79,9 +82,38 @@ func createQGFromSelect(sel *sqlparser.Select, semTable *semantics.SemTable) (*q
 	return qg, nil
 }
 
+func createQGFromSelectStatement(selStmt sqlparser.SelectStatement, semTable *semantics.SemTable) ([]*queryGraph, error) {
+	switch stmt := selStmt.(type) {
+	case *sqlparser.Select:
+		qg, err := createQGFromSelect(stmt, semTable)
+		if err != nil {
+			return nil, err
+		}
+		return []*queryGraph{qg}, err
+	case *sqlparser.Union:
+		qg, err := createQGFromSelectStatement(stmt.FirstStatement, semTable)
+		if err != nil {
+			return nil, err
+		}
+		for _, sel := range stmt.UnionSelects {
+			qgr, err := createQGFromSelectStatement(sel.Statement, semTable)
+			if err != nil {
+				return nil, err
+			}
+			qg = append(qg, qgr...)
+		}
+		return qg, nil
+	case *sqlparser.ParenSelect:
+		return createQGFromSelectStatement(stmt.Select, semTable)
+	}
+
+	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "BUG: not reachable %T", selStmt)
+}
+
 func newQueryGraph() *queryGraph {
 	return &queryGraph{
 		crossTable: map[semantics.TableSet][]sqlparser.Expr{},
+		subqueries: map[*sqlparser.Subquery][]*queryGraph{},
 	}
 }
 
@@ -156,7 +188,20 @@ func (qg *queryGraph) collectPredicate(predicate sqlparser.Expr, semTable *seman
 		}
 		qg.crossTable[deps] = allPredicates
 	}
-	return nil
+	err := sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		switch subQuery := node.(type) {
+		case *sqlparser.Subquery:
+
+			qgr, err := createQGFromSelectStatement(subQuery.Select, semTable)
+			if err != nil {
+				return false, err
+			}
+			qg.subqueries[subQuery] = qgr
+		}
+		return true, nil
+	}, predicate)
+
+	return err
 }
 
 func (qg *queryGraph) addToSingleTable(table semantics.TableSet, predicate sqlparser.Expr) bool {