Merge pull request #85 from gmac/fix-child-step-merging

Fix: fully traverse nested field selections while planning

plan.go

@@ -2,6 +2,7 @@ package bramble
 
 import (
 	"context"
+	"strings"
 	"encoding/json"
 	"fmt"
 
@@ -117,6 +118,7 @@ func createSteps(ctx *PlanningContext, insertionPoint []string, parentType, pare
 func extractSelectionSet(ctx *PlanningContext, insertionPoint []string, parentType string, input ast.SelectionSet, location string, childstep bool) (ast.SelectionSet, []*QueryPlanStep, error) {
 	var selectionSetResult []ast.Selection
 	var childrenStepsResult []*QueryPlanStep
+	var remoteSelections []ast.Selection
 	for _, selection := range input {
 		switch selection := selection.(type) {
 		case *ast.Field:
@@ -136,44 +138,29 @@ func extractSelectionSet(ctx *PlanningContext, insertionPoint []string, parentTy
 				childrenStepsResult = append(childrenStepsResult, steps...)
 				continue
 			}
-			if loc == location {
-				if selection.SelectionSet == nil {
-					selectionSetResult = append(selectionSetResult, selection)
-				} else {
-					newField := *selection
-					selectionSet, childrenSteps, err := extractSelectionSet(
-						ctx,
-						append(insertionPoint, selection.Alias),
-						selection.Definition.Type.Name(),
-						selection.SelectionSet,
-						location,
-						childstep,
-					)
-					if err != nil {
-						return nil, nil, err
-					}
-					newField.SelectionSet = selectionSet
-					selectionSetResult = append(selectionSetResult, &newField)
-					childrenStepsResult = append(childrenStepsResult, childrenSteps...)
-				}
+			if loc != location {
+				// field transitions to another service location
+				remoteSelections = append(remoteSelections, selection)
+			} else if selection.SelectionSet == nil {
+				// field is a leaf type in the current service
+				selectionSetResult = append(selectionSetResult, selection)
 			} else {
-				mergedWithExistingStep := false
-				for _, step := range childrenStepsResult {
-					if stringArraysEqual(step.InsertionPoint, insertionPoint) && step.ServiceURL == loc {
-						step.SelectionSet = append(step.SelectionSet, selection)
-						mergedWithExistingStep = true
-						break
-					}
-				}
-
-				if !mergedWithExistingStep {
-					newSelectionSet := []ast.Selection{selection}
-					childrenSteps, err := createSteps(ctx, insertionPoint, parentType, location, newSelectionSet, true)
-					if err != nil {
-						return nil, nil, err
-					}
-					childrenStepsResult = append(childrenStepsResult, childrenSteps...)
+				// field is a composite type in the current service
+				newField := *selection
+				selectionSet, childrenSteps, err := extractSelectionSet(
+					ctx,
+					append(insertionPoint, selection.Alias),
+					selection.Definition.Type.Name(),
+					selection.SelectionSet,
+					location,
+					childstep,
+				)
+				if err != nil {
+					return nil, nil, err
 				}
+				newField.SelectionSet = selectionSet
+				selectionSetResult = append(selectionSetResult, &newField)
+				childrenStepsResult = append(childrenStepsResult, childrenSteps...)
 			}
 		case *ast.InlineFragment:
 			selectionSet, childrenSteps, err := extractSelectionSet(
@@ -220,6 +207,32 @@ func extractSelectionSet(ctx *PlanningContext, insertionPoint []string, parentTy
 		}
 	}
 
+	if len(remoteSelections) > 0 {
+		// Create child steps for all remote field selections
+		childrenSteps, err := createSteps(ctx, insertionPoint, parentType, location, remoteSelections, true)
+		if err != nil {
+			return nil, nil, err
+		}
+		childrenStepsResult = append(childrenStepsResult, childrenSteps...)
+	}
+
+	if len(childrenStepsResult) > 1 {
+		// Merge steps targeting distinct service/path locations
+		mergedSteps := []*QueryPlanStep{}
+		mergedStepsMap := map[string]*QueryPlanStep{}
+		for _, step := range childrenStepsResult {
+			key := strings.Join(append([]string{step.ServiceURL}, step.InsertionPoint...), "/")
+			if existingStep, ok := mergedStepsMap[key]; ok {
+				existingStep.SelectionSet = append(existingStep.SelectionSet, step.SelectionSet...)
+				existingStep.Then = append(existingStep.Then, step.Then...)
+			} else {
+				mergedStepsMap[key] = step
+				mergedSteps = append(mergedSteps, step)
+			}
+		}
+		childrenStepsResult = mergedSteps
+	}
+
 	parentDef := ctx.Schema.Types[parentType]
 	// For abstract types, add an id fragment for all possible boundary
 	// implementations. This assures that abstract boundaries always return
@@ -246,6 +259,7 @@ func extractSelectionSet(ctx *PlanningContext, insertionPoint []string, parentTy
 					ObjectDefinition: implementationType,
 				}
 				selectionSetResult = append([]ast.Selection{possibleId}, selectionSetResult...)
+				break
 			}
 		}
 	// Otherwise, add an id selection to boundary types where the result
@@ -382,20 +396,6 @@ func (m FieldURLMap) keyFor(parent string, field string) string {
 	return fmt.Sprintf("%s.%s", parent, field)
 }
 
-func stringArraysEqual(a, b []string) bool {
-	if len(a) != len(b) {
-		return false
-	}
-
-	for i, v := range a {
-		if v != b[i] {
-			return false
-		}
-	}
-
-	return true
-}
-
 // BoundaryField contains the name and format for a boundary query
 type BoundaryField struct {
 	Field string

plan_fixtures_test.go

@@ -220,6 +220,47 @@ var PlanTestFixture5 = &PlanTestFixture{
 	},
 }
 
+var PlanTestFixture6 = &PlanTestFixture{
+	Schema: `
+	type Shop {
+		id: ID!
+		name: String!
+		products: [Product]!
+	}
+
+	type Product {
+		id: ID!
+		name: String!
+		collection: Collection
+	}
+
+	type Collection {
+		id: ID!
+		name: String!
+	}
+
+	type Query {
+		shop1: Shop!
+	}`,
+
+	Locations: map[string]string{
+		"Query.shop1":         "A",
+		"Shop.id":             "A",
+		"Shop.name":           "A",
+		"Shop.products":       "A",
+		"Product.name":        "B",
+		"Product.collection":  "B",
+		"Collection.name":     "C",
+	},
+
+	IsBoundary: map[string]bool{
+		"Shop":        false,
+		"Product":     true,
+		"Collection":  true,
+	},
+}
+
+
 func (f *PlanTestFixture) Check(t *testing.T, query, expectedJSON string) {
 	t.Helper()
 	schema := gqlparser.MustLoadSchema(&ast.Source{Name: "fixture", Input: f.Schema})

plan_test.go

@@ -122,14 +122,14 @@ func TestQueryPlanABA2(t *testing.T) {
 					"ServiceURL": "A",
 					"ParentType": "Movie",
 					"SelectionSet": "{ _id: id title }",
-					"InsertionPoint": ["movies", "compTitles"],
+					"InsertionPoint": ["movies", "compTitles", "compTitles"],
 					"Then": null
 				  },
 				  {
 					"ServiceURL": "A",
 					"ParentType": "Movie",
 					"SelectionSet": "{ _id: id title }",
-					"InsertionPoint": ["movies", "compTitles", "compTitles"],
+					"InsertionPoint": ["movies", "compTitles"],
 					"Then": null
 				  }
 				]
@@ -408,6 +408,83 @@ func TestQueryPlanFragmentSpread2(t *testing.T) {
 	PlanTestFixture1.Check(t, query, plan)
 }
 
+func TestQueryPlanCompleteDeepTraversal(t *testing.T) {
+	query := `
+	{
+		shop1 {
+			name
+			products {
+				name
+				collection {
+					name
+				}
+			}
+		}
+	}`
+	plan := `{
+		"RootSteps": [
+			{
+				"ServiceURL": "A",
+				"ParentType": "Query",
+				"SelectionSet": "{ shop1 { name products { _id: id } } }",
+				"InsertionPoint": null,
+				"Then": [
+					{
+					"ServiceURL": "B",
+					"ParentType": "Product",
+					"SelectionSet": "{ _id: id name collection { _id: id } }",
+					"InsertionPoint": ["shop1", "products"],
+					"Then": [
+							{
+							"ServiceURL": "C",
+							"ParentType": "Collection",
+							"SelectionSet": "{ _id: id name }",
+							"InsertionPoint": ["shop1", "products", "collection"],
+							"Then": null
+							}
+						]
+					}
+				]
+			}
+		]
+	}`
+	PlanTestFixture6.Check(t, query, plan)
+}
+
+func TestQueryPlanMergeInsertionPointSteps(t *testing.T) {
+	query := `
+	{
+		shop1 {
+			products {
+				name
+			}
+			products {
+				name
+			}
+		}
+	}`
+	plan := `{
+		"RootSteps": [
+			{
+				"ServiceURL": "A",
+				"ParentType": "Query",
+				"SelectionSet": "{ shop1 { products { _id: id } products { _id: id } } }",
+				"InsertionPoint": null,
+				"Then": [
+					{
+					"ServiceURL": "B",
+					"ParentType": "Product",
+					"SelectionSet": "{ _id: id name _id: id name }",
+					"InsertionPoint": ["shop1", "products"],
+					"Then": null
+					}
+				]
+			}
+		]
+	}`
+	PlanTestFixture6.Check(t, query, plan)
+}
+
 func TestQueryPlanExpandAbstractTypesWithPossibleBoundaryIds(t *testing.T) {
 	query := `
 	{