Skip to content

Commit 78a6e4e

Browse files
bbakermanbbakerman
authored
first class dataloader support (#1263)
* Added @OverRide as part of errorprone code health check * Revert "Added @OverRide as part of errorprone code health check" This reverts commit 38dfab1 * Initial commit of 1st class dataloader support * More tweaks and some extra tests * Added more examples. Need to write the words on these * Added doco * Found a bu and re-arranged the tests * Found a bug and then write a missing set of tests * Doco updates * PR feedback * Updated constructor - bloody groovy * Made the data loader hanging test create a new DLR per execution * Reverted back to old doco
1 parent ac9d5c6 commit 78a6e4e

32 files changed

+1094
-645
lines changed

‎src/main/java/graphql/ExecutionInput.java‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
package graphql;
22

3+
import org.dataloader.DataLoaderRegistry;
4+
35
import java.util.Collections;
46
import java.util.Map;
57
import java.util.function.Consumer;
68

9+
import static graphql.Assert.assertNotNull;
10+
711
/**
812
* This represents the series of values that can be input on a graphql query execution
913
*/
@@ -14,14 +18,21 @@ public class ExecutionInput {
1418
private final Object context;
1519
private final Object root;
1620
private final Map<String, Object> variables;
21+
private final DataLoaderRegistry dataLoaderRegistry;
1722

1823

1924
public ExecutionInput(String query, String operationName, Object context, Object root, Map<String, Object> variables) {
25+
this(query, operationName, context, root, variables, new DataLoaderRegistry());
26+
}
27+
28+
@Internal
29+
private ExecutionInput(String query, String operationName, Object context, Object root, Map<String, Object> variables, DataLoaderRegistry dataLoaderRegistry) {
2030
this.query = query;
2131
this.operationName = operationName;
2232
this.context = context;
2333
this.root = root;
2434
this.variables = variables;
35+
this.dataLoaderRegistry = dataLoaderRegistry;
2536
}
2637

2738
/**
@@ -59,6 +70,13 @@ public Map<String, Object> getVariables() {
5970
return variables;
6071
}
6172

73+
/**
74+
* @return the data loader registry associated with this execution
75+
*/
76+
public DataLoaderRegistry getDataLoaderRegistry() {
77+
return dataLoaderRegistry;
78+
}
79+
6280
/**
6381
* This helps you transform the current ExecutionInput object into another one by starting a builder with all
6482
* the current values and allows you to transform it how you want.
@@ -73,6 +91,7 @@ public ExecutionInput transform(Consumer<Builder> builderConsumer) {
7391
.operationName(this.operationName)
7492
.context(this.context)
7593
.root(this.root)
94+
.dataLoaderRegistry(this.dataLoaderRegistry)
7695
.variables(this.variables);
7796

7897
builderConsumer.accept(builder);
@@ -89,6 +108,7 @@ public String toString() {
89108
", context=" + context +
90109
", root=" + root +
91110
", variables=" + variables +
111+
", dataLoaderRegistry=" + dataLoaderRegistry +
92112
'}';
93113
}
94114

@@ -106,6 +126,7 @@ public static class Builder {
106126
private Object context;
107127
private Object root;
108128
private Map<String, Object> variables = Collections.emptyMap();
129+
private DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry();
109130

110131
public Builder query(String query) {
111132
this.query = query;
@@ -132,8 +153,21 @@ public Builder variables(Map<String, Object> variables) {
132153
return this;
133154
}
134155

156+
/**
157+
* You should create new {@link org.dataloader.DataLoaderRegistry}s and new {@link org.dataloader.DataLoader}s for each execution. Do not re-use
158+
* instances as this will create unexpected results.
159+
*
160+
* @param dataLoaderRegistry a registry of {@link org.dataloader.DataLoader}s
161+
*
162+
* @return this builder
163+
*/
164+
public Builder dataLoaderRegistry(DataLoaderRegistry dataLoaderRegistry) {
165+
this.dataLoaderRegistry = assertNotNull(dataLoaderRegistry);
166+
return this;
167+
}
168+
135169
public ExecutionInput build() {
136-
return new ExecutionInput(query, operationName, context, root, variables);
170+
return new ExecutionInput(query, operationName, context, root, variables, dataLoaderRegistry);
137171
}
138172
}
139173
}

