first class dataloader support

src/main/java/graphql/ExecutionInput.java

@@ -1,9 +1,13 @@
 package graphql;
 
+import org.dataloader.DataLoaderRegistry;
+
 import java.util.Collections;
 import java.util.Map;
 import java.util.function.Consumer;
 
+import static graphql.Assert.assertNotNull;
+
 /**
  * This represents the series of values that can be input on a graphql query execution
  */
@@ -14,14 +18,21 @@ public class ExecutionInput {
     private final Object context;
     private final Object root;
     private final Map<String, Object> variables;
+    private final DataLoaderRegistry dataLoaderRegistry;
 
 
     public ExecutionInput(String query, String operationName, Object context, Object root, Map<String, Object> variables) {
+        this(query, operationName, context, root, variables, new DataLoaderRegistry());
+    }
+
+    @Internal
+    private ExecutionInput(String query, String operationName, Object context, Object root, Map<String, Object> variables, DataLoaderRegistry dataLoaderRegistry) {
         this.query = query;
         this.operationName = operationName;
         this.context = context;
         this.root = root;
         this.variables = variables;
+        this.dataLoaderRegistry = dataLoaderRegistry;
     }
 
     /**
@@ -59,6 +70,13 @@ public Map<String, Object> getVariables() {
         return variables;
     }
 
+    /**
+     * @return the data loader registry associated with this execution
+     */
+    public DataLoaderRegistry getDataLoaderRegistry() {
+        return dataLoaderRegistry;
+    }
+
     /**
      * This helps you transform the current ExecutionInput object into another one by starting a builder with all
      * the current values and allows you to transform it how you want.
@@ -73,6 +91,7 @@ public ExecutionInput transform(Consumer<Builder> builderConsumer) {
                 .operationName(this.operationName)
                 .context(this.context)
                 .root(this.root)
+                .dataLoaderRegistry(this.dataLoaderRegistry)
                 .variables(this.variables);
 
         builderConsumer.accept(builder);
@@ -89,6 +108,7 @@ public String toString() {
                 ", context=" + context +
                 ", root=" + root +
                 ", variables=" + variables +
+                ", dataLoaderRegistry=" + dataLoaderRegistry +
                 '}';
     }
 
@@ -106,6 +126,7 @@ public static class Builder {
         private Object context;
         private Object root;
         private Map<String, Object> variables = Collections.emptyMap();
+        private DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry();
 
         public Builder query(String query) {
             this.query = query;
@@ -132,8 +153,21 @@ public Builder variables(Map<String, Object> variables) {
             return this;
         }
 
+        /**
+         * You should create new {@link org.dataloader.DataLoaderRegistry}s and new {@link org.dataloader.DataLoader}s for each execution.  Do not re-use
+         * instances as this will create unexpected results.
+         *
+         * @param dataLoaderRegistry a registry of {@link org.dataloader.DataLoader}s
+         *
+         * @return this builder
+         */
+        public Builder dataLoaderRegistry(DataLoaderRegistry dataLoaderRegistry) {
+            this.dataLoaderRegistry = assertNotNull(dataLoaderRegistry);
+            return this;
+        }
+
         public ExecutionInput build() {
-            return new ExecutionInput(query, operationName, context, root, variables);
+            return new ExecutionInput(query, operationName, context, root, variables, dataLoaderRegistry);
         }
     }
 }

src/main/java/graphql/GraphQL.java

@@ -8,10 +8,13 @@
 import graphql.execution.ExecutionIdProvider;
 import graphql.execution.ExecutionStrategy;
 import graphql.execution.SubscriptionExecutionStrategy;
+import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.SimpleInstrumentation;
+import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
+import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
 import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
@@ -26,6 +29,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -158,10 +162,25 @@ private GraphQL(GraphQLSchema graphQLSchema, ExecutionStrategy queryStrategy, Ex
         this.mutationStrategy = mutationStrategy != null ? mutationStrategy : new AsyncSerialExecutionStrategy();
         this.subscriptionStrategy = subscriptionStrategy != null ? subscriptionStrategy : new SubscriptionExecutionStrategy();
         this.idProvider = assertNotNull(idProvider, "idProvider must be non null");
-        this.instrumentation = instrumentation;
+        this.instrumentation = checkInstrumentation(assertNotNull(instrumentation));
         this.preparsedDocumentProvider = assertNotNull(preparsedDocumentProvider, "preparsedDocumentProvider must be non null");
     }
 
