Skip to content

Commit af3ebf4

Browse files
AlexandreCarltonAlexandreCarlton
committed
Add strict mode for DataLoaderRegistry
In a similar vein to graphql-java/graphql-java#3565 which enabled "strict mode" for type wiring (preventing multiple `DataFetcher`s being registered to the same field on a GraphQL type), we add "strict mode" to the `DataLoaderRegistry` so that we don't accidentally register multiple DataLoaders to the same key (which would result in the last registration winning). This defaults to `false` to avoid breaking changes.
1 parent 918203c commit af3ebf4

File tree

5 files changed

+201
-7
lines changed

5 files changed

+201
-7
lines changed

‎src/main/java/org/dataloader/DataLoaderRegistry.java‎

Lines changed: 57 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.dataloader;
22

33
import org.dataloader.annotations.PublicApi;
4+
import org.dataloader.errors.StrictModeRegistryException;
45
import org.dataloader.impl.Assertions;
56
import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation;
67
import org.dataloader.instrumentation.DataLoaderInstrumentation;
@@ -20,6 +21,7 @@
2021
import java.util.concurrent.ConcurrentHashMap;
2122
import java.util.function.Function;
2223

24+
import static java.lang.String.format;
2325
import static org.dataloader.impl.Assertions.assertState;
2426