‎src/main/java/graphql/GraphQL.java‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@
88
import graphql.execution.ExecutionIdProvider;
99
import graphql.execution.ExecutionStrategy;
1010
import graphql.execution.SubscriptionExecutionStrategy;
11+
import graphql.execution.instrumentation.ChainedInstrumentation;
1112
import graphql.execution.instrumentation.Instrumentation;
1213
import graphql.execution.instrumentation.InstrumentationContext;
1314
import graphql.execution.instrumentation.InstrumentationState;
1415
import graphql.execution.instrumentation.SimpleInstrumentation;
16+
import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
17+
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
1518
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
1619
import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
1720
import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
@@ -26,6 +29,7 @@
2629
import org.slf4j.Logger;
2730
import org.slf4j.LoggerFactory;
2831

32+
import java.util.ArrayList;
2933
import java.util.List;
3034
import java.util.Map;
3135
import java.util.concurrent.CompletableFuture;
@@ -158,10 +162,25 @@ private GraphQL(GraphQLSchema graphQLSchema, ExecutionStrategy queryStrategy, Ex
158162
this.mutationStrategy = mutationStrategy != null ? mutationStrategy : new AsyncSerialExecutionStrategy();
159163
this.subscriptionStrategy = subscriptionStrategy != null ? subscriptionStrategy : new SubscriptionExecutionStrategy();
160164
this.idProvider = assertNotNull(idProvider, "idProvider must be non null");
161-
this.instrumentation = instrumentation;
165+
this.instrumentation = checkInstrumentation(assertNotNull(instrumentation));
162166
this.preparsedDocumentProvider = assertNotNull(preparsedDocumentProvider, "preparsedDocumentProvider must be non null");
163167
}
164168

