Skip to content

Commit 1eca677

Browse files
rstoyanchevrstoyanchev
committed
Support for default DataLoaderOptions
Closes gh-522
1 parent 0748179 commit 1eca677

File tree

3 files changed

+77
-35
lines changed

3 files changed

+77
-35
lines changed

‎spring-graphql-docs/src/docs/asciidoc/index.adoc‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -715,12 +715,16 @@ By default, the `DataLoader` name is based on the class name of the target entit
715715
This allows an `@SchemaMapping` method to declare a
716716
<<controllers-schema-mapping-data-loader,DataLoader argument>> with a generic type, and
717717
without the need for specifying a name. The name, however, can be customized through the
718-
`BatchLoaderRegistry` builder, if necessary, along with other `DataLoader` options.
718+
`BatchLoaderRegistry` builder, if necessary, along with other `DataLoaderOptions`.
719+
720+
To configure default `DataLoaderOptions` globally, to use as a starting point for any
721+
registration, you can override Boot's `BatchLoaderRegistry` bean and use the constructor
722+
for `DefaultBatchLoaderRegistry` that accepts `Supplier<DataLoaderOptions>`.
719723

720724
For many cases, when loading related entities, you can use
721725
<<controllers-batch-mapping,@BatchMapping>> controller methods, which are a shortcut
722726
for and replace the need to use `BatchLoaderRegistry` and `DataLoader` directly.
723-
s
727+
724728
`BatchLoaderRegistry` provides other important benefits too. It supports access to
725729
the same `GraphQLContext` from batch loading functions and from `@BatchMapping` methods,
726730
as well as ensures <<execution-context>> to them. This is why applications are expected

‎spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java‎

Lines changed: 65 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2021 the original author or authors.
2+
* Copyright 2002-2022 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
2323
import java.util.concurrent.CompletionStage;
2424
import java.util.function.BiFunction;
2525
import java.util.function.Consumer;
26+
import java.util.function.Supplier;
2627

2728
import graphql.GraphQLContext;
2829
import io.micrometer.context.ContextSnapshot;
@@ -55,6 +56,25 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
5556

5657
private final List<ReactorMappedBatchLoader<?,?>> mappedLoaders = new ArrayList<>();
5758

59+
private final Supplier<DataLoaderOptions> defaultOptionsSupplier;
60+
61+
62+
/**
63+
* Default constructor
64+
*/
65+
public DefaultBatchLoaderRegistry() {
66+
this(DataLoaderOptions::newOptions);
67+
}
68+
69+
/**
70+
* Constructor with a default {@link DataLoaderOptions} supplier to use as
71+
* a starting point for all registrations.
72+
* @since 1.1
73+
*/
74+
public DefaultBatchLoaderRegistry(Supplier<DataLoaderOptions> defaultOptionsSupplier) {
75+
this.defaultOptionsSupplier = defaultOptionsSupplier;
76+
}
77+
5878

