From 447ced850c47b0c845fc5d091521141109d372f0 Mon Sep 17 00:00:00 2001
From: Anastasios Basoukos <abasoukos@groupon.com>
Date: Sat, 29 Jul 2017 17:44:12 +0200
Subject: [PATCH] Allow schemas to contain custom directives for queries.

---
 .../java/graphql/schema/GraphQLSchema.java    | 20 +++++-
 .../ValidateCustomDirectives.groovy           | 61 +++++++++++++++++++
 2 files changed, 78 insertions(+), 3 deletions(-)
 create mode 100644 src/test/groovy/graphql/validation/ValidateCustomDirectives.groovy

diff --git a/src/main/java/graphql/schema/GraphQLSchema.java b/src/main/java/graphql/schema/GraphQLSchema.java
index 83b165e308..0677340ba8 100644
--- a/src/main/java/graphql/schema/GraphQLSchema.java
+++ b/src/main/java/graphql/schema/GraphQLSchema.java
@@ -10,6 +10,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -23,6 +24,7 @@ public class GraphQLSchema {
     private final GraphQLObjectType subscriptionType;
     private final Map<String, GraphQLType> typeMap;
     private Set<GraphQLType> additionalTypes;
+    private Set<GraphQLDirective> directives;
 
     public GraphQLSchema(GraphQLObjectType queryType) {
         this(queryType, null, Collections.emptySet());
@@ -33,12 +35,19 @@ public GraphQLSchema(GraphQLObjectType queryType, GraphQLObjectType mutationType
     }
 
     public GraphQLSchema(GraphQLObjectType queryType, GraphQLObjectType mutationType, GraphQLObjectType subscriptionType, Set<GraphQLType> dictionary) {
+        this(queryType, mutationType, subscriptionType, dictionary, Collections.emptySet());
+    }
+
+    public GraphQLSchema(GraphQLObjectType queryType, GraphQLObjectType mutationType, GraphQLObjectType subscriptionType, Set<GraphQLType> dictionary, Set<GraphQLDirective> directives) {
         assertNotNull(dictionary, "dictionary can't be null");
         assertNotNull(queryType, "queryType can't be null");
+        assertNotNull(directives, "directives can't be null");
         this.queryType = queryType;
         this.mutationType = mutationType;
         this.subscriptionType = subscriptionType;
         this.additionalTypes = dictionary;
+        this.directives = new HashSet<>(Arrays.asList(Directives.IncludeDirective, Directives.SkipDirective));
+        this.directives.addAll(directives);
         typeMap = new SchemaUtil().allTypes(this, dictionary);
     }
 
@@ -67,7 +76,7 @@ public GraphQLObjectType getSubscriptionType() {
     }
 
     public List<GraphQLDirective> getDirectives() {
-        return Arrays.asList(Directives.IncludeDirective, Directives.SkipDirective);
+        return new ArrayList<>(directives);
     }
 
     public GraphQLDirective getDirective(String name) {
@@ -122,12 +131,17 @@ public Builder subscription(GraphQLObjectType subscriptionType) {
         }
 
         public GraphQLSchema build() {
-            return build(Collections.emptySet());
+            return build(Collections.emptySet(), Collections.emptySet());
         }
 
         public GraphQLSchema build(Set<GraphQLType> additionalTypes) {
+            return build(additionalTypes, Collections.emptySet());
+        }
+
+        public GraphQLSchema build(Set<GraphQLType> additionalTypes, Set<GraphQLDirective> additionalDirectives) {
             assertNotNull(additionalTypes, "additionalTypes can't be null");
-            GraphQLSchema graphQLSchema = new GraphQLSchema(queryType, mutationType, subscriptionType, additionalTypes);
+            assertNotNull(additionalDirectives, "additionalDirectives can't be null");
+            GraphQLSchema graphQLSchema = new GraphQLSchema(queryType, mutationType, subscriptionType, additionalTypes, additionalDirectives);
             new SchemaUtil().replaceTypeReferences(graphQLSchema);
             Collection<SchemaValidationError> errors = new SchemaValidator().validateSchema(graphQLSchema);
             if (errors.size() > 0) {
diff --git a/src/test/groovy/graphql/validation/ValidateCustomDirectives.groovy b/src/test/groovy/graphql/validation/ValidateCustomDirectives.groovy
new file mode 100644
index 0000000000..7a4bca682a
--- /dev/null
+++ b/src/test/groovy/graphql/validation/ValidateCustomDirectives.groovy
@@ -0,0 +1,61 @@
+package graphql.validation
+
+import graphql.schema.GraphQLDirective;
+import graphql.schema.GraphQLSchema;
+import graphql.introspection.Introspection.DirectiveLocation;
+import graphql.parser.Parser
+import graphql.validation.ValidationError
+import graphql.validation.Validator
+import spock.lang.Specification
+/**
+ * Test that custom directives are validated properly.
+ */
+class ValidateCustomDirectives extends Specification {
+
+    GraphQLDirective customDirective = GraphQLDirective.newDirective()
+            .name("customDirective")
+            .description("Dummy Custom Directive")
+            .validLocations(DirectiveLocation.FIELD)
+            .build();
+
+    GraphQLSchema customDirectiveSchema = GraphQLSchema.newSchema()
+            .query(SpecValidationSchema.queryRoot)
+            .build(SpecValidationSchema.specValidationDictionary, [customDirective].toSet());
+
+    def 'Schema with custom directive validates query with same directive'() {
+        def query = """
+query {
+  dog {
+    name @customDirective
+  }
+}
+"""
+        when:
+        def validationErrors = validate(query)
+
+        then:
+        validationErrors.empty
+    }
+
+    def 'Schema with custom directive validates query parameters'() {
+        def query = """
+query {
+  dog {
+    name @customDirective(dummy: true)
+  }
+}
+"""
+        when:
+        def validationErrors = validate(query)
+
+        then:
+        validationErrors.size() == 1
+        validationErrors.get(0).getValidationErrorType() == ValidationErrorType.UnknownDirective
+        validationErrors.get(0).getDescription() =='Unknown directive argument dummy'
+    }
+
+    List<ValidationError> validate(String query) {
+        def document = new Parser().parseDocument(query)
+        return new Validator().validateDocument(customDirectiveSchema, document)
+    }
+}