Reworked the DataLoader instrumentation to become compatible with 6.0.0

Author: oliemansm

gradle.properties

@@ -1,2 +1,2 @@
-version = 6.0.0
+version = 6.1.0
 group = com.graphql-java

src/main/java/graphql/servlet/GraphQLQueryInvoker.java

@@ -3,8 +3,10 @@
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.GraphQL;
+import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.SimpleInstrumentation;
+import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
 import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
 import graphql.execution.preparsed.PreparsedDocumentProvider;
 import graphql.schema.GraphQLSchema;
@@ -13,7 +15,9 @@
 import javax.security.auth.Subject;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
+import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 import java.util.function.Supplier;
 
 /**
@@ -44,17 +48,32 @@ public void query(GraphQLBatchedInvocationInput batchedInvocationInput, Executio
         }
     }
 
-    private GraphQL newGraphQL(GraphQLSchema schema) {
+    private GraphQL newGraphQL(GraphQLSchema schema, Object context) {
         ExecutionStrategyProvider executionStrategyProvider = getExecutionStrategyProvider.get();
         return GraphQL.newGraphQL(schema)
             .queryExecutionStrategy(executionStrategyProvider.getQueryExecutionStrategy())
             .mutationExecutionStrategy(executionStrategyProvider.getMutationExecutionStrategy())
             .subscriptionExecutionStrategy(executionStrategyProvider.getSubscriptionExecutionStrategy())
-            .instrumentation(getInstrumentation.get())
+            .instrumentation(getInstrumentation(context))
             .preparsedDocumentProvider(getPreparsedDocumentProvider.get())
             .build();
     }
 
+    protected Instrumentation getInstrumentation(Object context) {
+        if (context instanceof GraphQLContext) {
+            return ((GraphQLContext) context).getDataLoaderRegistry()
+                    .map(registry -> {
+                        List<Instrumentation> instrumentations = new ArrayList<>();
+                        instrumentations.add(getInstrumentation.get());
+                        instrumentations.add(new DataLoaderDispatcherInstrumentation(registry));
+                        return new ChainedInstrumentation(instrumentations);
+                    })
+                    .map(Instrumentation.class::cast)
+                    .orElse(getInstrumentation.get());
+        }
+        return getInstrumentation.get();
+    }
+
     private ExecutionResult query(GraphQLInvocationInput invocationInput, ExecutionInput executionInput) {
         if (Subject.getSubject(AccessController.getContext()) == null && invocationInput.getSubject().isPresent()) {
             return Subject.doAs(invocationInput.getSubject().get(), (PrivilegedAction<ExecutionResult>) () -> {
@@ -70,7 +89,7 @@ private ExecutionResult query(GraphQLInvocationInput invocationInput, ExecutionI
     }
 
     private ExecutionResult query(GraphQLSchema schema, ExecutionInput executionInput) {
-        return newGraphQL(schema).execute(executionInput);
+        return newGraphQL(schema, executionInput.getContext()).execute(executionInput);
     }
 
     public static Builder newBuilder() {

src/main/java/graphql/servlet/OsgiGraphQLHttpServlet.java

@@ -236,19 +236,6 @@ public InstrumentationProvider getInstrumentationProvider() {
         return instrumentationProvider;
     }
 
-    protected Instrumentation getInstrumentation(GraphQLContext context) {
-        return context.getDataLoaderRegistry()
-                .map(registry -> {
-                    List<Instrumentation> instrumentations = new ArrayList<>();
-                    instrumentations.add(this.instrumentationProvider.getInstrumentation());
-                    instrumentations.add(new DataLoaderDispatcherInstrumentation(registry));
-
-                    return new ChainedInstrumentation(instrumentations);
-                })
-                .map(Instrumentation.class::cast)
-                .orElse(instrumentationProvider.getInstrumentation());
-    }
-
     public GraphQLErrorHandler getErrorHandler() {
         return errorHandler;
     }

src/test/groovy/graphql/servlet/AbstractGraphQLHttpServletSpec.groovy

@@ -2,6 +2,7 @@ package graphql.servlet
 
 import com.fasterxml.jackson.databind.ObjectMapper
 import graphql.Scalars
+import graphql.analysis.QueryVisitorInlineFragmentEnvironment
 import graphql.execution.ExecutionTypeInfo
 import graphql.execution.instrumentation.ChainedInstrumentation
 import graphql.execution.instrumentation.Instrumentation
@@ -44,7 +45,7 @@ class AbstractGraphQLHttpServletSpec extends Specification {
 
     def createServlet(DataFetcher queryDataFetcher = { env -> env.arguments.arg },
                       DataFetcher mutationDataFetcher = { env -> env.arguments.arg }) {
-        return SimpleGraphQLHttpServlet.newBuilder(createGraphQlSchema(queryDataFetcher, mutationDataFetcher))
+        return SimpleGraphQLHttpServlet.newBuilder(createGraphQlSchema(queryDataFetcher, mutationDataFetcher)).build()
     }
 
     def createGraphQlSchema(DataFetcher queryDataFetcher = { env -> env.arguments.arg },
@@ -866,13 +867,13 @@ class AbstractGraphQLHttpServletSpec extends Specification {
 
         setup:
             Instrumentation expectedInstrumentation = Mock()
-            GraphQLContext context = new GraphQLContext(Optional.of(request), Optional.of(response))
-            SimpleGraphQLServlet simpleGraphQLServlet = SimpleGraphQLServlet
-                    .builder(createGraphQlSchema())
-                    .withInstrumentation(expectedInstrumentation)
-                    .build();
+            GraphQLContext context = new GraphQLContext(request, null, null)
+            SimpleGraphQLHttpServlet simpleGraphQLServlet = SimpleGraphQLHttpServlet
+                    .newBuilder(createGraphQlSchema())
+                    .withQueryInvoker(GraphQLQueryInvoker.newBuilder().withInstrumentation(expectedInstrumentation).build())
+                    .build()
         when:
-            Instrumentation actualInstrumentation = simpleGraphQLServlet.getInstrumentation(context)
+            Instrumentation actualInstrumentation = simpleGraphQLServlet.getQueryInvoker().getInstrumentation(context)
         then:
         actualInstrumentation == expectedInstrumentation;
         ! (actualInstrumentation instanceof ChainedInstrumentation)
@@ -883,15 +884,15 @@ class AbstractGraphQLHttpServletSpec extends Specification {
 
         setup:
             Instrumentation servletInstrumentation = Mock()
-            GraphQLContext context = new GraphQLContext(Optional.of(request), Optional.of(response))
+            GraphQLContext context = new GraphQLContext(request, null, null)
             DataLoaderRegistry dlr = Mock()
-            context.setDataLoaderRegistry(Optional.of(dlr))
-            SimpleGraphQLServlet simpleGraphQLServlet = SimpleGraphQLServlet
-                    .builder(createGraphQlSchema())
-                    .withInstrumentation(servletInstrumentation)
+            context.setDataLoaderRegistry(dlr)
+            SimpleGraphQLHttpServlet simpleGraphQLServlet = SimpleGraphQLHttpServlet
+                    .newBuilder(createGraphQlSchema())
+                    .withQueryInvoker(GraphQLQueryInvoker.newBuilder().withInstrumentation(servletInstrumentation).build())
                     .build();
         when:
-            Instrumentation actualInstrumentation = simpleGraphQLServlet.getInstrumentation(context)
+            Instrumentation actualInstrumentation = simpleGraphQLServlet.getQueryInvoker().getInstrumentation(context)
         then:
             actualInstrumentation instanceof ChainedInstrumentation
             actualInstrumentation != servletInstrumentation