Skip to content

Commit 46950a9

Browse files
bbakermanbbakerman
committed
Added a loadImpl function to allow delegate overriding better
1 parent 6c214be commit 46950a9

File tree

3 files changed

+110
-30
lines changed

3 files changed

+110
-30
lines changed

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

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

162-
/**
163-
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
164-
* <p>
165-
* If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
166-
* start batch execution. If you forget this call the future will never be completed (unless already completed,
167-
* and returned from cache).
168-
*
169-
* @param key the key to load
170-
* @return the future of the value
171-
*/
172-
public CompletableFuture<V> load(K key) {
173-
return load(key, null);
174-
}
175-
176162
/**
177163
* 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.
178164
* <p>
@@ -209,6 +195,24 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
209195
}
210196

211197

198+
private CompletableFuture<V> loadImpl(@NonNull K key, @Nullable Object keyContext) {
199+
return helper.load(nonNull(key), keyContext);
200+
}
201+
202+
/**
203+
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
204+
* <p>
205+
* If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
206+
* start batch execution. If you forget this call the future will never be completed (unless already completed,
207+
* and returned from cache).
208+
*
209+
* @param key the key to load
210+
* @return the future of the value
211+
*/
212+
public CompletableFuture<V> load(K key) {
213+
return loadImpl(key, null);
214+
}
215+
212216
/**
213217
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
214218
* <p>
@@ -224,7 +228,7 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
224228
* @return the future of the value
225229
*/
226230
public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
227-
return helper.load(nonNull(key), keyContext);
231+
return loadImpl(key, keyContext);
228232
}
229233

230234
/**
@@ -269,7 +273,7 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
269273
if (i < keyContexts.size()) {
270274
keyContext = keyContexts.get(i);
271275
}
272-
collect.add(load(key, keyContext));
276+
collect.add(loadImpl(key, keyContext));
273277
}
274278
return CompletableFutureKit.allOf(collect);
275279
}
@@ -297,7 +301,7 @@ public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
297301
for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
298302
K key = entry.getKey();
299303
Object keyContext = entry.getValue();
300-
collect.put(key, load(key, keyContext));
304+
collect.put(key, loadImpl(key, keyContext));
301305
}
302306
return CompletableFutureKit.allOf(collect);
303307
}

‎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: 68 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,20 @@
33
import org.dataloader.fixtures.TestKit;
44
import org.dataloader.fixtures.parameterized.DelegatingDataLoaderFactory;
55
import org.jspecify.annotations.NonNull;
6+
import org.jspecify.annotations.NullMarked;
67
import org.jspecify.annotations.Nullable;
7-
import org.junit.jupiter.api.Assertions;
88
import org.junit.jupiter.api.Test;
99

1010
import java.util.List;
11+
import java.util.Map;
1112
import java.util.concurrent.CompletableFuture;
13+
import java.util.concurrent.atomic.AtomicInteger;
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}
@@ -73,8 +75,69 @@ void can_delegate_simple_properties() {
7375
DelegatingDataLoader<String, String> delegate = new DelegatingDataLoader<>(rawLoader);
7476

7577
assertNotNull(delegate.getName());
76-
assertThat(delegate.getName(),equalTo("name"));
77-
assertThat(delegate.getOptions(),equalTo(options));
78-
assertThat(delegate.getBatchLoadFunction(),equalTo(loadFunction));
78+
assertThat(delegate.getName(), equalTo("name"));
79+
assertThat(delegate.getOptions(), equalTo(options));
80+
assertThat(delegate.getBatchLoadFunction(), equalTo(loadFunction));
81+
}
82+
83+
@NullMarked
84+
@Test
85+
void can_create_a_delegate_class_that_has_post_side_effects() {
86+
DataLoaderOptions options = DataLoaderOptions.newOptions().build();
87+
BatchLoader<String, String> loadFunction = CompletableFuture::completedFuture;
88+
DataLoader<String, String> rawLoader = DataLoaderFactory.newDataLoader("name", loadFunction, options);
89+
90+
AtomicInteger loadCalled = new AtomicInteger(0);
91+
AtomicInteger loadManyCalled = new AtomicInteger(0);
92+
AtomicInteger loadManyMapCalled = new AtomicInteger(0);
93+
DelegatingDataLoader<String, String> delegate = new DelegatingDataLoader<>(rawLoader) {
94+
95+
@Override
96+
public CompletableFuture<String> load(String key) {
97+
CompletableFuture<String> cf = super.load(key);
98+
loadCalled.incrementAndGet();
99+
return cf;
100+
}
101+
102+
@Override
103+
public CompletableFuture<String> load(String key, @Nullable Object keyContext) {
104+
CompletableFuture<String> cf = super.load(key, keyContext);
105+
loadCalled.incrementAndGet();
106+
return cf;
107+
}
108+
109+
@Override
110+
public CompletableFuture<List<String>> loadMany(List<String> keys, List<Object> keyContexts) {
111+
CompletableFuture<List<String>> cf = super.loadMany(keys, keyContexts);
112+
loadManyCalled.incrementAndGet();
113+
return cf;
114+
}
115+
116+
@Override
117+
public CompletableFuture<List<String>> loadMany(List<String> keys) {
118+
CompletableFuture<List<String>> cf = super.loadMany(keys);
119+
loadManyCalled.incrementAndGet();
120+
return cf;
121+
}
122+
123+
@Override
124+
public CompletableFuture<Map<String, String>> loadMany(Map<String, ?> keysAndContexts) {
125+
CompletableFuture<Map<String, String>> cf = super.loadMany(keysAndContexts);
126+
loadManyMapCalled.incrementAndGet();
127+
return cf;
128+
}
129+
};
130+
131+
132+
delegate.load("L1");
133+
delegate.load("L2", null);
134+
delegate.loadMany(List.of("LM1", "LM2"), List.of());
135+
delegate.loadMany(List.of("LM3"));
136+
delegate.loadMany(Map.of("LMM1", "kc1", "LMM2", "kc2"));
137+
138+
assertNotNull(delegate.getDelegate());
139+
assertThat(loadCalled.get(), equalTo(2));
140+
assertThat(loadManyCalled.get(), equalTo(2));
141+
assertThat(loadManyMapCalled.get(), equalTo(1));
79142
}
80143
}

0 commit comments

Comments
 (0)