PR feedback - added getImmediate fields

bbakerman · bbakerman · commit 6401f124bf5d · 2020-10-22T15:17:41.000+11:00
diff --git a/src/main/java/graphql/schema/DataFetchingFieldSelectionSet.java b/src/main/java/graphql/schema/DataFetchingFieldSelectionSet.java
@@ -146,6 +146,16 @@ public interface DataFetchingFieldSelectionSet {
      */
     List<SelectedField> getFields();
 
+    /**
+     * This will return all selected fields that are immediate child fields
+     * of the field being fetched.
+     * <p>
+     * The fields are guaranteed to be in pre-order as they appear in the query.
+     *
+     * @return a list of all selected immediate child fields or empty list if none match
+     */
+    List<SelectedField> getImmediateFields();
+
     /**
      * This will return a list of selected fields that match a specified "glob" pattern matching ie
      * the glob pattern matching supported by {@link java.nio.file.FileSystem#getPathMatcher}.
diff --git a/src/main/java/graphql/schema/DataFetchingFieldSelectionSetImpl.java b/src/main/java/graphql/schema/DataFetchingFieldSelectionSetImpl.java
@@ -51,6 +51,11 @@ public List<SelectedField> getFields() {
             return emptyList();
         }
 
+        @Override
+        public List<SelectedField> getImmediateFields() {
+            return emptyList();
+        }
+
         @Override
         public List<SelectedField> getFields(String fieldGlobPattern, String... fieldGlobPatterns) {
             return Collections.emptyList();
@@ -79,6 +84,7 @@ public static DataFetchingFieldSelectionSet newCollector(GraphQLOutputType field
     private final Supplier<NormalizedField> normalizedFieldSupplier;
 
     private boolean computedValues;
+    private List<SelectedField> immediateFields;
     private Map<String, List<SelectedField>> normalisedSelectionSetFields;
     private Set<String> flattenedFieldsForGlobSearching;
 
@@ -159,6 +165,12 @@ public List<SelectedField> getFields() {
                 .collect(toList());
     }
 
+    @Override
+    public List<SelectedField> getImmediateFields() {
+        computeValuesLazily();
+        return immediateFields;
+    }
+
     @Override
     public Map<String, List<SelectedField>> getFieldsGroupedByResultKey() {
         return getFields().stream().collect(Collectors.groupingBy(SelectedField::getResultKey));
@@ -179,13 +191,14 @@ private void computeValuesLazily() {
             }
             flattenedFieldsForGlobSearching = new LinkedHashSet<>();
             normalisedSelectionSetFields = new LinkedHashMap<>();
-            traverseSubSelectedFields(normalizedFieldSupplier.get(), "", "");
+            immediateFields = new ArrayList<>();
+            traverseSubSelectedFields(normalizedFieldSupplier.get(), "", "", true);
             computedValues = true;
         }
     }
 
 
-    private void traverseSubSelectedFields(NormalizedField currentNormalisedField, String qualifiedFieldPrefix, String simpleFieldPrefix) {
+    private void traverseSubSelectedFields(NormalizedField currentNormalisedField, String qualifiedFieldPrefix, String simpleFieldPrefix, boolean firstLevel) {
         List<NormalizedField> children = currentNormalisedField.getChildren();
         for (NormalizedField normalizedSubSelectedField : children) {
 
@@ -200,13 +213,16 @@ private void traverseSubSelectedFields(NormalizedField currentNormalisedField, S
             flattenedFieldsForGlobSearching.add(globSimpleName);
 
             SelectedFieldImpl selectedField = new SelectedFieldImpl(globSimpleName, globQualifiedName, normalizedSubSelectedField);
+            if (firstLevel) {
+                immediateFields.add(selectedField);
+            }
             normalisedSelectionSetFields.computeIfAbsent(globQualifiedName, newList()).add(selectedField);
             normalisedSelectionSetFields.computeIfAbsent(globSimpleName, newList()).add(selectedField);
 
             GraphQLFieldDefinition fieldDefinition = normalizedSubSelectedField.getFieldDefinition();
             GraphQLType unwrappedType = GraphQLTypeUtil.unwrapAll(fieldDefinition.getType());
             if (!GraphQLTypeUtil.isLeaf(unwrappedType)) {
-                traverseSubSelectedFields(normalizedSubSelectedField, globQualifiedName, globSimpleName);
+                traverseSubSelectedFields(normalizedSubSelectedField, globQualifiedName, globSimpleName, false);
             }
         }
     }
diff --git a/src/test/groovy/graphql/schema/DataFetchingFieldSelectionSetImplTest.groovy b/src/test/groovy/graphql/schema/DataFetchingFieldSelectionSetImplTest.groovy
@@ -740,10 +740,47 @@ class DataFetchingFieldSelectionSetImplTest extends Specification {
 
         then:
         selectedFields.size() == 3
-        assertTheyAreExpected(selectedFields,["Bird.name", "Cat.name", "Dog.name"])
+        assertTheyAreExpected(selectedFields, ["Bird.name", "Cat.name", "Dog.name"])
 
     }
 
+    def "immediate fields can be found at all levels"() {
+        when:
+        def query = '''
+        {
+            pet {
+                name
+                ... on Dog {
+                    woof
+                }
+                ... on Cat {
+                    meow
+                }
+                lead {
+                    material
+                }
+            }
+         }
+        '''
+        def ei = ExecutionInput.newExecutionInput(query).build()
+        def er = petSchema.execute(ei)
+        then:
+        er.errors.isEmpty()
+        when:
+        def selectedFields = petSelectionSet.getImmediateFields()
+        then:
+        assertTheyAreExpected(selectedFields, [
+                "Bird.lead", "Bird.name",
+                "Cat.lead", "Cat.meow", "Cat.name",
+                "Dog.lead", "Dog.name", "Dog.woof"])
+
+        when:
+        selectedFields = petSelectionSet.getFields("lead")
+        selectedFields = selectedFields[0].getSelectionSet().getImmediateFields()
+        then:
+        assertTheyAreExpected(selectedFields, ["Lead.material"])
+    }
+
     def "simple fields do not have selection sets"() {
         when:
         def query = '''
