Skip to content

Commit 9812f98

Browse files
dondonzdondonz
authored
Merge pull request #240 from graphql-java/fixed-delegate-load-handling
Added a loadImpl function to allow delegate overriding better
2 parents b2f27b4 + 5a46931 commit 9812f98

File tree

3 files changed

+137
-35
lines changed

3 files changed

+137
-35
lines changed

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

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -160,21 +160,6 @@ public Duration getTimeSinceDispatch() {
160160
return Duration.between(helper.getLastDispatchTime(), helper.now());
161161
}
162162

163-
/**
164-
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
165-
* <p>
166-
* If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
167-
* start batch execution. If you forget this call the future will never be completed (unless already completed,
168-
* and returned from cache).
169-
*
170-
* @param key the key to load
171-
*
172-
* @return the future of the value
173-
*/
174-
public CompletableFuture<V> load(K key) {
175-
return load(key, null);
176-
}
177-
178163
/**
179164
* This will return an optional promise to a value previously loaded via a {@link #load(Object)} call or empty if not call has been made for that key.
180165
* <p>
@@ -213,6 +198,24 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
213198
}
214199

215200

201+
private CompletableFuture<V> loadImpl(@NonNull K key, @Nullable Object keyContext) {
202+
return helper.load(nonNull(key), keyContext);
203+
}
204+
205+
/**
206+
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
207+
* <p>
208+
* If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
209+
* start batch execution. If you forget this call the future will never be completed (unless already completed,
210+
* and returned from cache).
211+
*
212+
* @param key the key to load
213+
* @return the future of the value
214+
*/
215+
public CompletableFuture<V> load(K key) {
216+
return loadImpl(key, null);
217+
}
218+
216219
/**
217220
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
218221
* <p>
@@ -229,7 +232,7 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
229232
* @return the future of the value
230233
*/
231234
public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
232-
return helper.load(nonNull(key), keyContext);
235+
return loadImpl(key, keyContext);
233236
}
234237

235238
/**
@@ -275,7 +278,7 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
275278
if (i < keyContexts.size()) {
276279
keyContext = keyContexts.get(i);
277280
}
278-
collect.add(load(key, keyContext));
281+
collect.add(loadImpl(key, keyContext));
279282
}
280283
return CompletableFutureKit.allOf(collect);
281284
}
@@ -302,7 +305,7 @@ public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
302305
for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
303306
K key = entry.getKey();
304307
Object keyContext = entry.getValue();
305-
collect.put(key, load(key, keyContext));
308+
collect.put(key, loadImpl(key, keyContext));
306309
}
307310
return CompletableFutureKit.allOf(collect);
308311
}

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

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import java.time.Duration;
1010
import java.time.Instant;
1111
import java.util.List;
12+
import java.util.Map;
1213
import java.util.Optional;
1314
import java.util.concurrent.CompletableFuture;
1415
import java.util.function.BiConsumer;
@@ -66,19 +67,31 @@ public DataLoader<K, V> getDelegate() {
6667
return delegate;
6768
}
6869

69-
/**
70-
* The {@link DataLoader#load(Object)} and {@link DataLoader#loadMany(List)} type methods all call back
71-
* to the {@link DataLoader#load(Object, Object)} and hence we don't override them.
72-
*
73-
* @param key the key to load
74-
* @param keyContext a context object that is specific to this key
75-
* @return the future of the value
76-
*/
70+
@Override
71+
public CompletableFuture<V> load(K key) {
72+
return delegate.load(key);
73+
}
74+
7775
@Override
7876
public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
7977
return delegate.load(key, keyContext);
8078
}
8179