+    private Instrumentation checkInstrumentation(Instrumentation instrumentation) {
+        List<Instrumentation> instrumentationList = new ArrayList<>();
+        if (instrumentation instanceof ChainedInstrumentation) {
+            instrumentationList.addAll(((ChainedInstrumentation) instrumentation).getInstrumentations());
+        } else {
+            instrumentationList.add(instrumentation);
+        }
+        boolean containsDLInstrumentation = instrumentationList.stream().anyMatch(instr -> instr instanceof DataLoaderDispatcherInstrumentation);
+        // if we don't have a DataLoaderDispatcherInstrumentation in play, we add one.  We want DataLoader to be 1st class in graphql
+        if (!containsDLInstrumentation) {
+            instrumentationList.add(new DataLoaderDispatcherInstrumentation());
+        }
+        return new ChainedInstrumentation(instrumentationList);
+    }
+
     /**
      * Helps you build a GraphQL object ready to execute queries
      *
@@ -457,7 +476,7 @@ public CompletableFuture<ExecutionResult> executeAsync(ExecutionInput executionI
         try {
             log.debug("Executing request. operation name: '{}'. query: '{}'. variables '{}'", executionInput.getOperationName(), executionInput.getQuery(), executionInput.getVariables());
 
-            InstrumentationState instrumentationState = instrumentation.createState();
+            InstrumentationState instrumentationState = instrumentation.createState(new InstrumentationCreateStateParameters(this.graphQLSchema, executionInput));
 
             InstrumentationExecutionParameters inputInstrumentationParameters = new InstrumentationExecutionParameters(executionInput, this.graphQLSchema, instrumentationState);
             executionInput = instrumentation.instrumentExecutionInput(executionInput, inputInstrumentationParameters);

src/main/java/graphql/execution/Execution.java

@@ -90,6 +90,7 @@ public CompletableFuture<ExecutionResult> execute(Document document, GraphQLSche
                 .variables(coercedVariables)
                 .document(document)
                 .operationDefinition(operationDefinition)
+                .dataLoaderRegistry(executionInput.getDataLoaderRegistry())
                 .build();
 
 

src/main/java/graphql/execution/ExecutionContext.java

@@ -2,6 +2,7 @@
 
 
 import graphql.GraphQLError;
+import graphql.Internal;
 import graphql.PublicApi;
 import graphql.execution.defer.DeferSupport;
 import graphql.execution.instrumentation.Instrumentation;
@@ -10,6 +11,7 @@
 import graphql.language.FragmentDefinition;
 import graphql.language.OperationDefinition;
 import graphql.schema.GraphQLSchema;
+import org.dataloader.DataLoaderRegistry;
 
 import java.util.Collections;
 import java.util.List;
@@ -35,13 +37,11 @@ public class ExecutionContext {
     private final Object context;
     private final Instrumentation instrumentation;
     private final List<GraphQLError> errors = new CopyOnWriteArrayList<>();
+    private final DataLoaderRegistry dataLoaderRegistry;
     private final DeferSupport deferSupport = new DeferSupport();
 
-    public ExecutionContext(Instrumentation instrumentation, ExecutionId executionId, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState, ExecutionStrategy queryStrategy, ExecutionStrategy mutationStrategy, ExecutionStrategy subscriptionStrategy, Map<String, FragmentDefinition> fragmentsByName, Document document, OperationDefinition operationDefinition, Map<String, Object> variables, Object context, Object root) {
-        this(instrumentation, executionId, graphQLSchema, instrumentationState, queryStrategy, mutationStrategy, subscriptionStrategy, fragmentsByName, document, operationDefinition, variables, context, root, Collections.emptyList());
-    }
-
-    ExecutionContext(Instrumentation instrumentation, ExecutionId executionId, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState, ExecutionStrategy queryStrategy, ExecutionStrategy mutationStrategy, ExecutionStrategy subscriptionStrategy, Map<String, FragmentDefinition> fragmentsByName, Document document, OperationDefinition operationDefinition, Map<String, Object> variables, Object context, Object root, List<GraphQLError> startingErrors) {
+    @Internal
+    ExecutionContext(Instrumentation instrumentation, ExecutionId executionId, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState, ExecutionStrategy queryStrategy, ExecutionStrategy mutationStrategy, ExecutionStrategy subscriptionStrategy, Map<String, FragmentDefinition> fragmentsByName, Document document, OperationDefinition operationDefinition, Map<String, Object> variables, Object context, Object root, DataLoaderRegistry dataLoaderRegistry, List<GraphQLError> startingErrors) {
         this.graphQLSchema = graphQLSchema;
         this.executionId = executionId;
         this.instrumentationState = instrumentationState;
@@ -55,6 +55,7 @@ public ExecutionContext(Instrumentation instrumentation, ExecutionId executionId
         this.context = context;
         this.root = root;
         this.instrumentation = instrumentation;
+        this.dataLoaderRegistry = dataLoaderRegistry;
         this.errors.addAll(startingErrors);
     }
 
@@ -104,6 +105,10 @@ public FragmentDefinition getFragment(String name) {
         return fragmentsByName.get(name);
     }
 
+    public DataLoaderRegistry getDataLoaderRegistry() {
+        return dataLoaderRegistry;
+    }
+
     /**
      * This method will only put one error per field path.
      *

src/main/java/graphql/execution/ExecutionContextBuilder.java

@@ -9,6 +9,7 @@
 import graphql.language.FragmentDefinition;
 import graphql.language.OperationDefinition;
 import graphql.schema.GraphQLSchema;
+import org.dataloader.DataLoaderRegistry;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -33,6 +34,7 @@ public class ExecutionContextBuilder {
     private OperationDefinition operationDefinition;
     private Map<String, Object> variables = new HashMap<>();
     private Map<String, FragmentDefinition> fragmentsByName = new HashMap<>();
+    private DataLoaderRegistry dataLoaderRegistry;
     private List<GraphQLError> errors = new ArrayList<>();
 
     /**
@@ -72,6 +74,7 @@ public ExecutionContextBuilder() {
         operationDefinition = other.getOperationDefinition();
         variables = new HashMap<>(other.getVariables());
         fragmentsByName = new HashMap<>(other.getFragmentsByName());
+        dataLoaderRegistry = other.getDataLoaderRegistry();
         errors = new ArrayList<>(other.getErrors());
     }
 
@@ -141,6 +144,11 @@ public ExecutionContextBuilder operationDefinition(OperationDefinition operation
     }
 
 
+    public ExecutionContextBuilder dataLoaderRegistry(DataLoaderRegistry dataLoaderRegistry) {
+        this.dataLoaderRegistry = assertNotNull(dataLoaderRegistry);
+        return this;
+    }
+
     public ExecutionContext build() {
         // preconditions
         assertNotNull(executionId, "You must provide a query identifier");
@@ -159,6 +167,7 @@ public ExecutionContext build() {
                 variables,
                 context,
                 root,
-                errors);
+                dataLoaderRegistry, errors
+        );
     }
 }

src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java

@@ -6,6 +6,7 @@
 import graphql.execution.Async;
 import graphql.execution.ExecutionContext;
 import graphql.execution.FieldValueInfo;
+import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationDeferredFieldParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
@@ -20,6 +21,7 @@
 import graphql.schema.GraphQLSchema;
 import graphql.validation.ValidationError;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -48,14 +50,21 @@ public ChainedInstrumentation(List<Instrumentation> instrumentations) {
         this.instrumentations = Collections.unmodifiableList(assertNotNull(instrumentations));
     }
 
+    /**
+     * @return the list of instrumentations in play
+     */
+    public List<Instrumentation> getInstrumentations() {
+        return instrumentations;
+    }
+
     private InstrumentationState getState(Instrumentation instrumentation, InstrumentationState parametersInstrumentationState) {
         ChainedInstrumentationState chainedInstrumentationState = (ChainedInstrumentationState) parametersInstrumentationState;
         return chainedInstrumentationState.getState(instrumentation);
     }
 
     @Override
