Added support for a delegating data loader

build.gradle

@@ -176,6 +176,7 @@ nexusPublishing {
 }
 
 signing {
+    required { !project.hasProperty('publishToMavenLocal') }
     def signingKey = System.env.MAVEN_CENTRAL_PGP_KEY
     useInMemoryPgpKeys(signingKey, "")
     sign publishing.publications

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

@@ -21,6 +21,7 @@
 import org.dataloader.impl.CompletableFutureKit;
 import org.dataloader.stats.Statistics;
 import org.dataloader.stats.StatisticsCollector;
+import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -517,8 +518,8 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
      * @param keyContext a context object that is specific to this key
      * @return the future of the value
      */
-    public CompletableFuture<V> load(K key, Object keyContext) {
-        return helper.load(key, keyContext);
+    public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
+        return helper.load(nonNull(key), keyContext);
     }
 
     /**

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

@@ -0,0 +1,188 @@
+package org.dataloader;
+
+import org.dataloader.annotations.PublicApi;
+import org.dataloader.stats.Statistics;
+import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+
+/**
+ * This delegating {@link DataLoader} makes it easier to create wrappers of {@link DataLoader}s in case you want to change how
+ * values are returned for example.
+ * <p>
+ * The most common way would be to make a new {@link DelegatingDataLoader} subclass that overloads the {@link DelegatingDataLoader#load(Object, Object)}
+ * method.
+ * <p>
+ * For example the following allows you to change the returned value in some way :
+ * <pre>{@code
+ * DataLoader<String, String> rawLoader = createDataLoader();
+ * DelegatingDataLoader<String, String> delegatingDataLoader = new DelegatingDataLoader<>(rawLoader) {
+ *    public CompletableFuture<String> load(@NonNull String key, @Nullable Object keyContext) {
+ *       CompletableFuture<String> cf = super.load(key, keyContext);
+ *       return cf.thenApply(v -> "|" + v + "|");
+ *    }
+ *};
+ *}</pre>
+ *
+ * @param <K> type parameter indicating the type of the data load keys
+ * @param <V> type parameter indicating the type of the data that is returned
+ */
+@PublicApi
+@NullMarked
+public class DelegatingDataLoader<K, V> extends DataLoader<K, V> {
+
+    protected final DataLoader<K, V> delegate;
+
+    /**
+     * This can be called to unwrap a given {@link DataLoader} such that if it's a {@link DelegatingDataLoader} the underlying
+     * {@link DataLoader} is returned otherwise it's just passed in data loader
+     *
+     * @param dataLoader the dataLoader to unwrap
+     * @param <K>        type parameter indicating the type of the data load keys
+     * @param <V>        type parameter indicating the type of the data that is returned
+     * @return the delegate dataLoader OR just this current one if it's not wrapped
+     */
+    public static <K, V> DataLoader<K, V> unwrap(DataLoader<K, V> dataLoader) {
+        if (dataLoader instanceof DelegatingDataLoader) {
+            return ((DelegatingDataLoader<K, V>) dataLoader).getDelegate();
+        }
+        return dataLoader;
+    }
+
+    public DelegatingDataLoader(DataLoader<K, V> delegate) {
+        super(delegate.getBatchLoadFunction(), delegate.getOptions());
+        this.delegate = delegate;
+    }
+
+    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(@NonNull K key, @Nullable Object keyContext) {
+        return delegate.load(key, keyContext);
+    }
+
+    @Override
+    public DataLoader<K, V> transform(Consumer<DataLoaderFactory.Builder<K, V>> builderConsumer) {
+        return delegate.transform(builderConsumer);
+    }
+
+    @Override
+    public Instant getLastDispatchTime() {
+        return delegate.getLastDispatchTime();
+    }
+
+    @Override
+    public Duration getTimeSinceDispatch() {
+        return delegate.getTimeSinceDispatch();
+    }
+
+    @Override
+    public Optional<CompletableFuture<V>> getIfPresent(K key) {
+        return delegate.getIfPresent(key);
+    }
+
+    @Override
+    public Optional<CompletableFuture<V>> getIfCompleted(K key) {
+        return delegate.getIfCompleted(key);
+    }
+
+    @Override
+    public CompletableFuture<List<V>> dispatch() {
+        return delegate.dispatch();
+    }
+
+    @Override
+    public DispatchResult<V> dispatchWithCounts() {
+        return delegate.dispatchWithCounts();
+    }
+
+    @Override
+    public List<V> dispatchAndJoin() {
+        return delegate.dispatchAndJoin();
+    }
+
+    @Override
+    public int dispatchDepth() {
+        return delegate.dispatchDepth();
+    }
+
+    @Override
+    public Object getCacheKey(K key) {
+        return delegate.getCacheKey(key);
+    }
+
+    @Override
+    public Statistics getStatistics() {
+        return delegate.getStatistics();
+    }
+
+    @Override
+    public CacheMap<Object, V> getCacheMap() {
+        return delegate.getCacheMap();
+    }
+
+    @Override
+    public ValueCache<K, V> getValueCache() {
+        return delegate.getValueCache();
+    }
+
+    @Override
+    public DataLoader<K, V> clear(K key) {
+        delegate.clear(key);
+        return this;
+    }
+
+    @Override
+    public DataLoader<K, V> clear(K key, BiConsumer<Void, Throwable> handler) {
+        delegate.clear(key, handler);
+        return this;
+    }
+
+    @Override
+    public DataLoader<K, V> clearAll() {
+        delegate.clearAll();
+        return this;
+    }
+
+    @Override
+    public DataLoader<K, V> clearAll(BiConsumer<Void, Throwable> handler) {
+        delegate.clearAll(handler);
+        return this;
+    }
+
+    @Override
+    public DataLoader<K, V> prime(K key, V value) {
+        delegate.prime(key, value);
+        return this;
+    }
+
+    @Override
+    public DataLoader<K, V> prime(K key, Exception error) {
+        delegate.prime(key, error);
+        return this;
+    }
+
+    @Override
+    public DataLoader<K, V> prime(K key, CompletableFuture<V> value) {
+        delegate.prime(key, value);
+        return this;
+    }
+}

src/test/java/org/dataloader/DataLoaderTest.java

@@ -785,7 +785,7 @@ public void should_work_with_duplicate_keys_when_caching_disabled(TestDataLoader
         assertThat(future1.get(), equalTo("A"));
         assertThat(future2.get(), equalTo("B"));
         assertThat(future3.get(), equalTo("A"));
-        if (factory instanceof MappedDataLoaderFactory || factory instanceof MappedPublisherDataLoaderFactory) {
+        if (factory.unwrap() instanceof MappedDataLoaderFactory || factory.unwrap() instanceof MappedPublisherDataLoaderFactory) {
             assertThat(loadCalls, equalTo(singletonList(asList("A", "B"))));
         } else {
             assertThat(loadCalls, equalTo(singletonList(asList("A", "B", "A"))));
@@ -1152,12 +1152,12 @@ public void when_values_size_are_less_then_key_size(TestDataLoaderFactory factor
 
         await().atMost(Duration.FIVE_SECONDS).until(() -> areAllDone(cf1, cf2, cf3, cf4));
 
-        if (factory instanceof ListDataLoaderFactory) {
+        if (factory.unwrap() instanceof ListDataLoaderFactory) {
             assertThat(cause(cf1), instanceOf(DataLoaderAssertionException.class));
             assertThat(cause(cf2), instanceOf(DataLoaderAssertionException.class));
             assertThat(cause(cf3), instanceOf(DataLoaderAssertionException.class));
             assertThat(cause(cf4), instanceOf(DataLoaderAssertionException.class));
-        } else if (factory instanceof PublisherDataLoaderFactory) {
+        } else if (factory.unwrap() instanceof PublisherDataLoaderFactory) {
             // some have completed progressively but the other never did
             assertThat(cf1.join(), equalTo("A"));
             assertThat(cf2.join(), equalTo("B"));
@@ -1187,7 +1187,7 @@ public void when_values_size_are_more_then_key_size(TestDataLoaderFactory factor
         await().atMost(Duration.FIVE_SECONDS).until(() -> areAllDone(cf1, cf2, cf3, cf4));
 
 
-        if (factory instanceof ListDataLoaderFactory) {
+        if (factory.unwrap() instanceof ListDataLoaderFactory) {
             assertThat(cause(cf1), instanceOf(DataLoaderAssertionException.class));
             assertThat(cause(cf2), instanceOf(DataLoaderAssertionException.class));
             assertThat(cause(cf3), instanceOf(DataLoaderAssertionException.class));

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

@@ -0,0 +1,64 @@
+package org.dataloader;
+
+import org.dataloader.fixtures.TestKit;
+import org.dataloader.fixtures.parameterized.DelegatingDataLoaderFactory;
+import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.Nullable;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.awaitility.Awaitility.await;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * There are WAY more tests via the {@link DelegatingDataLoaderFactory}
+ * parameterized tests.  All the basic {@link DataLoader} tests pass when wrapped in a {@link DelegatingDataLoader}
+ */
+public class DelegatingDataLoaderTest {
+
+    @Test
+    void canUnwrapDataLoaders() {
+        DataLoader<Object, Object> rawLoader = TestKit.idLoader();
+        DataLoader<Object, Object> delegateLoader = new DelegatingDataLoader<>(rawLoader);
+
+        assertThat(DelegatingDataLoader.unwrap(rawLoader), is(rawLoader));
+        assertThat(DelegatingDataLoader.unwrap(delegateLoader), is(rawLoader));
+    }
+
+    @Test
+    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);
+                return cf.thenApply(v -> "|" + v + "|");
+            }
+        };
+
+        assertThat(delegatingDataLoader.getDelegate(), is(rawLoader));
+
+
+        CompletableFuture<String> cfA = delegatingDataLoader.load("A");
+        CompletableFuture<String> cfB = delegatingDataLoader.load("B");
+        CompletableFuture<List<String>> cfCD = delegatingDataLoader.loadMany(List.of("C", "D"));
+
+        CompletableFuture<List<String>> dispatch = delegatingDataLoader.dispatch();
+
+        await().until(dispatch::isDone);
+
+        assertThat(cfA.join(), equalTo("|A|"));
+        assertThat(cfB.join(), equalTo("|B|"));
+        assertThat(cfCD.join(), equalTo(List.of("|C|", "|D|")));
+
+        assertThat(delegatingDataLoader.getIfPresent("A").isEmpty(), equalTo(false));
+        assertThat(delegatingDataLoader.getIfPresent("X").isEmpty(), equalTo(true));
+
+        assertThat(delegatingDataLoader.getIfCompleted("A").isEmpty(), equalTo(false));
+        assertThat(delegatingDataLoader.getIfCompleted("X").isEmpty(), equalTo(true));
+    }
+}

src/test/java/org/dataloader/fixtures/parameterized/DelegatingDataLoaderFactory.java

@@ -0,0 +1,71 @@
+package org.dataloader.fixtures.parameterized;
+
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderOptions;
+import org.dataloader.DelegatingDataLoader;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+public class DelegatingDataLoaderFactory implements TestDataLoaderFactory {
+    // its delegates all the way down to the turtles
+    private final TestDataLoaderFactory delegateFactory;
+
+    public DelegatingDataLoaderFactory(TestDataLoaderFactory delegateFactory) {
+        this.delegateFactory = delegateFactory;
+    }
+
+    @Override
+    public String toString() {
+        return "DelegatingDataLoaderFactory{" +
+                "delegateFactory=" + delegateFactory +
+                '}';
+    }
+
+    @Override
+    public TestDataLoaderFactory unwrap() {
+        return delegateFactory.unwrap();
+    }
+
+    private <K, V> DataLoader<K, V> mkDelegateDataLoader(DataLoader<K, V> dataLoader) {
+        return new DelegatingDataLoader<>(dataLoader);
+    }
+
+    @Override
+    public <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<K>> loadCalls) {
+        return mkDelegateDataLoader(delegateFactory.idLoader(options, loadCalls));
+    }
+
+    @Override
+    public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderDelayed(options, loadCalls, delay));
+    }
+
+    @Override
+    public <K> DataLoader<K, K> idLoaderBlowsUps(
+            DataLoaderOptions options, List<Collection<K>> loadCalls) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderBlowsUps(options, loadCalls));
+    }
+
+    @Override
+    public <K> DataLoader<K, Object> idLoaderAllExceptions(DataLoaderOptions options, List<Collection<K>> loadCalls) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderAllExceptions(options, loadCalls));
+    }
+
+    @Override
+    public DataLoader<Integer, Object> idLoaderOddEvenExceptions(DataLoaderOptions options, List<Collection<Integer>> loadCalls) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderOddEvenExceptions(options, loadCalls));
+    }
+
+    @Override
+    public DataLoader<String, String> onlyReturnsNValues(int N, DataLoaderOptions options, ArrayList<Object> loadCalls) {
+        return mkDelegateDataLoader(delegateFactory.onlyReturnsNValues(N, options, loadCalls));
+    }
+
+    @Override
+    public DataLoader<String, String> idLoaderReturnsTooMany(int howManyMore, DataLoaderOptions options, ArrayList<Object> loadCalls) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderReturnsTooMany(howManyMore, options, loadCalls));
+    }
+}