DataLoader and CompletableFutureKit: add support for loadMany given a Map of keys and context objects.

Author: Harrison Cole

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

@@ -25,10 +25,7 @@
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Optional;
+import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.BiConsumer;
 
@@ -574,6 +571,35 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
         }
     }
 
+    /**
+     * Requests to load the map of data provided by the specified keys asynchronously, and returns a composite future
+     * of the resulting values.
+     * <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).
+     * <p>
+     * The key context object may be useful in the batch loader interfaces such as {@link org.dataloader.BatchLoaderWithContext} or
+     * {@link org.dataloader.MappedBatchLoaderWithContext} to help retrieve data.
+     *
+     * @param keysAndContexts the map of keys to their respective contexts
+     *
+     * @return the composite future of the map of keys and values
+     */
+    public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
+        nonNull(keysAndContexts);
+
+        synchronized (this) {
+            Map<K, CompletableFuture<V>> collect = new HashMap<>(keysAndContexts.size());
+            for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
+                K key = entry.getKey();
+                Object keyContext = entry.getValue();
+                collect.put(key, load(key, keyContext));
+            }
+            return CompletableFutureKit.allOf(collect);
+        }
+    }
+
     /**
      * Dispatches the queued load requests to the batch execution function and returns a promise of the result.
      * <p>

src/main/java/org/dataloader/impl/CompletableFutureKit.java

@@ -3,8 +3,10 @@
 import org.dataloader.annotations.Internal;
 
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
 
 import static java.util.stream.Collectors.toList;
 
@@ -48,10 +50,21 @@ public static <V> boolean failed(CompletableFuture<V> future) {
     }
 
     public static <T> CompletableFuture<List<T>> allOf(List<CompletableFuture<T>> cfs) {
-        return CompletableFuture.allOf(cfs.toArray(new CompletableFuture[0]))
+        return CompletableFuture.allOf(cfs.toArray(CompletableFuture[]::new))
                 .thenApply(v -> cfs.stream()
                         .map(CompletableFuture::join)
                         .collect(toList())
                 );
     }
+
+    public static <K, V> CompletableFuture<Map<K, V>> allOf(Map<K, CompletableFuture<V>> cfs) {
+        return CompletableFuture.allOf(cfs.values().toArray(CompletableFuture[]::new))
+                .thenApply(v -> cfs.entrySet().stream()
+                        .collect(
+                                Collectors.toMap(
+                                        Map.Entry::getKey,
+                                        task -> task.getValue().join())
+                        )
+                );
+    }
 }