Added a loadImpl function to allow delegate overriding better

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.
-     * <p>
-     * 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<V> 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.
      * <p>
@@ -213,6 +198,24 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
     }
 
 
+    private CompletableFuture<V> 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.
+     * <p>
+     * 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<V> 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.
      * <p>
@@ -229,7 +232,7 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
      * @return the future of the value
      */
     public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
-        return helper.load(nonNull(key), keyContext);
+        return loadImpl(key, keyContext);
     }
 
     /**
@@ -275,7 +278,7 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> 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<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
         for (Map.Entry<K, ?> 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);
     }

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<K, V> 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<V> load(K key) {
+        return delegate.load(key);
+    }
+
     @Override
     public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
         return delegate.load(key, keyContext);
     }
 
+    @Override
+    public CompletableFuture<List<V>> loadMany(List<K> keys) {
+        return delegate.loadMany(keys);
+    }
+
+    @Override
+    public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContexts) {
+        return delegate.loadMany(keys, keyContexts);
+    }
+
+    @Override
+    public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
+        return delegate.loadMany(keysAndContexts);
+    }
+
     @Override
     public DataLoader<K, V> transform(Consumer<DataLoaderFactory.Builder<K, V>> builderConsumer) {
         return delegate.transform(builderConsumer);

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<String, String> rawLoader = TestKit.idLoader();
         DelegatingDataLoader<String, String> delegatingDataLoader = new DelegatingDataLoader<>(rawLoader) {
-            @Override
-            public CompletableFuture<String> load(@NonNull String key, @Nullable Object keyContext) {
-                CompletableFuture<String> cf = super.load(key, keyContext);
+            private CompletableFuture<String> enhance(CompletableFuture<String> cf) {
                 return cf.thenApply(v -> "|" + v + "|");
             }
+
+            private CompletableFuture<List<String>> enhanceList(CompletableFuture<List<String>> cf) {
+                return cf.thenApply(v -> v.stream().map(s -> "|" + s + "|").collect(Collectors.toList()));
+            }
+
+            @Override
+            public CompletableFuture<String> load(String key, @Nullable Object keyContext) {
+                return enhance(super.load(key, keyContext));
+            }
+
+            @Override
+            public CompletableFuture<String> load(String key) {
+                return enhance(super.load(key));
+            }
+
+            @Override
+            public CompletableFuture<List<String>> loadMany(List<String> keys) {
+                return enhanceList(super.loadMany(keys));
+            }
+
+            @Override
+            public CompletableFuture<List<String>> loadMany(List<String> keys, List<Object> keyContexts) {
+                return enhanceList(super.loadMany(keys, keyContexts));
+            }
         };
 
         assertThat(delegatingDataLoader.getDelegate(), is(rawLoader));
@@ -73,8 +98,69 @@ void can_delegate_simple_properties() {
         DelegatingDataLoader<String, String> 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<String, String> loadFunction = CompletableFuture::completedFuture;
+        DataLoader<String, String> rawLoader = DataLoaderFactory.newDataLoader("name", loadFunction, options);
+
+        AtomicInteger loadCalled = new AtomicInteger(0);
+        AtomicInteger loadManyCalled = new AtomicInteger(0);
+        AtomicInteger loadManyMapCalled = new AtomicInteger(0);
+        DelegatingDataLoader<String, String> delegate = new DelegatingDataLoader<>(rawLoader) {
+
+            @Override
+            public CompletableFuture<String> load(String key) {
+                CompletableFuture<String> cf = super.load(key);
+                loadCalled.incrementAndGet();
+                return cf;
+            }
+
+            @Override
+            public CompletableFuture<String> load(String key, @Nullable Object keyContext) {
+                CompletableFuture<String> cf = super.load(key, keyContext);
+                loadCalled.incrementAndGet();
+                return cf;
+            }
+
+            @Override
+            public CompletableFuture<List<String>> loadMany(List<String> keys, List<Object> keyContexts) {
+                CompletableFuture<List<String>> cf = super.loadMany(keys, keyContexts);
+                loadManyCalled.incrementAndGet();
+                return cf;
+            }
+
+            @Override
+            public CompletableFuture<List<String>> loadMany(List<String> keys) {
+                CompletableFuture<List<String>> cf = super.loadMany(keys);
+                loadManyCalled.incrementAndGet();
+                return cf;
+            }
+
+            @Override
+            public CompletableFuture<Map<String, String>> loadMany(Map<String, ?> keysAndContexts) {
+                CompletableFuture<Map<String, String>> 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));
     }
 }