169+
private Instrumentation checkInstrumentation(Instrumentation instrumentation) {
170+
List<Instrumentation> instrumentationList = new ArrayList<>();
171+
if (instrumentation instanceof ChainedInstrumentation) {
172+
instrumentationList.addAll(((ChainedInstrumentation) instrumentation).getInstrumentations());
173+
} else {
174+
instrumentationList.add(instrumentation);
175+
}
176+
boolean containsDLInstrumentation = instrumentationList.stream().anyMatch(instr -> instr instanceof DataLoaderDispatcherInstrumentation);
177+
// if we don't have a DataLoaderDispatcherInstrumentation in play, we add one. We want DataLoader to be 1st class in graphql
178+
if (!containsDLInstrumentation) {
179+
instrumentationList.add(new DataLoaderDispatcherInstrumentation());
180+
}
181+
return new ChainedInstrumentation(instrumentationList);
182+
}
183+
165184
/**
166185
* Helps you build a GraphQL object ready to execute queries
167186
*
@@ -457,7 +476,7 @@ public CompletableFuture<ExecutionResult> executeAsync(ExecutionInput executionI
457476
try {
458477
log.debug("Executing request. operation name: '{}'. query: '{}'. variables '{}'", executionInput.getOperationName(), executionInput.getQuery(), executionInput.getVariables());
459478

460-
InstrumentationState instrumentationState = instrumentation.createState();
479+
InstrumentationState instrumentationState = instrumentation.createState(new InstrumentationCreateStateParameters(this.graphQLSchema, executionInput));
461480

462481
InstrumentationExecutionParameters inputInstrumentationParameters = new InstrumentationExecutionParameters(executionInput, this.graphQLSchema, instrumentationState);
463482
executionInput = instrumentation.instrumentExecutionInput(executionInput, inputInstrumentationParameters);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ public CompletableFuture<ExecutionResult> execute(Document document, GraphQLSche
9090
.variables(coercedVariables)
9191
.document(document)
9292
.operationDefinition(operationDefinition)
93+
.dataLoaderRegistry(executionInput.getDataLoaderRegistry())
9394
.build();
9495

9596

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33

44
import graphql.GraphQLError;
5+
import graphql.Internal;
56
import graphql.PublicApi;
67
import graphql.execution.defer.DeferSupport;
78
import graphql.execution.instrumentation.Instrumentation;
@@ -10,6 +11,7 @@
1011
import graphql.language.FragmentDefinition;
1112
import graphql.language.OperationDefinition;
1213
import graphql.schema.GraphQLSchema;
14+
import org.dataloader.DataLoaderRegistry;
1315

1416
import java.util.Collections;
1517
import java.util.List;
@@ -35,13 +37,11 @@ public class ExecutionContext {
3537
private final Object context;
3638
private final Instrumentation instrumentation;
3739
private final List<GraphQLError> errors = new CopyOnWriteArrayList<>();
40+
private final DataLoaderRegistry dataLoaderRegistry;
3841
private final DeferSupport deferSupport = new DeferSupport();
3942

40-
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) {
41-
this(instrumentation, executionId, graphQLSchema, instrumentationState, queryStrategy, mutationStrategy, subscriptionStrategy, fragmentsByName, document, operationDefinition, variables, context, root, Collections.emptyList());
42-
}
43-
44-
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) {
43+
@Internal
44+
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) {
4545
this.graphQLSchema = graphQLSchema;
4646
this.executionId = executionId;
4747
this.instrumentationState = instrumentationState;
@@ -55,6 +55,7 @@ public ExecutionContext(Instrumentation instrumentation, ExecutionId executionId
5555
this.context = context;
5656
this.root = root;
5757
this.instrumentation = instrumentation;
58+
this.dataLoaderRegistry = dataLoaderRegistry;
5859
this.errors.addAll(startingErrors);
5960
}
6061

@@ -104,6 +105,10 @@ public FragmentDefinition getFragment(String name) {
104105
return fragmentsByName.get(name);
105106
}
106107

108+
public DataLoaderRegistry getDataLoaderRegistry() {
109+
return dataLoaderRegistry;
110+
}
111+
107112
/**
108113
* This method will only put one error per field path.
109114
*

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import graphql.language.FragmentDefinition;
1010
import graphql.language.OperationDefinition;
1111
import graphql.schema.GraphQLSchema;
12+
import org.dataloader.DataLoaderRegistry;
1213

1314
import java.util.ArrayList;
1415
import java.util.HashMap;
@@ -33,6 +34,7 @@ public class ExecutionContextBuilder {
3334
private OperationDefinition operationDefinition;
3435
private Map<String, Object> variables = new HashMap<>();
3536
private Map<String, FragmentDefinition> fragmentsByName = new HashMap<>();
37+
private DataLoaderRegistry dataLoaderRegistry;
3638
private List<GraphQLError> errors = new ArrayList<>();
3739

3840
/**
@@ -72,6 +74,7 @@ public ExecutionContextBuilder() {
7274
operationDefinition = other.getOperationDefinition();
7375
variables = new HashMap<>(other.getVariables());
7476
fragmentsByName = new HashMap<>(other.getFragmentsByName());
77+
dataLoaderRegistry = other.getDataLoaderRegistry();
7578
errors = new ArrayList<>(other.getErrors());
7679
}
7780

@@ -141,6 +144,11 @@ public ExecutionContextBuilder operationDefinition(OperationDefinition operation
141144
}
142145

143146

147+
public ExecutionContextBuilder dataLoaderRegistry(DataLoaderRegistry dataLoaderRegistry) {
148+
this.dataLoaderRegistry = assertNotNull(dataLoaderRegistry);
149+
return this;
150+
}
151+
144152
public ExecutionContext build() {
145153
// preconditions
146154
assertNotNull(executionId, "You must provide a query identifier");
@@ -159,6 +167,7 @@ public ExecutionContext build() {
159167
variables,
160168
context,
161169
root,
162-
errors);
170+
dataLoaderRegistry, errors
171+
);
163172
}
164173
}

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

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import graphql.execution.Async;
77
import graphql.execution.ExecutionContext;
88
import graphql.execution.FieldValueInfo;
9+
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
910
import graphql.execution.instrumentation.parameters.InstrumentationDeferredFieldParameters;
1011
import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
1112
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
@@ -20,6 +21,7 @@
2021
import graphql.schema.GraphQLSchema;
2122
import graphql.validation.ValidationError;
2223

24+
import java.util.ArrayList;
2325
import java.util.Collections;
2426
import java.util.LinkedHashMap;
2527
import java.util.List;
@@ -48,14 +50,21 @@ public ChainedInstrumentation(List<Instrumentation> instrumentations) {
4850
this.instrumentations = Collections.unmodifiableList(assertNotNull(instrumentations));
4951
}
5052

53+
/**
54+
* @return the list of instrumentations in play
55+
*/
56+
public List<Instrumentation> getInstrumentations() {
57+
return instrumentations;
58+
}
59+
5160
private InstrumentationState getState(Instrumentation instrumentation, InstrumentationState parametersInstrumentationState) {
5261
ChainedInstrumentationState chainedInstrumentationState = (ChainedInstrumentationState) parametersInstrumentationState;
5362
return chainedInstrumentationState.getState(instrumentation);
5463
}
5564

