graphql-java · bbakerman · Aug 27, 2018 · Aug 19, 2018 · Aug 19, 2018 · Aug 19, 2018
diff --git a/README.md b/README.md
@@ -168,16 +168,17 @@ That said, with key caching turn on (the default), it will still be more efficie
 ### Calling the batch loader function with call context environment
 
 Often there is a need to call the batch loader function with some sort of call context environment, such as the calling users security
-credentials or the database connection parameters.  You can do this by implementing a 
-`org.dataloader.BatchLoaderEnvironmentProvider` and using one of the `xxxWithContext` batch loading interfaces
-such as `org.dataloader.BatchLoaderWithContext`.
+credentials or the database connection parameters.  
 
-```java
-        BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment()
-                .context(SecurityCtx.getCallingUserCtx()).build();
+You can do this by implementing a `org.dataloader.BatchLoaderContextProvider` and using one of 
+the batch loading interfaces such as `org.dataloader.BatchLoaderWithContext`.
+
+It will be given a `org.dataloader.BatchLoaderEnvironment` parameter and it can then ask it
+for the context object.
 
+```java
         DataLoaderOptions options = DataLoaderOptions.newOptions()
-                .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment);
+                .setBatchLoaderEnvironmentProvider(() -> SecurityCtx.getCallingUserCtx());
 
         BatchLoaderWithContext<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
             @Override

diff --git a/src/main/java/org/dataloader/BatchLoaderContextProvider.java b/src/main/java/org/dataloader/BatchLoaderContextProvider.java
@@ -0,0 +1,13 @@
+package org.dataloader;
+
+/**
+ * A BatchLoaderContextProvider is used by the {@link org.dataloader.DataLoader} code to
+ * provide overall calling context to the {@link org.dataloader.BatchLoader} call.  A common use
+ * case is for propagating user security credentials or database connection parameters for example.
+ */
+public interface BatchLoaderContextProvider {
+    /**
+     * @return a context object that may be needed in batch load calls
+     */
+    Object getContext();
+}
diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironment.java b/src/main/java/org/dataloader/BatchLoaderEnvironment.java
@@ -1,28 +1,54 @@
 package org.dataloader;
 
+import java.util.Collections;
+import java.util.Map;
+
+import static org.dataloader.impl.Assertions.nonNull;
+
 /**
  * This object is passed to a batch loader as calling context.  It could contain security credentials
  * of the calling users for example or database parameters that allow the data layer call to succeed.
  */
 public class BatchLoaderEnvironment {
 
     private final Object context;
+    private final Map<Object, Object> keyContexts;
 
-    private BatchLoaderEnvironment(Object context) {
+    private BatchLoaderEnvironment(Object context, Map<Object, Object> keyContexts) {
         this.context = context;
+        this.keyContexts = keyContexts;
     }
 
+    /**
+     * Returns the overall context object provided by {@link org.dataloader.BatchLoaderContextProvider}
+     *
+     * @param <T> the type you would like the object to be
+     *
+     * @return a context object or null if there isn't one
+     */
     @SuppressWarnings("unchecked")
     public <T> T getContext() {
         return (T) context;
     }
 
+    /**
+     * Each call to {@link org.dataloader.DataLoader#load(Object, Object)} or
+     * {@link org.dataloader.DataLoader#loadMany(java.util.List, java.util.List)} can be given
+     * a context object when it is invoked.  A map of them is present by this method.
+     *
+     * @return a map of key context objects
+     */
+    public Map<Object, Object> getKeyContexts() {
+        return keyContexts;
+    }
+
     public static Builder newBatchLoaderEnvironment() {
         return new Builder();
     }
 
     public static class Builder {
         private Object context;
+        private Map<Object, Object> keyContexts = Collections.emptyMap();
 
         private Builder() {
 
@@ -33,8 +59,13 @@ public Builder context(Object context) {
             return this;
         }
 
+        public Builder keyContexts(Map<Object, Object> keyContexts) {
+            this.keyContexts = nonNull(keyContexts);
+            return this;
+        }
+
         public BatchLoaderEnvironment build() {
-            return new BatchLoaderEnvironment(context);
+            return new BatchLoaderEnvironment(context, keyContexts);
         }
     }
 }
diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java
@@ -20,11 +20,10 @@
 import org.dataloader.stats.Statistics;
 import org.dataloader.stats.StatisticsCollector;
 
-import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
-import java.util.stream.Collectors;
 
 import static org.dataloader.impl.Assertions.nonNull;
 
@@ -60,7 +59,6 @@ public class DataLoader<K, V> {
     private final DataLoaderHelper<K, V> helper;
     private final DataLoaderOptions loaderOptions;
     private final CacheMap<Object, CompletableFuture<V>> futureCache;
-    private final List<SimpleImmutableEntry<K, CompletableFuture<V>>> loaderQueue;
     private final StatisticsCollector stats;
 
     /**
@@ -237,6 +235,7 @@ public static <K, V> DataLoader<K, V> newMappedDataLoader(MappedBatchLoader<K, V
      * Using Try objects allows you to capture a value returned or an exception that might
      * have occurred trying to get a value. .
      * <p>
+     *
      * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects
      * @param <K>               the key type
      * @param <V>               the value type
@@ -357,10 +356,9 @@ private DataLoader(Object batchLoadFunction, DataLoaderOptions options) {
         this.loaderOptions = options == null ? new DataLoaderOptions() : options;
         this.futureCache = determineCacheMap(loaderOptions);
         // order of keys matter in data loader
-        this.loaderQueue = new ArrayList<>();
         this.stats = nonNull(this.loaderOptions.getStatisticsCollector());
 
-        this.helper = new DataLoaderHelper<>(this, batchLoadFunction, this.loaderOptions, this.futureCache, this.loaderQueue, this.stats);
+        this.helper = new DataLoaderHelper<>(this, batchLoadFunction, this.loaderOptions, this.futureCache, this.stats);
     }
 
     @SuppressWarnings("unchecked")
@@ -380,7 +378,26 @@ private CacheMap<Object, CompletableFuture<V>> determineCacheMap(DataLoaderOptio
      * @return the future of the value
      */
     public CompletableFuture<V> load(K key) {
-        return helper.load(key);
+        return load(key, null);
+    }
+
+    /**
+     * Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
+     * <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 key        the key to load
+     * @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);
     }
 
     /**
@@ -396,11 +413,39 @@ public CompletableFuture<V> load(K key) {
      * @return the composite future of the list of values
      */
     public CompletableFuture<List<V>> loadMany(List<K> keys) {
-        synchronized (this) {
-            List<CompletableFuture<V>> collect = keys.stream()
-                    .map(this::load)
-                    .collect(Collectors.toList());
+        return loadMany(keys, Collections.emptyList());
+    }
 
+    /**
+     * Requests to load the list 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 keys        the list of keys to load
+     * @param keyContexts the list of key calling context objects
+     *
+     * @return the composite future of the list of values
+     */
+    public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContexts) {
+        nonNull(keys);
+        nonNull(keyContexts);
+
+        synchronized (this) {
+            List<CompletableFuture<V>> collect = new ArrayList<>();
+            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(load(key, keyContext));
+            }
             return CompletableFutureKit.allOf(collect);
         }
     }
@@ -441,9 +486,7 @@ public List<V> dispatchAndJoin() {
      * @return the depth of the batched key loads that need to be dispatched
      */
     public int dispatchDepth() {
-        synchronized (this) {
-            return loaderQueue.size();
-        }
+        return helper.dispatchDepth();
     }