graphql-java · bbakerman · Sep 7, 2022 · Sep 7, 2022
diff --git a/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java b/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java
@@ -1,10 +1,12 @@
 package graphql.schema.transform;
 
+import com.google.common.collect.ImmutableList;
 import graphql.PublicApi;
 import graphql.schema.GraphQLEnumType;
 import graphql.schema.GraphQLFieldDefinition;
 import graphql.schema.GraphQLImplementingType;
 import graphql.schema.GraphQLInputObjectField;
+import graphql.schema.GraphQLInputObjectType;
 import graphql.schema.GraphQLInterfaceType;
 import graphql.schema.GraphQLNamedSchemaElement;
 import graphql.schema.GraphQLNamedType;
@@ -60,7 +62,7 @@ public final GraphQLSchema apply(GraphQLSchema schema) {
         Set<GraphQLType> markedForRemovalTypes = new HashSet<>();
 
         // query, mutation, and subscription types should not be removed
-        final Set<String> protectedTypeNames = getRootTypes(schema).stream()
+        final Set<String> protectedTypeNames = getOperationTypes(schema).stream()
                 .map(GraphQLObjectType::getName)
                 .collect(Collectors.toSet());
 
@@ -216,6 +218,7 @@ public TraversalControl visitGraphQLType(GraphQLSchemaElement node,
                     !observedAfterTransform.contains(node) &&
                     (node instanceof GraphQLObjectType ||
                             node instanceof GraphQLEnumType ||
+                            node instanceof GraphQLInputObjectType ||
                             node instanceof GraphQLInterfaceType ||
                             node instanceof GraphQLUnionType)) {
 
@@ -250,12 +253,21 @@ public TraversalControl visitGraphQLType(GraphQLSchemaElement node,
         }
     }
 
-    private List<GraphQLObjectType> getRootTypes(GraphQLSchema schema) {
+    private List<GraphQLSchemaElement> getRootTypes(GraphQLSchema schema) {
+        return ImmutableList.<GraphQLSchemaElement>builder()
+                .addAll(getOperationTypes(schema))
+                // Include directive definitions as roots, since they won't be removed in the filtering process.
+                // Some types (enums, input types, etc.) might be reachable only by directive definitions (and
+                // not by other types or fields).
+                .addAll(schema.getDirectives())
+                .build();
+    }
+
+    private List<GraphQLObjectType> getOperationTypes(GraphQLSchema schema) {
         return Stream.of(
                 schema.getQueryType(),
                 schema.getSubscriptionType(),
                 schema.getMutationType()
         ).filter(Objects::nonNull).collect(Collectors.toList());
     }
-
 }
diff --git a/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy b/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy
@@ -1178,4 +1178,67 @@ class FieldVisibilitySchemaTransformationTest extends Specification {
         (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
         restrictedSchema.getType("BillingStatus") != null
     }
+
+    def "can remove a field with a directive containing enum argument"() {
+        given:
+        GraphQLSchema schema = TestUtil.schema("""
+
+        directive @private(privateType: SecretType) on FIELD_DEFINITION
+        enum SecretType {
+            SUPER_SECRET
+            NOT_SO_SECRET
+        }
+
+        type Query {
+            account: Account
+        }
+
+        type Account {
+            name: String
+            billingStatus: BillingStatus @private(privateType: NOT_SO_SECRET)
+        }
+
+        type BillingStatus {
+            accountNumber: String
+        }
+        """)
+
+        when:
+        GraphQLSchema restrictedSchema = visibilitySchemaTransformation.apply(schema)
+
+        then:
+        (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
+        restrictedSchema.getType("BillingStatus") == null
+    }
+
+    def "can remove a field with a directive containing type argument"() {
+        given:
+        GraphQLSchema schema = TestUtil.schema("""
+
+        directive @private(privateType: SecretType) on FIELD_DEFINITION
+        input SecretType {
+            description: String
+        }
+
+        type Query {
+            account: Account
+        }
+
+        type Account {
+            name: String
+            billingStatus: BillingStatus @private(privateType: { description: "secret" })
+        }
+
+        type BillingStatus {
+            accountNumber: String
+        }
+        """)
+
+        when:
+        GraphQLSchema restrictedSchema = visibilitySchemaTransformation.apply(schema)
+
+        then:
+        (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
+        restrictedSchema.getType("BillingStatus") == null
+    }
 }