@@ -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; Expand All @@ -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> { Expand All @@ -32,48 +26,40 @@ 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; }
@Override public CompletableFuture<List<V>> loadMany(List<K> keys) { CompletableFuture<List<V>> result = super.loadMany(keys); newDataLoaderInvocation(); return result; }
@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; } Expand Down