5665
@Override
57-
public InstrumentationState createState() {
58-
return new ChainedInstrumentationState(instrumentations);
66+
public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
67+
return new ChainedInstrumentationState(instrumentations, parameters);
5968
}
6069

6170
@Override
@@ -208,9 +217,9 @@ private static class ChainedInstrumentationState implements InstrumentationState
208217
private final Map<Instrumentation, InstrumentationState> instrumentationStates;
209218

210219

211-
private ChainedInstrumentationState(List<Instrumentation> instrumentations) {
220+
private ChainedInstrumentationState(List<Instrumentation> instrumentations, InstrumentationCreateStateParameters parameters) {
212221
instrumentationStates = new LinkedHashMap<>(instrumentations.size());
213-
instrumentations.forEach(i -> instrumentationStates.put(i, i.createState()));
222+
instrumentations.forEach(i -> instrumentationStates.put(i, i.createState(parameters)));
214223
}
215224

216225
private InstrumentationState getState(Instrumentation instrumentation) {

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import graphql.ExecutionInput;
44
import graphql.ExecutionResult;
55
import graphql.execution.ExecutionContext;
6+
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
67
import graphql.execution.instrumentation.parameters.InstrumentationDeferredFieldParameters;
78
import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
89
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
@@ -44,6 +45,18 @@ default InstrumentationState createState() {
4445
return null;
4546
}
4647

48+
/**
49+
* This will be called just before execution to create an object that is given back to all instrumentation methods
50+
* to allow them to have per execution request state
51+
*
52+
* @param parameters the parameters to this step
53+
*
54+
* @return a state object that is passed to each method
55+
*/
56+
default InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
57+
return createState();
58+
}
59+
4760
/**
4861
* This is called right at the start of query execution and its the first step in the instrumentation chain.
4962
*

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
* An {@link Instrumentation} implementation can create this as a stateful object that is then passed
55
* to each instrumentation method, allowing state to be passed down with the request execution
66
*
7-
* @see Instrumentation#createState()
7+
* @see Instrumentation#createState(graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters)
88
*/
99
public interface InstrumentationState {
1010
}

0 commit comments

Comments
 (0)