5979
@Override
6080
public <K, V> RegistrationSpec<K, V> forTypePair(Class<K> keyType, Class<V> valueType) {
@@ -69,14 +89,16 @@ public <K, V> RegistrationSpec<K, V> forName(String name) {
6989
@Override
7090
public void registerDataLoaders(DataLoaderRegistry registry, GraphQLContext context) {
7191
BatchLoaderContextProvider contextProvider = () -> context;
72-
DataLoaderOptions defaultOptions = DataLoaderOptions.newOptions().setBatchLoaderContextProvider(contextProvider);
92+
DataLoaderOptions defaultOptions = this.defaultOptionsSupplier.get();
7393
for (ReactorBatchLoader<?, ?> loader : this.loaders) {
74-
DataLoaderOptions options = loader.getOptionsOrDefault(contextProvider, defaultOptions);
94+
DataLoaderOptions options = loader.getOptions();
95+
options = (options != null ? options : defaultOptions).setBatchLoaderContextProvider(contextProvider);
7596
DataLoader<?, ?> dataLoader = DataLoaderFactory.newDataLoader(loader, options);
7697
registerDataLoader(loader.getName(), dataLoader, registry);
7798
}
7899
for (ReactorMappedBatchLoader<?, ?> loader : this.mappedLoaders) {
79-
DataLoaderOptions options = loader.getOptionsOrDefault(contextProvider, defaultOptions);
100+
DataLoaderOptions options = loader.getOptions();
101+
options = (options != null ? options : defaultOptions).setBatchLoaderContextProvider(contextProvider);
80102
DataLoader<?, ?> dataLoader = DataLoaderFactory.newMappedDataLoader(loader, options);
81103
registerDataLoader(loader.getName(), dataLoader, registry);
82104
}
@@ -101,6 +123,9 @@ private class DefaultRegistrationSpec<K, V> implements RegistrationSpec<K, V> {
101123
@Nullable
102124
private DataLoaderOptions options;
103125

126+
@Nullable
127+
private Consumer<DataLoaderOptions> optionsConsumer;
128+
104129
public DefaultRegistrationSpec(Class<V> valueType) {
105130
this.valueType = valueType;
106131
}
@@ -118,8 +143,8 @@ public RegistrationSpec<K, V> withName(String name) {
118143

119144
@Override
120145
public RegistrationSpec<K, V> withOptions(Consumer<DataLoaderOptions> optionsConsumer) {
121-
this.options = (this.options != null ? this.options : DataLoaderOptions.newOptions());
122-
optionsConsumer.accept(this.options);
146+
this.optionsConsumer = (this.optionsConsumer != null ?
147+
this.optionsConsumer.andThen(optionsConsumer) : optionsConsumer);
123148
return this;
124149
}
125150

@@ -132,13 +157,33 @@ public RegistrationSpec<K, V> withOptions(DataLoaderOptions options) {
132157
@Override
133158
public void registerBatchLoader(BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader) {
134159
DefaultBatchLoaderRegistry.this.loaders.add(
135-
new ReactorBatchLoader<>(initName(), loader, this.options));
160+
new ReactorBatchLoader<>(initName(), loader, initOptionsSupplier()));
136161
}
137162

138163
@Override
139164
public void registerMappedBatchLoader(BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader) {
140165
DefaultBatchLoaderRegistry.this.mappedLoaders.add(
141-
new ReactorMappedBatchLoader<>(initName(), loader, this.options));
166+
new ReactorMappedBatchLoader<>(initName(), loader, initOptionsSupplier()));
167+
}
168+
169+
@Nullable
170+
private Supplier<DataLoaderOptions> initOptionsSupplier() {
171+
if (this.options == null && this.optionsConsumer == null) {
172+
return null;
173+
}
174+
175+
Supplier<DataLoaderOptions> optionsSupplier =
176+
(this.options != null ? () -> this.options : defaultOptionsSupplier);
177+
178+
if (this.optionsConsumer == null) {
179+
return optionsSupplier;
180+
}
181+
182+
return () -> {
183+
DataLoaderOptions options = optionsSupplier.get();
184+
this.optionsConsumer.accept(options);
185+
return options;
186+
};
142187
}
143188

144189
private String initName() {
@@ -162,29 +207,24 @@ private static class ReactorBatchLoader<K, V> implements BatchLoaderWithContext<
162207
private final BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader;
163208

164209
@Nullable
165-
private final DataLoaderOptions options;
210+
private final Supplier<DataLoaderOptions> optionsSupplier;
166211

167212
private ReactorBatchLoader(String name,
168213
BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader,
169-
@Nullable DataLoaderOptions options) {
214+
@Nullable Supplier<DataLoaderOptions> optionsSupplier) {
170215

171216
this.name = name;
172217
this.loader = loader;
173-
this.options = options;
218+
this.optionsSupplier = optionsSupplier;
174219
}
175220

176221
public String getName() {
177222
return this.name;
178223
}
179224

180-
public DataLoaderOptions getOptionsOrDefault(
181-
BatchLoaderContextProvider provider, DataLoaderOptions defaultOptions) {
182-
183-
if (this.options != null) {
184-
return new DataLoaderOptions(this.options).setBatchLoaderContextProvider(provider);
185-
}
186-
187-
return defaultOptions;
225+
@Nullable
226+
public DataLoaderOptions getOptions() {
227+
return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
188228
}
189229

190230
@Override
@@ -217,29 +257,24 @@ private static class ReactorMappedBatchLoader<K, V> implements MappedBatchLoader
217257
private final BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader;
218258

219259
@Nullable
220-
private final DataLoaderOptions options;
260+
private final Supplier<DataLoaderOptions> optionsSupplier;
221261

222262
private ReactorMappedBatchLoader(String name,
223263
BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader,
224-
@Nullable DataLoaderOptions options) {
264+
@Nullable Supplier<DataLoaderOptions> optionsSupplier) {
225265

226266
this.name = name;
227267
this.loader = loader;
228-
this.options = options;
268+
this.optionsSupplier = optionsSupplier;
229269
}
230270

231271
public String getName() {
232272
return this.name;
233273
}
234274

235-
public DataLoaderOptions getOptionsOrDefault(
236-
BatchLoaderContextProvider provider, DataLoaderOptions defaultOptions) {
237-
238-
if (this.options != null) {
239-
return new DataLoaderOptions(this.options).setBatchLoaderContextProvider(provider);
240-
}
241-
242-
return defaultOptions;
275+
@Nullable
276+
public DataLoaderOptions getOptions() {
277+
return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
243278
}
244279

245280
@Override

‎spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import graphql.ExecutionInput;
2323
import graphql.GraphQLContext;
2424
import org.dataloader.DataLoader;
25+
import org.dataloader.DataLoaderOptions;
2526
import org.dataloader.DataLoaderRegistry;
2627
import org.dataloader.stats.NoOpStatisticsCollector;
2728
import org.dataloader.stats.StatisticsCollector;
@@ -40,7 +41,11 @@
4041
*/
4142
public class DefaultBatchLoaderRegistryTests {
4243

43-
private final BatchLoaderRegistry batchLoaderRegistry = new DefaultBatchLoaderRegistry();
44+
private final BatchLoaderRegistry batchLoaderRegistry =
45+
new DefaultBatchLoaderRegistry(() -> {
46+
// Disable batching, so we can test loading immediately
47+
return DataLoaderOptions.newOptions().setBatchingEnabled(false);
48+
});
4449

4550
private final DataLoaderRegistry dataLoaderRegistry = DataLoaderRegistry.newRegistry().build();
4651

@@ -50,7 +55,6 @@ void batchLoader() throws Exception {
5055
AtomicReference<String> valueRef = new AtomicReference<>();
5156

5257
this.batchLoaderRegistry.forTypePair(Long.class, Book.class)
53-
.withOptions(options -> options.setBatchingEnabled(false)) // DataLoader invoked immediately
5458
.registerBatchLoader((ids, environment) ->
5559
Flux.deferContextual(contextView -> {
5660
valueRef.set(contextView.get("key"));
@@ -76,7 +80,6 @@ void mappedBatchLoader() throws Exception {
7680
AtomicReference<String> valueRef = new AtomicReference<>();
7781

7882
this.batchLoaderRegistry.forTypePair(Long.class, Book.class)
79-
.withOptions(options -> options.setBatchingEnabled(false)) // DataLoader invoked immediately
8083
.registerMappedBatchLoader((ids, environment) ->
8184
Mono.deferContextual(contextView -> {
8285
valueRef.set(contextView.get("key"));

0 commit comments

Comments
 (0)