Added the ability to have call context on batch loader calls

Author: bbakerman

README.md

@@ -165,6 +165,34 @@ a list of user ids in one call.
 
 That said, with key caching turn on (the default), it will still be more efficient using `dataloader` than without it.
 
+### Calling the batch loader function with context
+
+Often there is a need to call the batch loader function with some sort of context, such as the calling users security
+credentials or the database connection parameters.  You can do this by implementing a 
+`org.dataloader.BatchContextProvider`.
+
+```java
+        BatchLoader<String, String> batchLoader = new BatchLoader<String, String>() {
+            @Override
+            public CompletionStage<List<String>> load(List<String> keys) {
+                throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method");
+            }
+
+            @Override
+            public CompletionStage<List<String>> load(List<String> keys, Object context) {
+                SecurityCtx callCtx = (SecurityCtx) context;
+                return callDatabaseForResults(callCtx, keys);
+            }
+
+        };
+        DataLoaderOptions options = DataLoaderOptions.newOptions()
+                .setBatchContextProvider(() -> SecurityCtx.getCallingUserCtx());
+        DataLoader<String, String> loader = new DataLoader<>(batchLoader, options);
+
+```
+
+The batch loading code will now receive this context object and it can be used to get to data layers or
+to connect to other systems. 
 
 ### Error object is not a thing in a type safe Java world
 

src/main/java/org/dataloader/BatchContextProvider.java

@@ -0,0 +1,13 @@
+package org.dataloader;
+
+/**
+ * A BatchContextProvider is used by the {@link org.dataloader.DataLoader} code to
+ * provide context to the {@link org.dataloader.BatchLoader} call.  A common use
+ * case is for propagating user security credentials or database connection parameters.
+ */
+public interface BatchContextProvider {
+    /**
+     * @return a context object that may be needed in batch calls
+     */
+    Object get();
+}

src/main/java/org/dataloader/BatchLoader.java

@@ -74,11 +74,33 @@
 public interface BatchLoader<K, V> {
 
     /**
-     * Called to batch load the provided keys and return a promise to a list of values
+     * Called to batch load the provided keys and return a promise to a list of values.
+     *
+     * If you need calling context then implement the {@link #load(java.util.List, Object)} method
+     * instead.
      *
      * @param keys the collection of keys to load
      *
      * @return a promise of the values for those keys
      */
     CompletionStage<List<V>> load(List<K> keys);
+
+    /**
+     * Called to batch load the provided keys and return a promise to a list of values.  This default
+     * version can be given a context object to that maybe be useful during the call.  A typical use case
+     * is passing in security credentials or database connection details say.
+     *
+     * This method is implemented as a default method in order to preserve the API for previous
+     * callers.  It is always called first by the {@link org.dataloader.DataLoader} code and simply
+     * delegates to the {@link #load(java.util.List)} method.
+     *
+     * @param keys    the collection of keys to load
+     * @param context a context object that can help with the call
+     *
+     * @return a promise of the values for those keys
+     */
+    @SuppressWarnings("unused")
+    default CompletionStage<List<V>> load(List<K> keys, Object context) {
+        return load(keys);
+    }
 }

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

