upgrade to dataloader 6

andimarek · andimarek · commit b3b3bea5f4e1 · 2025-11-10T08:51:29.000+10:00
diff --git a/src/main/java/graphql/schema/DataLoaderWithContext.java b/src/main/java/graphql/schema/DataLoaderWithContext.java
@@ -1,8 +1,6 @@
 package graphql.schema;
 
-import com.google.common.collect.Maps;
 import graphql.Internal;
-import graphql.execution.Async;
 import graphql.execution.incremental.AlternativeCallContext;
 import graphql.execution.instrumentation.dataloader.ExhaustedDataLoaderDispatchStrategy;
 import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategy;
@@ -12,14 +10,10 @@
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
-import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
-import static graphql.Assert.assertNotNull;
-
 @Internal
 @NullMarked
 public class DataLoaderWithContext<K, V> extends DelegatingDataLoader<K, V> {
@@ -32,10 +26,18 @@ public DataLoaderWithContext(DataFetchingEnvironment dfe, String dataLoaderName,
         this.dfe = dfe;
     }
 
+    // general note: calling super.load() is important, because otherwise the data loader will sometimes called
+    // later than the dispatch, which results in a hanging DL
+
+    @Override
+    public CompletableFuture<V> load(K key) {
+        CompletableFuture<V> result = super.load(key);
+        newDataLoaderInvocation();
+        return result;
+    }
+
     @Override
     public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
-        // calling super.load() is important, because otherwise the data loader will sometimes called
-        // later than the dispatch, which results in a hanging DL
         CompletableFuture<V> result = super.load(key, keyContext);
         newDataLoaderInvocation();
         return result;
@@ -44,41 +46,20 @@ public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
 
     @Override
     public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContexts) {
-        assertNotNull(keys);
-        assertNotNull(keyContexts);
-
-        CompletableFuture<List<V>> result;
-        List<CompletableFuture<V>> collect = new ArrayList<>(keys.size());
-        for (int i = 0; i < keys.size(); i++) {
-            K key = keys.get(i);
-            Object keyContext = null;
-            if (i < keyContexts.size()) {
-                keyContext = keyContexts.get(i);
-            }
-            collect.add(delegate.load(key, keyContext));
-        }
-        result = Async.allOf(collect);
+        CompletableFuture<List<V>> result = super.loadMany(keys, keyContexts);
         newDataLoaderInvocation();
         return result;
     }
 
     @Override
     public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
-        assertNotNull(keysAndContexts);
-
-        CompletableFuture<Map<K, V>> result;
-        Map<K, CompletableFuture<V>> collect = Maps.newHashMapWithExpectedSize(keysAndContexts.size());
-        for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
-            K key = entry.getKey();
-            Object keyContext = entry.getValue();
-            collect.put(key, delegate.load(key, keyContext));
-        }
-        result = Async.allOf(collect);
+        CompletableFuture<Map<K, V>> result = super.loadMany(keysAndContexts);
         newDataLoaderInvocation();
         return result;
     }
 
     private void newDataLoaderInvocation() {
+        System.out.println("newDataLoaderInvocation");
         DataFetchingEnvironmentImpl dfeImpl = (DataFetchingEnvironmentImpl) dfe;
         DataFetchingEnvironmentImpl.DFEInternalState dfeInternalState = (DataFetchingEnvironmentImpl.DFEInternalState) dfeImpl.toInternal();
         if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof PerLevelDataLoaderDispatchStrategy) {
