diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java
index d80823f..125efe1 100644
--- a/src/main/java/org/dataloader/DataLoader.java
+++ b/src/main/java/org/dataloader/DataLoader.java
@@ -160,21 +160,6 @@ public Duration getTimeSinceDispatch() {
return Duration.between(helper.getLastDispatchTime(), helper.now());
}
- /**
- * Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
- *
- * If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
- * start batch execution. If you forget this call the future will never be completed (unless already completed,
- * and returned from cache).
- *
- * @param key the key to load
- *
- * @return the future of the value
- */
- public CompletableFuture load(K key) {
- return load(key, null);
- }
-
/**
* 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.
*
@@ -213,6 +198,24 @@ public Optional> getIfCompleted(K key) {
}
+ private CompletableFuture loadImpl(@NonNull K key, @Nullable Object keyContext) {
+ return helper.load(nonNull(key), keyContext);
+ }
+
+ /**
+ * Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
+ *
+ * If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
+ * start batch execution. If you forget this call the future will never be completed (unless already completed,
+ * and returned from cache).
+ *
+ * @param key the key to load
+ * @return the future of the value
+ */
+ public CompletableFuture load(K key) {
+ return loadImpl(key, null);
+ }
+
/**
* Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
*
@@ -229,7 +232,7 @@ public Optional> getIfCompleted(K key) {
* @return the future of the value
*/
public CompletableFuture load(@NonNull K key, @Nullable Object keyContext) {
- return helper.load(nonNull(key), keyContext);
+ return loadImpl(key, keyContext);
}
/**
@@ -275,7 +278,7 @@ public CompletableFuture> loadMany(List keys, List keyContext
if (i < keyContexts.size()) {
keyContext = keyContexts.get(i);
}
- collect.add(load(key, keyContext));
+ collect.add(loadImpl(key, keyContext));
}
return CompletableFutureKit.allOf(collect);
}
@@ -302,7 +305,7 @@ public CompletableFuture> loadMany(Map keysAndContexts) {
for (Map.Entry entry : keysAndContexts.entrySet()) {
K key = entry.getKey();
Object keyContext = entry.getValue();
- collect.put(key, load(key, keyContext));
+ collect.put(key, loadImpl(key, keyContext));
}
return CompletableFutureKit.allOf(collect);
}
diff --git a/src/main/java/org/dataloader/DelegatingDataLoader.java b/src/main/java/org/dataloader/DelegatingDataLoader.java
index 7cffced..0175059 100644
--- a/src/main/java/org/dataloader/DelegatingDataLoader.java
+++ b/src/main/java/org/dataloader/DelegatingDataLoader.java
@@ -9,6 +9,7 @@
import java.time.Duration;
import java.time.Instant;
import java.util.List;
+import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.function.BiConsumer;
@@ -66,19 +67,31 @@ public DataLoader getDelegate() {
return delegate;
}
- /**
- * The {@link DataLoader#load(Object)} and {@link DataLoader#loadMany(List)} type methods all call back
- * to the {@link DataLoader#load(Object, Object)} and hence we don't override them.
- *
- * @param key the key to load
- * @param keyContext a context object that is specific to this key
- * @return the future of the value
- */
+ @Override
+ public CompletableFuture load(K key) {
+ return delegate.load(key);
+ }
+
@Override
public CompletableFuture load(@NonNull K key, @Nullable Object keyContext) {
return delegate.load(key, keyContext);
}
+ @Override
+ public CompletableFuture> loadMany(List keys) {
+ return delegate.loadMany(keys);
+ }
+
+ @Override
+ public CompletableFuture> loadMany(List keys, List keyContexts) {
+ return delegate.loadMany(keys, keyContexts);
+ }
+
+ @Override
+ public CompletableFuture> loadMany(Map keysAndContexts) {
+ return delegate.loadMany(keysAndContexts);
+ }
+
@Override
public DataLoader transform(Consumer> builderConsumer) {
return delegate.transform(builderConsumer);
diff --git a/src/test/java/org/dataloader/DelegatingDataLoaderTest.java b/src/test/java/org/dataloader/DelegatingDataLoaderTest.java
index 8849752..0f51b5f 100644
--- a/src/test/java/org/dataloader/DelegatingDataLoaderTest.java
+++ b/src/test/java/org/dataloader/DelegatingDataLoaderTest.java
@@ -2,19 +2,21 @@
import org.dataloader.fixtures.TestKit;
import org.dataloader.fixtures.parameterized.DelegatingDataLoaderFactory;
-import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.NullMarked;
import org.jspecify.annotations.Nullable;
-import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import java.util.List;
+import java.util.Map;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
import static org.awaitility.Awaitility.await;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
/**
* There are WAY more tests via the {@link DelegatingDataLoaderFactory}
@@ -32,14 +34,37 @@ void canUnwrapDataLoaders() {
}
@Test
+ @NullMarked
void canCreateAClassOk() {
DataLoader rawLoader = TestKit.idLoader();
DelegatingDataLoader delegatingDataLoader = new DelegatingDataLoader<>(rawLoader) {
- @Override
- public CompletableFuture load(@NonNull String key, @Nullable Object keyContext) {
- CompletableFuture cf = super.load(key, keyContext);
+ private CompletableFuture enhance(CompletableFuture cf) {
return cf.thenApply(v -> "|" + v + "|");
}
+
+ private CompletableFuture> enhanceList(CompletableFuture> cf) {
+ return cf.thenApply(v -> v.stream().map(s -> "|" + s + "|").collect(Collectors.toList()));
+ }
+
+ @Override
+ public CompletableFuture load(String key, @Nullable Object keyContext) {
+ return enhance(super.load(key, keyContext));
+ }
+
+ @Override
+ public CompletableFuture load(String key) {
+ return enhance(super.load(key));
+ }
+
+ @Override
+ public CompletableFuture> loadMany(List keys) {
+ return enhanceList(super.loadMany(keys));
+ }
+
+ @Override
+ public CompletableFuture> loadMany(List keys, List keyContexts) {
+ return enhanceList(super.loadMany(keys, keyContexts));
+ }
};
assertThat(delegatingDataLoader.getDelegate(), is(rawLoader));
@@ -73,8 +98,69 @@ void can_delegate_simple_properties() {
DelegatingDataLoader delegate = new DelegatingDataLoader<>(rawLoader);
assertNotNull(delegate.getName());
- assertThat(delegate.getName(),equalTo("name"));
- assertThat(delegate.getOptions(),equalTo(options));
- assertThat(delegate.getBatchLoadFunction(),equalTo(loadFunction));
+ assertThat(delegate.getName(), equalTo("name"));
+ assertThat(delegate.getOptions(), equalTo(options));
+ assertThat(delegate.getBatchLoadFunction(), equalTo(loadFunction));
+ }
+
+ @NullMarked
+ @Test
+ void can_create_a_delegate_class_that_has_post_side_effects() {
+ DataLoaderOptions options = DataLoaderOptions.newOptions().build();
+ BatchLoader loadFunction = CompletableFuture::completedFuture;
+ DataLoader rawLoader = DataLoaderFactory.newDataLoader("name", loadFunction, options);
+
+ AtomicInteger loadCalled = new AtomicInteger(0);
+ AtomicInteger loadManyCalled = new AtomicInteger(0);
+ AtomicInteger loadManyMapCalled = new AtomicInteger(0);
+ DelegatingDataLoader delegate = new DelegatingDataLoader<>(rawLoader) {
+
+ @Override
+ public CompletableFuture load(String key) {
+ CompletableFuture cf = super.load(key);
+ loadCalled.incrementAndGet();
+ return cf;
+ }
+
+ @Override
+ public CompletableFuture load(String key, @Nullable Object keyContext) {
+ CompletableFuture cf = super.load(key, keyContext);
+ loadCalled.incrementAndGet();
+ return cf;
+ }
+
+ @Override
+ public CompletableFuture> loadMany(List keys, List keyContexts) {
+ CompletableFuture> cf = super.loadMany(keys, keyContexts);
+ loadManyCalled.incrementAndGet();
+ return cf;
+ }
+
+ @Override
+ public CompletableFuture> loadMany(List keys) {
+ CompletableFuture> cf = super.loadMany(keys);
+ loadManyCalled.incrementAndGet();
+ return cf;
+ }
+
+ @Override
+ public CompletableFuture> loadMany(Map keysAndContexts) {
+ CompletableFuture> cf = super.loadMany(keysAndContexts);
+ loadManyMapCalled.incrementAndGet();
+ return cf;
+ }
+ };
+
+
+ delegate.load("L1");
+ delegate.load("L2", null);
+ delegate.loadMany(List.of("LM1", "LM2"), List.of());
+ delegate.loadMany(List.of("LM3"));
+ delegate.loadMany(Map.of("LMM1", "kc1", "LMM2", "kc2"));
+
+ assertNotNull(delegate.getDelegate());
+ assertThat(loadCalled.get(), equalTo(2));
+ assertThat(loadManyCalled.get(), equalTo(2));
+ assertThat(loadManyMapCalled.get(), equalTo(1));
}
}
\ No newline at end of file