@@ -196,8 +196,9 @@ public CompletableFuture<V> load(K key) {
             } else {
                 stats.incrementBatchLoadCountBy(1);
                 // immediate execution of batch function
+                Object context = loaderOptions.getBatchContextProvider().get();
                 CompletableFuture<List<V>> batchedLoad = batchLoadFunction
-                        .load(singletonList(key))
+                        .load(singletonList(key), context)
                         .toCompletableFuture();
                 future = batchedLoad
                         .thenApply(list -> list.get(0));
@@ -303,7 +304,8 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Complet
         stats.incrementBatchLoadCountBy(keys.size());
         CompletionStage<List<V>> batchLoad;
         try {
-            batchLoad = nonNull(batchLoadFunction.load(keys), "Your batch loader function MUST return a non null CompletionStage promise");
+            Object context = loaderOptions.getBatchContextProvider().get();
+            batchLoad = nonNull(batchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise");
         } catch (Exception e) {
             batchLoad = CompletableFutureKit.failedFuture(e);
         }

src/main/java/org/dataloader/DataLoaderOptions.java

@@ -31,12 +31,15 @@
  */
 public class DataLoaderOptions {
 
+    private static final BatchContextProvider NULL_PROVIDER = () -> null;
+
     private boolean batchingEnabled;
     private boolean cachingEnabled;
     private CacheKey cacheKeyFunction;
     private CacheMap cacheMap;
     private int maxBatchSize;
     private Supplier<StatisticsCollector> statisticsCollector;
+    private BatchContextProvider contextProvider;
 
     /**
      * Creates a new data loader options with default settings.
@@ -46,6 +49,7 @@ public DataLoaderOptions() {
         cachingEnabled = true;
         maxBatchSize = -1;
         statisticsCollector = SimpleStatisticsCollector::new;
+        contextProvider = NULL_PROVIDER;
     }
 
     /**
@@ -61,6 +65,7 @@ public DataLoaderOptions(DataLoaderOptions other) {
         this.cacheMap = other.cacheMap;
         this.maxBatchSize = other.maxBatchSize;
         this.statisticsCollector = other.statisticsCollector;
+        this.contextProvider = other.contextProvider;
     }
 
     /**
@@ -202,5 +207,22 @@ public DataLoaderOptions setStatisticsCollector(Supplier<StatisticsCollector> st
         return this;
     }
 
+    /**
+     * @return the batch context provider that will be used to give context to batch load functions
+     */
+    public BatchContextProvider getBatchContextProvider() {
+        return contextProvider;
+    }
 
+    /**
+     * Sets the batch context provider that will be used to give context to batch load functions
+     *
+     * @param contextProvider the batch context provider
+     *
+     * @return the data loader options for fluent coding
+     */
+    public DataLoaderOptions setBatchContextProvider(BatchContextProvider contextProvider) {
+        this.contextProvider = nonNull(contextProvider);
+        return this;
+    }
 }

src/test/java/ReadmeExamples.java

@@ -78,6 +78,36 @@ public CompletionStage<List<User>> load(List<Long> userIds) {
         userLoader.dispatchAndJoin();
     }
 
+    private static class SecurityCtx {
+
+        public static Object getCallingUserCtx() {
+            return null;
+        }
+    }
+
+    private void callContextExample() {
+        BatchLoader<String, String> batchLoader = new BatchLoader<String, String>() {
+            @Override
+            public CompletionStage<List<String>> load(List<String> keys) {
+                throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method");
+            }
+
+            @Override
+            public CompletionStage<List<String>> load(List<String> keys, Object context) {
+                SecurityCtx callCtx = (SecurityCtx) context;
+                return callDatabaseForResults(callCtx, keys);
+            }
+
+        };
+        DataLoaderOptions options = DataLoaderOptions.newOptions()
+                .setBatchContextProvider(() -> SecurityCtx.getCallingUserCtx());
+        DataLoader<String, String> loader = new DataLoader<>(batchLoader, options);
+    }
+
+    private CompletionStage<List<String>> callDatabaseForResults(SecurityCtx callCtx, List<String> keys) {
+        return null;
+    }
+
 
     private void tryExample() {
         Try<String> tryS = Try.tryCall(() -> {

src/test/java/org/dataloader/DataLoaderContextTest.java

@@ -0,0 +1,59 @@
+package org.dataloader;
+
+import org.junit.Test;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.stream.Collectors;
+
+import static java.util.Arrays.asList;
+import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertThat;
+
+/**
+ * Tests related to context.  DataLoaderTest is getting to big and needs refactoring
+ */
+public class DataLoaderContextTest {
+
+    @Test
+    public void context_is_passed_to_batch_loader_function() throws Exception {
+        BatchLoader<String, String> batchLoader = new BatchLoader<String, String>() {
+            @Override
+            public CompletionStage<List<String>> load(List<String> keys) {
+                throw new UnsupportedOperationException("this wont be called");
+            }
+
+            @Override
+            public CompletionStage<List<String>> load(List<String> keys, Object context) {
+                List<String> list = keys.stream().map(k -> k + "-" + context).collect(Collectors.toList());
+                return CompletableFuture.completedFuture(list);
+            }
+        };
+        DataLoaderOptions options = DataLoaderOptions.newOptions()
+                .setBatchContextProvider(() -> "ctx");
+        DataLoader<String, String> loader = new DataLoader<>(batchLoader, options);
+
+        loader.load("A");
+        loader.load("B");
+        loader.loadMany(asList("C", "D"));
+
+        List<String> results = loader.dispatchAndJoin();
+
+        assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx")));
+    }
+
+    @Test
+    public void null_is_passed_as_context_if_you_do_nothing() throws Exception {
+        BatchLoader<String, String> batchLoader = CompletableFuture::completedFuture;
+        DataLoader<String, String> loader = new DataLoader<>(batchLoader);
+
+        loader.load("A");
+        loader.load("B");
+        loader.loadMany(asList("C", "D"));
+
+        List<String> results = loader.dispatchAndJoin();
+
+        assertThat(results, equalTo(asList("A", "B", "C", "D")));
+    }
+}