80+
@Override
81+
public CompletableFuture<List<V>> loadMany(List<K> keys) {
82+
return delegate.loadMany(keys);
83+
}
84+
85+
@Override
86+
public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContexts) {
87+
return delegate.loadMany(keys, keyContexts);
88+
}
89+
90+
@Override
91+
public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
92+
return delegate.loadMany(keysAndContexts);
93+
}
94+
8295
@Override
8396
public DataLoader<K, V> transform(Consumer<DataLoaderFactory.Builder<K, V>> builderConsumer) {
8497
return delegate.transform(builderConsumer);

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

Lines changed: 95 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,21 @@
22

33
import org.dataloader.fixtures.TestKit;
44
import org.dataloader.fixtures.parameterized.DelegatingDataLoaderFactory;
5-
import org.jspecify.annotations.NonNull;
5+
import org.jspecify.annotations.NullMarked;
66
import org.jspecify.annotations.Nullable;
7-
import org.junit.jupiter.api.Assertions;
87
import org.junit.jupiter.api.Test;
98

109
import java.util.List;
10+
import java.util.Map;
1111
import java.util.concurrent.CompletableFuture;
12+
import java.util.concurrent.atomic.AtomicInteger;
13+
import java.util.stream.Collectors;
1214

1315
import static org.awaitility.Awaitility.await;
1416
import static org.hamcrest.CoreMatchers.equalTo;
1517
import static org.hamcrest.CoreMatchers.is;
1618
import static org.hamcrest.MatcherAssert.assertThat;
17-
import static org.junit.jupiter.api.Assertions.*;
19+
import static org.junit.jupiter.api.Assertions.assertNotNull;
1820

1921
/**
2022
* There are WAY more tests via the {@link DelegatingDataLoaderFactory}
@@ -32,14 +34,37 @@ void canUnwrapDataLoaders() {
3234
}
3335

3436
@Test
37+
@NullMarked
3538
void canCreateAClassOk() {
3639
DataLoader<String, String> rawLoader = TestKit.idLoader();
3740
DelegatingDataLoader<String, String> delegatingDataLoader = new DelegatingDataLoader<>(rawLoader) {
38-
@Override
39-
public CompletableFuture<String> load(@NonNull String key, @Nullable Object keyContext) {
40-
CompletableFuture<String> cf = super.load(key, keyContext);
41+
private CompletableFuture<String> enhance(CompletableFuture<String> cf) {
4142
return cf.thenApply(v -> "|" + v + "|");
4243
}
44+
45+
private CompletableFuture<List<String>> enhanceList(CompletableFuture<List<String>> cf) {
46+
return cf.thenApply(v -> v.stream().map(s -> "|" + s + "|").collect(Collectors.toList()));
47+
}
48+
49+
@Override
50+
public CompletableFuture<String> load(String key, @Nullable Object keyContext) {
51+
return enhance(super.load(key, keyContext));
52+
}
53+
54+
@Override
55+
public CompletableFuture<String> load(String key) {
56+
return enhance(super.load(key));
57+
}
58+
59+
@Override
60+
public CompletableFuture<List<String>> loadMany(List<String> keys) {
61+
return enhanceList(super.loadMany(keys));
62+
}
63+
64+
@Override
65+
public CompletableFuture<List<String>> loadMany(List<String> keys, List<Object> keyContexts) {
66+
return enhanceList(super.loadMany(keys, keyContexts));
67+
}
4368
};
4469

4570
assertThat(delegatingDataLoader.getDelegate(), is(rawLoader));
@@ -73,8 +98,69 @@ void can_delegate_simple_properties() {
7398
DelegatingDataLoader<String, String> delegate = new DelegatingDataLoader<>(rawLoader);
7499

75100
assertNotNull(delegate.getName());
76-
assertThat(delegate.getName(),equalTo("name"));
77-
assertThat(delegate.getOptions(),equalTo(options));
78-
assertThat(delegate.getBatchLoadFunction(),equalTo(loadFunction));
101+
assertThat(delegate.getName(), equalTo("name"));
102+
assertThat(delegate.getOptions(), equalTo(options));
103+
assertThat(delegate.getBatchLoadFunction(), equalTo(loadFunction));
104+
}
105+
106+
@NullMarked
107+
@Test
108+
void can_create_a_delegate_class_that_has_post_side_effects() {
109+
DataLoaderOptions options = DataLoaderOptions.newOptions().build();
110+
BatchLoader<String, String> loadFunction = CompletableFuture::completedFuture;
111+
DataLoader<String, String> rawLoader = DataLoaderFactory.newDataLoader("name", loadFunction, options);
112+
113+
AtomicInteger loadCalled = new AtomicInteger(0);
114+
AtomicInteger loadManyCalled = new AtomicInteger(0);
115+
AtomicInteger loadManyMapCalled = new AtomicInteger(0);
116+
DelegatingDataLoader<String, String> delegate = new DelegatingDataLoader<>(rawLoader) {
117+
118+
@Override
119+
public CompletableFuture<String> load(String key) {
120+
CompletableFuture<String> cf = super.load(key);
121+
loadCalled.incrementAndGet();
122+
return cf;
123+
}
124+
125+
@Override
126+
public CompletableFuture<String> load(String key, @Nullable Object keyContext) {
127+
CompletableFuture<String> cf = super.load(key, keyContext);
128+
loadCalled.incrementAndGet();
129+
return cf;
130+
}
131+
132+
@Override
133+
public CompletableFuture<List<String>> loadMany(List<String> keys, List<Object> keyContexts) {
134+
CompletableFuture<List<String>> cf = super.loadMany(keys, keyContexts);
135+
loadManyCalled.incrementAndGet();
136+
return cf;
137+
}
138+
139+
@Override
140+
public CompletableFuture<List<String>> loadMany(List<String> keys) {
141+
CompletableFuture<List<String>> cf = super.loadMany(keys);
142+
loadManyCalled.incrementAndGet();
143+
return cf;
144+
}
145+
146+
@Override
147+
public CompletableFuture<Map<String, String>> loadMany(Map<String, ?> keysAndContexts) {
148+
CompletableFuture<Map<String, String>> cf = super.loadMany(keysAndContexts);
149+
loadManyMapCalled.incrementAndGet();
150+
return cf;
151+
}
152+
};
153+
154+
155+
delegate.load("L1");
156+
delegate.load("L2", null);
157+
delegate.loadMany(List.of("LM1", "LM2"), List.of());
158+
delegate.loadMany(List.of("LM3"));
159+
delegate.loadMany(Map.of("LMM1", "kc1", "LMM2", "kc2"));
160+
161+
assertNotNull(delegate.getDelegate());
162+
assertThat(loadCalled.get(), equalTo(2));
163+
assertThat(loadManyCalled.get(), equalTo(2));
164+
assertThat(loadManyMapCalled.get(), equalTo(1));
79165
}
80166
}

0 commit comments

Comments
 (0)