improve implementation to avoid starting nested traversals

felipe-gdr · felipe-gdr · commit 00bc31770c99 · 2022-06-29T14:47:50.000+10:00
diff --git a/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java b/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java
@@ -3,6 +3,7 @@
 import graphql.PublicApi;
 import graphql.schema.GraphQLEnumType;
 import graphql.schema.GraphQLFieldDefinition;
+import graphql.schema.GraphQLImplementingType;
 import graphql.schema.GraphQLInputObjectField;
 import graphql.schema.GraphQLInterfaceType;
 import graphql.schema.GraphQLNamedSchemaElement;
@@ -14,20 +15,22 @@
 import graphql.schema.GraphQLTypeVisitorStub;
 import graphql.schema.GraphQLUnionType;
 import graphql.schema.SchemaTraverser;
+import graphql.schema.impl.SchemaUtil;
 import graphql.schema.transform.VisibleFieldPredicateEnvironment.VisibleFieldPredicateEnvironmentImpl;
 import graphql.util.TraversalControl;
 import graphql.util.TraverserContext;
 
-import java.util.Collections;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static graphql.schema.SchemaTransformer.transformSchema;
-import static graphql.util.TreeTransformerUtil.deleteNode;
 
 /**
  * Transforms a schema by applying a visibility predicate to every field.
@@ -63,13 +66,13 @@ public final GraphQLSchema apply(GraphQLSchema schema) {
 
         beforeTransformationHook.run();
 
-        new SchemaTraverser().depthFirst(new TypeObservingVisitor(observedBeforeTransform, schema), getRootTypes(schema));
+        new SchemaTraverser(getChildrenFn(schema)).depthFirst(new TypeObservingVisitor(observedBeforeTransform), getRootTypes(schema));
 
         // remove fields
         GraphQLSchema interimSchema = transformSchema(schema,
                 new FieldRemovalVisitor(visibleFieldPredicate, markedForRemovalTypes));
 
-        new SchemaTraverser().depthFirst(new TypeObservingVisitor(observedAfterTransform, interimSchema), getRootTypes(interimSchema));
+        new SchemaTraverser(getChildrenFn(interimSchema)).depthFirst(new TypeObservingVisitor(observedAfterTransform), getRootTypes(interimSchema));
 
         // remove types that are not used after removing fields - (connected schema only)
         GraphQLSchema connectedSchema = transformSchema(interimSchema,
@@ -84,6 +87,23 @@ public final GraphQLSchema apply(GraphQLSchema schema) {
         return finalSchema;
     }
 
+    // Creates a getChildrenFn that includes interface
+    private Function<GraphQLSchemaElement, List<GraphQLSchemaElement>> getChildrenFn(GraphQLSchema schema) {
+        Map<String, List<GraphQLImplementingType>> interfaceImplementations = new SchemaUtil().groupImplementationsForInterfacesAndObjects(schema);
+
+        return graphQLSchemaElement -> {
+            if (!(graphQLSchemaElement instanceof GraphQLInterfaceType)) {
+                return graphQLSchemaElement.getChildren();
+            }
+            ArrayList<GraphQLSchemaElement> children = new ArrayList<>(graphQLSchemaElement.getChildren());
+            List<GraphQLImplementingType> implementations = interfaceImplementations.get(((GraphQLInterfaceType) graphQLSchemaElement).getName());
+            if (implementations != null) {
+                children.addAll(implementations);
+            }
+            return children;
+        };
+    }
+
     private GraphQLSchema removeUnreferencedTypes(Set<GraphQLType> markedForRemovalTypes, GraphQLSchema connectedSchema) {
         GraphQLSchema withoutAdditionalTypes = connectedSchema.transform(builder -> {
             Set<GraphQLType> additionalTypes = new HashSet<>(connectedSchema.getAdditionalTypes());
@@ -110,12 +130,10 @@ protected TraversalControl visitGraphQLType(GraphQLSchemaElement node, Traverser
     private static class TypeObservingVisitor extends GraphQLTypeVisitorStub {
 
         private final Set<GraphQLType> observedTypes;
-        private GraphQLSchema graphQLSchema;
 
 
-        private TypeObservingVisitor(Set<GraphQLType> observedTypes, GraphQLSchema graphQLSchema) {
+        private TypeObservingVisitor(Set<GraphQLType> observedTypes) {
             this.observedTypes = observedTypes;
-            this.graphQLSchema = graphQLSchema;
         }
 
         @Override
@@ -124,21 +142,6 @@ protected TraversalControl visitGraphQLType(GraphQLSchemaElement node,
             if (node instanceof GraphQLType) {
                 observedTypes.add((GraphQLType) node);
             }
-            if (node instanceof GraphQLInterfaceType) {
-                final GraphQLSchemaElement parentNode = context.getParentNode();
-                if(parentNode instanceof GraphQLObjectType && observedTypes.contains(parentNode)) {
-                    // This means the traversal reached this interface via a type that implements it. If that type
-                    // has already been observed, we don't need to continue traversing so we quit early.
-                    // This is just to avoid the streaming/filtering bellow
-                    return TraversalControl.QUIT;
-                }
-
-                final List<GraphQLObjectType> implementations = graphQLSchema.getImplementations((GraphQLInterfaceType) node);
-
-                implementations.stream()
-                        .filter(impl -> !observedTypes.contains(impl))
-                        .forEach(impl -> new SchemaTraverser().depthFirst(this, Collections.singleton(impl)));
-            }
 
             return TraversalControl.CONTINUE;
         }
diff --git a/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy b/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy
@@ -1134,4 +1134,48 @@ class FieldVisibilitySchemaTransformationTest extends Specification {
         restrictedSchema.getType("BillingStatus") != null
     }
 
+    def "handles types that become visible via types reachable by interface that implements interface"() {
+        given:
+        GraphQLSchema schema = TestUtil.schema("""
+
+        directive @private on FIELD_DEFINITION
+
+        type Query {
+            account: Account
+            node: Node
+        }
+        
+        type Account {
+            name: String
+            billingStatus: BillingStatus @private
+        }
+        
+        type BillingStatus {
+            accountNumber: String
+        }
+       
+        interface Node {
+            id: ID!
+        } 
+        
+        interface NamedNode implements Node {
+            id: ID!
+            name: String 
+        }
+        
+        type Billing implements Node & NamedNode {
+            id: ID!
+            name: String
+            status: BillingStatus
+        }
+        
+        """)
+
+        when:
+        GraphQLSchema restrictedSchema = visibilitySchemaTransformation.apply(schema)
+
+        then:
+        (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
+        restrictedSchema.getType("BillingStatus") != null
+    }
 }