2527
/**
@@ -43,21 +45,28 @@
4345
@PublicApi
4446
@NullMarked
4547
public class DataLoaderRegistry {
48+
4649
protected final Map<String, DataLoader<?, ?>> dataLoaders;
4750
protected final @Nullable DataLoaderInstrumentation instrumentation;
51+
protected final boolean strictMode;
4852

4953

5054
public DataLoaderRegistry() {
51-
this(new ConcurrentHashMap<>(), null);
55+
this(new ConcurrentHashMap<>(), null, false);
5256
}
5357

5458
private DataLoaderRegistry(Builder builder) {
55-
this(builder.dataLoaders, builder.instrumentation);
59+
this(builder.dataLoaders, builder.instrumentation, builder.strictMode);
5660
}
5761

58-
protected DataLoaderRegistry(Map<String, DataLoader<?, ?>> dataLoaders, @Nullable DataLoaderInstrumentation instrumentation) {
62+
protected DataLoaderRegistry(
63+
Map<String, DataLoader<?, ?>> dataLoaders,
64+
@Nullable DataLoaderInstrumentation instrumentation,
65+
boolean strictMode
66+
) {
5967
this.dataLoaders = instrumentDLs(dataLoaders, instrumentation);
6068
this.instrumentation = instrumentation;
69+
this.strictMode = strictMode;
6170
}
6271

6372
private Map<String, DataLoader<?, ?>> instrumentDLs(Map<String, DataLoader<?, ?>> incomingDataLoaders, @Nullable DataLoaderInstrumentation registryInstrumentation) {
@@ -144,6 +153,9 @@ private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options,
144153
*/
145154
public DataLoaderRegistry register(DataLoader<?, ?> dataLoader) {
146155
String name = Assertions.nonNull(dataLoader.getName(), () -> "The DataLoader must have a non null name");
156+
if (strictMode) {
157+
assertKeyStrictly(name);
158+
}
147159
dataLoaders.put(name, nameAndInstrumentDL(name, instrumentation, dataLoader));
148160
return this;
149161
}
@@ -160,6 +172,9 @@ public DataLoaderRegistry register(DataLoader<?, ?> dataLoader) {
160172
* @return this registry
161173
*/
162174
public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
175+
if (strictMode) {
176+
assertKeyStrictly(key);
177+
}
163178
dataLoaders.put(key, nameAndInstrumentDL(key, instrumentation, dataLoader));
164179
return this;
165180
}
@@ -176,6 +191,9 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
176191
* @return the data loader instance that was registered
177192
*/
178193
public <K, V> DataLoader<K, V> registerAndGet(String key, DataLoader<?, ?> dataLoader) {
194+
if (strictMode) {
195+
assertKeyStrictly(key);
196+
}
179197
dataLoaders.put(key, nameAndInstrumentDL(key, instrumentation, dataLoader));
180198
return Objects.requireNonNull(getDataLoader(key));
181199
}
@@ -214,7 +232,9 @@ public <K, V> DataLoader<K, V> computeIfAbsent(final String key,
214232
* @return a new combined registry
215233
*/
216234
public DataLoaderRegistry combine(DataLoaderRegistry registry) {
217-
DataLoaderRegistry combined = new DataLoaderRegistry();
235+
DataLoaderRegistry combined = new Builder()
236+
.strictMode(strictMode)
237+
.build();
218238

219239
this.dataLoaders.forEach(combined::register);
220240
registry.dataLoaders.forEach(combined::register);
@@ -312,6 +332,12 @@ public Statistics getStatistics() {
312332
return stats;
313333
}
314334

335+
protected void assertKeyStrictly(String key) {
336+
if (dataLoaders.containsKey(key)) {
337+
throw new StrictModeRegistryException(format("The key %s already has a DataLoader defined", key));
338+
}
339+
}
340+
315341
/**
316342
* @return A builder of {@link DataLoaderRegistry}s
317343
*/
@@ -323,6 +349,19 @@ public static class Builder {
323349

324350
private final Map<String, DataLoader<?, ?>> dataLoaders = new HashMap<>();
325351
private @Nullable DataLoaderInstrumentation instrumentation;
352+
private boolean strictMode;
353+
354+
/**
355+
* This puts the builder into strict mode, so if things get defined twice, for example, it
356+
* will throw a {@link org.dataloader.errors.StrictModeRegistryException}.
357+
*
358+
* @param strictMode whether strict mode is enabled
359+
* @return this builder
360+
*/
361+
public Builder strictMode(boolean strictMode) {
362+
this.strictMode = strictMode;
363+
return this;
364+
}
326365

327366
/**
328367
* This will register a new dataloader
@@ -332,6 +371,9 @@ public static class Builder {
332371
* @return this builder for a fluent pattern
333372
*/
334373
public Builder register(String key, DataLoader<?, ?> dataLoader) {
374+
if (strictMode) {
375+
assertKeyStrictly(key);
376+
}
335377
dataLoaders.put(key, dataLoader);
336378
return this;
337379
}
@@ -344,6 +386,11 @@ public Builder register(String key, DataLoader<?, ?> dataLoader) {
344386
* @return this builder for a fluent pattern
345387
*/
346388
public Builder registerAll(DataLoaderRegistry otherRegistry) {
389+
if (strictMode) {
390+
otherRegistry.dataLoaders.forEach((key, dataLoader) -> {
391+
assertKeyStrictly(key);
392+
});
393+
}
347394
dataLoaders.putAll(otherRegistry.dataLoaders);
348395
return this;
349396
}
@@ -353,6 +400,12 @@ public Builder instrumentation(DataLoaderInstrumentation instrumentation) {
353400
return this;
354401
}
355402

403+
private void assertKeyStrictly(String key) {
404+
if (dataLoaders.containsKey(key)) {
405+
throw new StrictModeRegistryException(format("The key %s already has a DataLoader defined", key));
406+
}
407+
}
408+
356409
/**
357410
* @return the newly built {@link DataLoaderRegistry}
358411
*/

‎src/main/java/org/dataloader/errors/StrictModeRegistryException.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package org.dataloader.errors;
2+
3+
import org.dataloader.annotations.PublicApi;
4+
5+
/**
6+
* An exception that is thrown when {@link org.dataloader.DataLoaderRegistry.Builder#strictMode(boolean)} is true and multiple
7+
* DataLoaders are registered to the same key.
8+
*/
9+
@PublicApi
10+
public class StrictModeRegistryException extends RuntimeException {
11+
public StrictModeRegistryException(String msg) {
12+
super(msg);
13+
}
14+
}

‎src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java‎

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import org.dataloader.DataLoader;
44
import org.dataloader.DataLoaderRegistry;
55
import org.dataloader.annotations.ExperimentalApi;
6+
import org.dataloader.errors.StrictModeRegistryException;
67
import org.dataloader.impl.Assertions;
78
import org.dataloader.instrumentation.DataLoaderInstrumentation;
89
import org.jspecify.annotations.NullMarked;
@@ -16,6 +17,7 @@
1617
import java.util.concurrent.ScheduledExecutorService;
1718
import java.util.concurrent.TimeUnit;
1819

20+
import static java.lang.String.format;
1921
import static org.dataloader.impl.Assertions.nonNull;
2022

2123
/**
@@ -69,7 +71,7 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A
6971
private volatile boolean closed;
7072

7173
private ScheduledDataLoaderRegistry(Builder builder) {
72-
super(builder.dataLoaders, builder.instrumentation);
74+
super(builder.dataLoaders, builder.instrumentation, builder.strictMode);
7375
this.scheduledExecutorService = Assertions.nonNull(builder.scheduledExecutorService);
7476
this.defaultExecutorUsed = builder.defaultExecutorUsed;
7577
this.schedule = builder.schedule;
@@ -120,7 +122,8 @@ public boolean isTickerMode() {
120122
*/
121123
public ScheduledDataLoaderRegistry combine(DataLoaderRegistry registry) {
122124
Builder combinedBuilder = ScheduledDataLoaderRegistry.newScheduledRegistry()
123-
.dispatchPredicate(this.dispatchPredicate);
125+
.dispatchPredicate(this.dispatchPredicate)
126+
.strictMode(this.strictMode);
124127
combinedBuilder.registerAll(this);
125128
combinedBuilder.registerAll(registry);
126129
return combinedBuilder.build();
@@ -166,6 +169,9 @@ public DispatchPredicate getDispatchPredicate() {
166169
* @return this registry
167170
*/
168171
public ScheduledDataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader, DispatchPredicate dispatchPredicate) {
172+
if (strictMode) {
173+
assertKeyStrictly(key);
174+
}
169175
dataLoaders.put(key, dataLoader);
170176
dataLoaderPredicates.put(dataLoader, dispatchPredicate);
171177
return this;
@@ -272,6 +278,7 @@ public static class Builder {
272278
private Duration schedule = Duration.ofMillis(10);
273279
private boolean tickerMode = false;
274280
private @Nullable DataLoaderInstrumentation instrumentation;
281+
private boolean strictMode;
275282

276283

277284
/**
@@ -291,6 +298,18 @@ public Builder schedule(Duration schedule) {
291298
return this;
292299
}
293300

301+
/**
302+
* This puts the builder into strict mode, so if things get defined twice, for example, it
303+
* will throw a {@link org.dataloader.errors.StrictModeRegistryException}.
304+
*
305+
* @param strictMode whether strict mode is enabled
306+
* @return this builder
307+
*/
308+
public Builder strictMode(boolean strictMode) {
309+
this.strictMode = strictMode;
310+
return this;
311+
}
312+
294313
/**
295314
* This will register a new dataloader
296315
*
@@ -299,6 +318,9 @@ public Builder schedule(Duration schedule) {
299318
* @return this builder for a fluent pattern
300319
*/
301320
public Builder register(String key, DataLoader<?, ?> dataLoader) {
321+
if (strictMode) {
322+
assertKeyStrictly(key);
323+
}
302324
dataLoaders.put(key, dataLoader);
303325
return this;
304326
}
@@ -326,7 +348,13 @@ public Builder register(String key, DataLoader<?, ?> dataLoader, DispatchPredica
326348
* @return this builder for a fluent pattern
327349
*/
328350
public Builder registerAll(DataLoaderRegistry otherRegistry) {
329-
dataLoaders.putAll(otherRegistry.getDataLoadersMap());
351+
Map<String, DataLoader<?, ?>> otherDataLoaders = otherRegistry.getDataLoadersMap();
352+
if (strictMode) {
353+
otherDataLoaders.forEach((key, dataLoader) -> {
354+
assertKeyStrictly(key);
355+
});
356+
}
357+
dataLoaders.putAll(otherDataLoaders);
330358
if (otherRegistry instanceof ScheduledDataLoaderRegistry) {
331359
ScheduledDataLoaderRegistry other = (ScheduledDataLoaderRegistry) otherRegistry;
332360
dataLoaderPredicates.putAll(other.dataLoaderPredicates);
@@ -364,6 +392,12 @@ public Builder instrumentation(DataLoaderInstrumentation instrumentation) {
364392
return this;
365393
}
366394

395+
private void assertKeyStrictly(String key) {
396+
if (dataLoaders.containsKey(key)) {
397+
throw new StrictModeRegistryException(format("The key %s already has a DataLoader defined", key));
398+
}
399+
}
400+
367401
/**
368402
* @return the newly built {@link ScheduledDataLoaderRegistry}
369403
*/

‎src/test/java/org/dataloader/DataLoaderRegistryTest.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.dataloader;
22

3+
import org.dataloader.errors.StrictModeRegistryException;
34
import org.dataloader.stats.SimpleStatisticsCollector;
45
import org.dataloader.stats.Statistics;
56
import org.junit.jupiter.api.Assertions;
@@ -14,6 +15,7 @@
1415
import static org.hamcrest.Matchers.equalTo;
1516
import static org.hamcrest.Matchers.hasItems;
1617
import static org.hamcrest.Matchers.sameInstance;
18+
import static org.junit.jupiter.api.Assertions.assertThrows;
1719

1820
public class DataLoaderRegistryTest {
1921
final BatchLoader<Object, Object> identityBatchLoader = CompletableFuture::completedFuture;
@@ -219,4 +221,48 @@ public void builder_works() {
219221
assertThat(registry.getDataLoader("c"), equalTo(dlC));
220222

221223
}
224+
225+
@Test
226+
public void strictMode_works() {
227+
228+
DataLoader<Object, Object> dlA = newDataLoader(identityBatchLoader);
229+
DataLoader<Object, Object> dlB = newDataLoader(identityBatchLoader);
230+
231+
assertThrows(StrictModeRegistryException.class, () -> {
232+
DataLoaderRegistry.newRegistry()
233+
.strictMode(true)
234+
.register("a", dlA)
235+
.register("a", dlB)
236+
.build();
237+
});
238+
assertThrows(StrictModeRegistryException.class, () -> {
239+
DataLoaderRegistry.newRegistry()
240+
.strictMode(true)
241+
.register("a", dlA)
242+
.registerAll(DataLoaderRegistry.newRegistry()
243+
.register("a", dlB)
244+
.build())
245+
.build();
246+
});
247+
248+
DataLoaderRegistry registry = DataLoaderRegistry.newRegistry()
249+
.strictMode(true)
250+
.build();
251+
registry.register("a", dlA);
252+
253+
assertThrows(StrictModeRegistryException.class, () -> {
254+
registry.register("a", dlB);
255+
});
256+
assertThrows(StrictModeRegistryException.class, () -> {
257+
registry.register(newDataLoader("a", identityBatchLoader));
258+
});
259+
assertThrows(StrictModeRegistryException.class, () -> {
260+
registry.registerAndGet("a", dlB);
261+
});
262+
assertThrows(StrictModeRegistryException.class, () -> {
263+
registry.combine(DataLoaderRegistry.newRegistry()
264+
.register("a", dlB)
265+
.build());
266+
});
267+
}
222268
}

0 commit comments

Comments
 (0)