-    public InstrumentationState createState() {
-        return new ChainedInstrumentationState(instrumentations);
+    public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
+        return new ChainedInstrumentationState(instrumentations, parameters);
     }
 
     @Override
@@ -208,9 +217,9 @@ private static class ChainedInstrumentationState implements InstrumentationState
         private final Map<Instrumentation, InstrumentationState> instrumentationStates;
 
 
-        private ChainedInstrumentationState(List<Instrumentation> instrumentations) {
+        private ChainedInstrumentationState(List<Instrumentation> instrumentations, InstrumentationCreateStateParameters parameters) {
             instrumentationStates = new LinkedHashMap<>(instrumentations.size());
-            instrumentations.forEach(i -> instrumentationStates.put(i, i.createState()));
+            instrumentations.forEach(i -> instrumentationStates.put(i, i.createState(parameters)));
         }
 
         private InstrumentationState getState(Instrumentation instrumentation) {

src/main/java/graphql/execution/instrumentation/Instrumentation.java

@@ -3,6 +3,7 @@
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.execution.ExecutionContext;
+import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationDeferredFieldParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
@@ -44,6 +45,18 @@ default InstrumentationState createState() {
         return null;
     }
 
+    /**
+     * This will be called just before execution to create an object that is given back to all instrumentation methods
+     * to allow them to have per execution request state
+     *
+     * @param parameters the parameters to this step
+     *
+     * @return a state object that is passed to each method
+     */
+    default InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
+        return createState();
+    }
+
     /**
      * This is called right at the start of query execution and its the first step in the instrumentation chain.
      *

src/main/java/graphql/execution/instrumentation/InstrumentationState.java

@@ -4,7 +4,7 @@
  * An {@link Instrumentation} implementation can create this as a stateful object that is then passed
  * to each instrumentation method, allowing state to be passed down with the request execution
  *
- * @see Instrumentation#createState()
+ * @see Instrumentation#createState(graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters)
  */
 public interface InstrumentationState {
 }