Skip to content

Commit 0ed786e

Browse files
author
agairol
committed
Add DataLoaderRegistry to GraphQLContext and return Instrumentation accordingly
1 parent 763769b commit 0ed786e

File tree

5 files changed

+103
-37
lines changed

5 files changed

+103
-37
lines changed

‎src/main/java/graphql/servlet/GraphQLContext.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package graphql.servlet;
22

3+
import org.dataloader.DataLoaderRegistry;
4+
35
import javax.security.auth.Subject;
46
import javax.servlet.http.HttpServletRequest;
57
import javax.servlet.http.HttpServletResponse;
@@ -15,6 +17,8 @@ public class GraphQLContext {
1517
private Optional<Subject> subject = Optional.empty();
1618
private Optional<Map<String, List<Part>>> files = Optional.empty();
1719

20+
private Optional<DataLoaderRegistry> dataLoaderRegistry = Optional.empty();
21+
1822
public GraphQLContext(Optional<HttpServletRequest> request, Optional<HttpServletResponse> response) {
1923
this.request = request;
2024
this.response = response;
@@ -51,4 +55,12 @@ public Optional<Map<String, List<Part>>> getFiles() {
5155
public void setFiles(Optional<Map<String, List<Part>>> files) {
5256
this.files = files;
5357
}
58+
59+
public Optional<DataLoaderRegistry> getDataLoaderRegistry() {
60+
return dataLoaderRegistry;
61+
}
62+
63+
public void setDataLoaderRegistry(Optional<DataLoaderRegistry> dataLoaderRegistry) {
64+
this.dataLoaderRegistry = dataLoaderRegistry;
65+
}
5466
}

‎src/main/java/graphql/servlet/GraphQLServlet.java‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ public abstract class GraphQLServlet extends HttpServlet implements Servlet, Gra
5757

5858
protected abstract ExecutionStrategyProvider getExecutionStrategyProvider();
5959

60-
protected abstract Instrumentation getInstrumentation();
60+
protected abstract Instrumentation getInstrumentation(GraphQLContext context);
6161

6262
protected abstract GraphQLErrorHandler getGraphQLErrorHandler();
6363

@@ -68,7 +68,7 @@ public abstract class GraphQLServlet extends HttpServlet implements Servlet, Gra
6868

6969
private final HttpRequestHandler getHandler;
7070
private final HttpRequestHandler postHandler;
71-
71+
7272
private final boolean asyncServletMode;
7373

7474
public GraphQLServlet() {
@@ -246,7 +246,8 @@ public String[] getMutations() {
246246
@Override
247247
public String executeQuery(String query) {
248248
try {
249-
final ExecutionResult result = newGraphQL(getSchemaProvider().getSchema()).execute(new ExecutionInput(query, null, createContext(Optional.empty(), Optional.empty()), createRootObject(Optional.empty(), Optional.empty()), new HashMap<>()));
249+
GraphQLContext context = createContext(Optional.empty(), Optional.empty());
250+
final ExecutionResult result = newGraphQL(getSchemaProvider().getSchema(), context).execute(new ExecutionInput(query, null, context, createRootObject(Optional.empty(), Optional.empty()), new HashMap<>()));
250251
return getMapper().writeValueAsString(createResultFromDataErrorsAndExtensions(result.getData(), result.getErrors(), result.getExtensions()));
251252
} catch (Exception e) {
252253
return e.getMessage();
@@ -299,13 +300,13 @@ private Optional<Part> getFileItem(Map<String, List<Part>> fileItems, String nam
299300
return Optional.ofNullable(fileItems.get(name)).filter(list -> !list.isEmpty()).map(list -> list.get(0));
300301
}
301302

302-
private GraphQL newGraphQL(GraphQLSchema schema) {
303+
private GraphQL newGraphQL(GraphQLSchema schema, GraphQLContext context) {
303304
ExecutionStrategyProvider executionStrategyProvider = getExecutionStrategyProvider();
304305
return GraphQL.newGraphQL(schema)
305306
.queryExecutionStrategy(executionStrategyProvider.getQueryExecutionStrategy())
306307
.mutationExecutionStrategy(executionStrategyProvider.getMutationExecutionStrategy())
307308
.subscriptionExecutionStrategy(executionStrategyProvider.getSubscriptionExecutionStrategy())
308-
.instrumentation(getInstrumentation())
309+
.instrumentation(getInstrumentation(context))
309310
.preparsedDocumentProvider(getPreparsedDocumentProvider())
310311
.build();
311312
}
@@ -353,7 +354,7 @@ private void query(String query, String operationName, Map<String, Object> varia
353354
} else {
354355
List<GraphQLServletListener.OperationCallback> operationCallbacks = runListeners(l -> l.onOperation(context, operationName, query, variables));
355356

356-
final ExecutionResult executionResult = newGraphQL(schema).execute(new ExecutionInput(query, operationName, context, rootObject, variables));
357+
final ExecutionResult executionResult = newGraphQL(schema, context).execute(new ExecutionInput(query, operationName, context, rootObject, variables));
357358
final List<GraphQLError> errors = executionResult.getErrors();
358359
final Object data = executionResult.getData();
359360
final Object extensions = executionResult.getExtensions();

‎src/main/java/graphql/servlet/OsgiGraphQLServlet.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package graphql.servlet;
22

33
import graphql.execution.instrumentation.Instrumentation;
4+
import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
45
import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
56
import graphql.execution.preparsed.PreparsedDocumentProvider;
67
import graphql.schema.GraphQLObjectType;
@@ -204,8 +205,11 @@ protected ExecutionStrategyProvider getExecutionStrategyProvider() {
204205
}
205206

206207
@Override
207-
protected Instrumentation getInstrumentation() {
208-
return instrumentationProvider.getInstrumentation();
208+
protected Instrumentation getInstrumentation(GraphQLContext context) {
209+
return context.getDataLoaderRegistry()
210+
.map(DataLoaderDispatcherInstrumentation::new)
211+
.map(Instrumentation.class::cast)
212+
.orElse(instrumentationProvider.getInstrumentation());
209213
}
210214

211215
@Override

‎src/main/java/graphql/servlet/SimpleGraphQLServlet.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import graphql.execution.ExecutionStrategy;
1010
import graphql.execution.instrumentation.Instrumentation;
1111
import graphql.execution.instrumentation.SimpleInstrumentation;
12+
import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
1213
import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
1314
import graphql.execution.preparsed.PreparsedDocumentProvider;
1415
import graphql.schema.GraphQLSchema;
@@ -220,8 +221,11 @@ protected ExecutionStrategyProvider getExecutionStrategyProvider() {
220221
}
221222

222223
@Override
223-
protected Instrumentation getInstrumentation() {
224-
return instrumentation;
224+
protected Instrumentation getInstrumentation(GraphQLContext context) {
225+
return context.getDataLoaderRegistry()
226+
.map(DataLoaderDispatcherInstrumentation::new)
227+
.map(Instrumentation.class::cast)
228+
.orElse(this.instrumentation);
225229
}
226230

227231
@Override

‎src/test/groovy/graphql/servlet/GraphQLServletSpec.groovy‎

Lines changed: 72 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,14 @@ package graphql.servlet
33
import com.fasterxml.jackson.databind.ObjectMapper
44
import graphql.Scalars
55
import graphql.execution.ExecutionTypeInfo
6+
import graphql.execution.instrumentation.Instrumentation
7+
import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation
68
import graphql.schema.DataFetcher
79
import graphql.schema.GraphQLFieldDefinition
810
import graphql.schema.GraphQLNonNull
911
import graphql.schema.GraphQLObjectType
1012
import graphql.schema.GraphQLSchema
13+
import org.dataloader.DataLoaderRegistry
1114
import org.springframework.mock.web.MockHttpServletRequest
1215
import org.springframework.mock.web.MockHttpServletResponse
1316
import spock.lang.Shared
@@ -38,39 +41,45 @@ class GraphQLServletSpec extends Specification {
3841
response = new MockHttpServletResponse()
3942
}
4043

41-
def createServlet(DataFetcher queryDataFetcher = { env -> env.arguments.arg }, DataFetcher mutationDataFetcher = { env -> env.arguments.arg }) {
44+
def createServlet(DataFetcher queryDataFetcher = { env -> env.arguments.arg },
45+
DataFetcher mutationDataFetcher = { env -> env.arguments.arg }) {
46+
return new SimpleGraphQLServlet(createGraphQlSchema(queryDataFetcher, mutationDataFetcher))
47+
}
48+
49+
def createGraphQlSchema(DataFetcher queryDataFetcher = { env -> env.arguments.arg },
50+
DataFetcher mutationDataFetcher = { env -> env.arguments.arg }) {
4251
GraphQLObjectType query = GraphQLObjectType.newObject()
43-
.name("Query")
44-
.field { GraphQLFieldDefinition.Builder field ->
45-
field.name("echo")
46-
field.type(Scalars.GraphQLString)
47-
field.argument { argument ->
48-
argument.name("arg")
49-
argument.type(Scalars.GraphQLString)
50-
}
51-
field.dataFetcher(queryDataFetcher)
52+
.name("Query")
53+
.field { GraphQLFieldDefinition.Builder field ->
54+
field.name("echo")
55+
field.type(Scalars.GraphQLString)
56+
field.argument { argument ->
57+
argument.name("arg")
58+
argument.type(Scalars.GraphQLString)
5259
}
53-
.field { GraphQLFieldDefinition.Builder field ->
54-
field.name("returnsNullIncorrectly")
55-
field.type(new GraphQLNonNull(Scalars.GraphQLString))
56-
field.dataFetcher({env -> null})
57-
}
58-
.build()
60+
field.dataFetcher(queryDataFetcher)
61+
}
62+
.field { GraphQLFieldDefinition.Builder field ->
63+
field.name("returnsNullIncorrectly")
64+
field.type(new GraphQLNonNull(Scalars.GraphQLString))
65+
field.dataFetcher({env -> null})
66+
}
67+
.build()
5968

6069
GraphQLObjectType mutation = GraphQLObjectType.newObject()
61-
.name("Mutation")
62-
.field { field ->
63-
field.name("echo")
64-
field.type(Scalars.GraphQLString)
65-
field.argument { argument ->
66-
argument.name("arg")
67-
argument.type(Scalars.GraphQLString)
68-
}
69-
field.dataFetcher(mutationDataFetcher)
70+
.name("Mutation")
71+
.field { field ->
72+
field.name("echo")
73+
field.type(Scalars.GraphQLString)
74+
field.argument { argument ->
75+
argument.name("arg")
76+
argument.type(Scalars.GraphQLString)
7077
}
71-
.build()
78+
field.dataFetcher(mutationDataFetcher)
79+
}
80+
.build()
7281

73-
return new SimpleGraphQLServlet(new GraphQLSchema(query, mutation, [query, mutation].toSet()))
82+
return new GraphQLSchema(query, mutation, [query, mutation].toSet())
7483
}
7584

7685
Map<String, Object> getResponseContent() {
@@ -855,4 +864,40 @@ class GraphQLServletSpec extends Specification {
855864
then:
856865
1 * mockInputStream.reset()
857866
}
867+
868+
def "getInstrumentation returns the set Instrumentation if none is provided in the context"() {
869+
870+
setup:
871+
Instrumentation expectedInstrumentation = Mock()
872+
GraphQLContext context = new GraphQLContext(Optional.of(request), Optional.of(response))
873+
SimpleGraphQLServlet simpleGraphQLServlet = SimpleGraphQLServlet
874+
.builder(createGraphQlSchema())
875+
.withInstrumentation(expectedInstrumentation)
876+
.build();
877+
when:
878+
Instrumentation actualInstrumentation = simpleGraphQLServlet.getInstrumentation(context)
879+
then:
880+
actualInstrumentation == expectedInstrumentation;
881+
! (actualInstrumentation instanceof DataLoaderDispatcherInstrumentation)
882+
883+
}
884+
885+
def "getInstrumentation returns the DataLoaderDispatcherInstrumentation if DataLoader provided in context"() {
886+
887+
setup:
888+
Instrumentation servletInstrumentation = Mock()
889+
GraphQLContext context = new GraphQLContext(Optional.of(request), Optional.of(response))
890+
DataLoaderRegistry dlr = Mock()
891+
context.setDataLoaderRegistry(Optional.of(dlr))
892+
SimpleGraphQLServlet simpleGraphQLServlet = SimpleGraphQLServlet
893+
.builder(createGraphQlSchema())
894+
.withInstrumentation(servletInstrumentation)
895+
.build();
896+
when:
897+
Instrumentation actualInstrumentation = simpleGraphQLServlet.getInstrumentation(context)
898+
then:
899+
actualInstrumentation instanceof DataLoaderDispatcherInstrumentation
900+
actualInstrumentation != servletInstrumentation;
901+
902+
}
858903
}

0 commit comments

Comments
 (0)