Merge pull request graphql-java#33 from zourzouvillys/additional-types-in-schema

andimarek · andimarek · commit 0c69290a5158 · 2015-09-25T04:06:09.000+02:00
Allow providing additional types to schema build.
diff --git a/src/main/java/graphql/schema/GraphQLSchema.java b/src/main/java/graphql/schema/GraphQLSchema.java
@@ -1,31 +1,40 @@
 package graphql.schema;
 
 
-import graphql.Directives;
+import static graphql.Assert.assertNotNull;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
-import static graphql.Assert.assertNotNull;
+import graphql.Assert;
+import graphql.Directives;
 
 public class GraphQLSchema {
 
     private final GraphQLObjectType queryType;
     private final GraphQLObjectType mutationType;
     private final Map<String, GraphQLType> typeMap;
+    private Set<GraphQLType> dictionary;
 
     public GraphQLSchema(GraphQLObjectType queryType) {
-        this(queryType, null);
+        this(queryType, null, Collections.<GraphQLType>emptySet());
     }
 
+    public Set<GraphQLType> getDictionary() {
+      return dictionary;
+    }
 
-    public GraphQLSchema(GraphQLObjectType queryType, GraphQLObjectType mutationType) {
+    public GraphQLSchema(GraphQLObjectType queryType, GraphQLObjectType mutationType, Set<GraphQLType> dictionary) {
+        assertNotNull(dictionary, "dictionary can't be null");
         assertNotNull(queryType, "queryType can't be null");
         this.queryType = queryType;
         this.mutationType = mutationType;
-        typeMap = new SchemaUtil().allTypes(this);
+        this.dictionary = dictionary;
+        typeMap = new SchemaUtil().allTypes(this, dictionary);
     }
 
     public GraphQLType getType(String typeName) {
@@ -80,11 +89,17 @@ public Builder mutation(GraphQLObjectType mutationType) {
         }
 
         public GraphQLSchema build() {
-            GraphQLSchema graphQLSchema = new GraphQLSchema(queryType, mutationType);
-            new SchemaUtil().replaceTypeReferences(graphQLSchema);
-            return graphQLSchema;
-        }
+          return build(Collections.<GraphQLType>emptySet());
+      }
+
+        public GraphQLSchema build(Set<GraphQLType> dictionary) {
+          Assert.assertNotNull(dictionary, "dictionary can't be null");
+          GraphQLSchema graphQLSchema = new GraphQLSchema(queryType, mutationType, dictionary);
+          new SchemaUtil().replaceTypeReferences(graphQLSchema);
+          return graphQLSchema;
+      }
 
 
     }
+
 }
diff --git a/src/main/java/graphql/schema/SchemaUtil.java b/src/main/java/graphql/schema/SchemaUtil.java
@@ -8,6 +8,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class SchemaUtil {
 
@@ -94,18 +95,23 @@ private void collectTypesForObjects(GraphQLObjectType objectType, Map<String, Gr
         }
     }
 
-    public Map<String, GraphQLType> allTypes(GraphQLSchema schema) {
+    public Map<String, GraphQLType> allTypes(GraphQLSchema schema,  Set<GraphQLType> dictionary) {
         Map<String, GraphQLType> typesByName = new LinkedHashMap<>();
         collectTypes(schema.getQueryType(), typesByName);
         if (schema.isSupportingMutations()) {
             collectTypes(schema.getMutationType(), typesByName);
         }
+        if (dictionary != null) {
+          for (GraphQLType type : dictionary) {
+            collectTypes(type, typesByName);
+          }
+        }
         collectTypes(Introspection.__Schema, typesByName);
         return typesByName;
     }
 
     public List<GraphQLObjectType> findImplementations(GraphQLSchema schema, GraphQLInterfaceType interfaceType) {
-        Map<String, GraphQLType> allTypes = allTypes(schema);
+        Map<String, GraphQLType> allTypes = allTypes(schema, schema.getDictionary());
         List<GraphQLObjectType> result = new ArrayList<>();
         for (GraphQLType type : allTypes.values()) {
             if (!(type instanceof GraphQLObjectType)) {
@@ -119,7 +125,7 @@ public List<GraphQLObjectType> findImplementations(GraphQLSchema schema, GraphQL
 
 
     void replaceTypeReferences(GraphQLSchema schema) {
-        Map<String, GraphQLType> typeMap = allTypes(schema);
+        Map<String, GraphQLType> typeMap = allTypes(schema, schema.getDictionary());
         for (GraphQLType type : typeMap.values()) {
             if (type instanceof GraphQLFieldsContainer) {
                 resolveTypeReferencesForFieldsContainer((GraphQLFieldsContainer) type, typeMap);
diff --git a/src/test/groovy/graphql/schema/SchemaUtilTest.groovy b/src/test/groovy/graphql/schema/SchemaUtilTest.groovy
@@ -3,6 +3,8 @@ package graphql.schema
 import graphql.introspection.Introspection
 import spock.lang.Specification
 
+import java.util.Collections;
+
 import static graphql.Scalars.GraphQLBoolean
 import static graphql.Scalars.GraphQLString
 import static graphql.StarWarsSchema.*
@@ -11,7 +13,7 @@ class SchemaUtilTest extends Specification {
 
     def "collectAllTypes"() {
         when:
-        Map<String, GraphQLType> types = new SchemaUtil().allTypes(starWarsSchema)
+        Map<String, GraphQLType> types = new SchemaUtil().allTypes(starWarsSchema, Collections.emptySet())
         then:
         types == [(droidType.name)                 : droidType,
                   (humanType.name)                 : humanType,
