Merge pull request #166 from Hc747/master

Add support for batch loading of a Map of key-value pairs.

Author: bbakerman

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())
+ )
+ );
+ }
 }

src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java

@@ -2,10 +2,7 @@
 
 import org.junit.jupiter.api.Test;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
@@ -50,10 +47,14 @@ public void context_is_passed_to_batch_loader_function() {
  loader.load("A");
  loader.load("B");
  loader.loadMany(asList("C", "D"));
+ Map<String, ?> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", null);
+ keysAndContexts.put("F", null);
+ loader.loadMany(keysAndContexts);
 
  List<String> results = loader.dispatchAndJoin();
 
- assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx")));
+ assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx", "E-ctx", "F-ctx")));
  }
 
  @Test
@@ -66,10 +67,14 @@ public void key_contexts_are_passed_to_batch_loader_function() {
  loader.load("A", "aCtx");
  loader.load("B", "bCtx");
  loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx"));
+ Map<String, String> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", "eCtx");
+ keysAndContexts.put("F", "fCtx");
+ loader.loadMany(keysAndContexts);
 
  List<String> results = loader.dispatchAndJoin();
 
- assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx")));
+ assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:fCtx-l:fCtx")));
  }
 
  @Test
@@ -82,12 +87,17 @@ public void key_contexts_are_passed_to_batch_loader_function_when_batching_disab
 
  CompletableFuture<String> aLoad = loader.load("A", "aCtx");
  CompletableFuture<String> bLoad = loader.load("B", "bCtx");
- CompletableFuture<List<String>> canDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx"));
+ CompletableFuture<List<String>> cAndDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx"));
+ Map<String, String> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", "eCtx");
+ keysAndContexts.put("F", "fCtx");
+ CompletableFuture<Map<String, String>> eAndFLoad = loader.loadMany(keysAndContexts);
 
  List<String> results = new ArrayList<>(asList(aLoad.join(), bLoad.join()));
- results.addAll(canDLoad.join());
+ results.addAll(cAndDLoad.join());
+ results.addAll(eAndFLoad.join().values());
 
- assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx")));
+ assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:fCtx-l:fCtx")));
  }
 
  @Test
@@ -101,9 +111,14 @@ public void missing_key_contexts_are_passed_to_batch_loader_function() {
  loader.load("B");
  loader.loadMany(asList("C", "D"), singletonList("cCtx"));
 
+ Map<String, String> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", "eCtx");
+ keysAndContexts.put("F", null);
+ loader.loadMany(keysAndContexts);
+
  List<String> results = loader.dispatchAndJoin();
 
- assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null")));
+ assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:null-l:null")));
  }
 
  @Test
@@ -125,9 +140,14 @@ public void context_is_passed_to_map_batch_loader_function() {
  loader.load("B");
  loader.loadMany(asList("C", "D"), singletonList("cCtx"));
 
+ Map<String, String> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", "eCtx");
+ keysAndContexts.put("F", null);
+ loader.loadMany(keysAndContexts);
+
  List<String> results = loader.dispatchAndJoin();
 
- assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null")));
+ assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null", "E-ctx-eCtx", "F-ctx-null")));
  }
 
  @Test
@@ -142,9 +162,14 @@ public void null_is_passed_as_context_if_you_do_nothing() {
  loader.load("B");
  loader.loadMany(asList("C", "D"));
 
+ Map<String, String> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", null);
+ keysAndContexts.put("F", null);
+ loader.loadMany(keysAndContexts);
+
  List<String> results = loader.dispatchAndJoin();
 
- assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null")));
+ assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null", "E-null", "F-null")));
  }
 
  @Test
@@ -160,9 +185,14 @@ public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() {
  loader.load("B");
  loader.loadMany(asList("C", "D"));
 
+ Map<String, String> keysAndContexts = new LinkedHashMap<>();
+ keysAndContexts.put("E", null);
+ keysAndContexts.put("F", null);
+ loader.loadMany(keysAndContexts);
+
  List<String> results = loader.dispatchAndJoin();
 
- assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null")));
+ assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null", "E-null", "F-null")));
  }
 
  @Test

src/test/java/org/dataloader/DataLoaderStatsTest.java

@@ -13,6 +13,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
 import static java.util.Arrays.asList;
@@ -118,19 +119,20 @@ public void stats_are_collected_with_caching_disabled() {
  loader.load("A");
  loader.load("B");
  loader.loadMany(asList("C", "D"));
+ loader.loadMany(Map.of("E", "E", "F", "F"));
 
  Statistics stats = loader.getStatistics();
- assertThat(stats.getLoadCount(), equalTo(4L));
+ assertThat(stats.getLoadCount(), equalTo(6L));
  assertThat(stats.getBatchInvokeCount(), equalTo(0L));
  assertThat(stats.getBatchLoadCount(), equalTo(0L));
  assertThat(stats.getCacheHitCount(), equalTo(0L));
 
  loader.dispatch();
 
  stats = loader.getStatistics();
- assertThat(stats.getLoadCount(), equalTo(4L));
+ assertThat(stats.getLoadCount(), equalTo(6L));
  assertThat(stats.getBatchInvokeCount(), equalTo(1L));
- assertThat(stats.getBatchLoadCount(), equalTo(4L));
+ assertThat(stats.getBatchLoadCount(), equalTo(6L));
  assertThat(stats.getCacheHitCount(), equalTo(0L));
 
  loader.load("A");
@@ -139,9 +141,9 @@ public void stats_are_collected_with_caching_disabled() {
  loader.dispatch();
 
  stats = loader.getStatistics();
- assertThat(stats.getLoadCount(), equalTo(6L));
+ assertThat(stats.getLoadCount(), equalTo(8L));
  assertThat(stats.getBatchInvokeCount(), equalTo(2L));
- assertThat(stats.getBatchLoadCount(), equalTo(6L));
+ assertThat(stats.getBatchLoadCount(), equalTo(8L));
  assertThat(stats.getCacheHitCount(), equalTo(0L));
  }