Added the BatchLoadingEnvironment as a way to hang future objects into the call context

Author: bbakerman

README.md

@@ -165,33 +165,40 @@ 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
+### Calling the batch loader function with call context environment
 
-Often there is a need to call the batch loader function with some sort of context, such as the calling users security
+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.BatchContextProvider`.
+`org.dataloader.BatchLoaderEnvironmentProvider`.
 
 ```java
         BatchLoader<String, String> batchLoader = new BatchLoader<String, String>() {
+            //
+            // the reason this method exists is for backwards compatibility.  There are a large
+            // number of existing dataloader clients out there that used this method before the call context was invented
+            // and hence to preserve compatibility we have this unfortunate method declaration.
+            //
             @Override
             public CompletionStage<List<String>> load(List<String> keys) {
-                throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method");
+                throw new UnsupportedOperationException();
             }
 
             @Override
-            public CompletionStage<List<String>> load(List<String> keys, Object context) {
-                SecurityCtx callCtx = (SecurityCtx) context;
+            public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironment environment) {
+                SecurityCtx callCtx = environment.getContext();
                 return callDatabaseForResults(callCtx, keys);
             }
-
         };
+
+        BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment()
+                .context(SecurityCtx.getCallingUserCtx()).build();
+
         DataLoaderOptions options = DataLoaderOptions.newOptions()
-                .setBatchContextProvider(() -> SecurityCtx.getCallingUserCtx());
+                .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment);
         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
+The batch loading code will now receive this environment object and it can be used to get context perhaps allowing it
 to connect to other systems. 
 
 ### Returning a Map of results from your batch loader
@@ -219,8 +226,8 @@ For example, let's assume you want to load users from a database, you could prob
 ```java
         MapBatchLoader<Long, User> mapBatchLoader = new MapBatchLoader<Long, User>() {
             @Override
-            public CompletionStage<Map<Long, User>> load(List<Long> userIds, Object context) {
-                SecurityCtx callCtx = (SecurityCtx) context;
+            public CompletionStage<Map<Long, User>> load(List<Long> userIds, BatchLoaderEnvironment environment) {
+                SecurityCtx callCtx = environment.getContext();
                 return CompletableFuture.supplyAsync(() -> userManager.loadMapOfUsersById(callCtx, userIds));
             }
         };
@@ -267,7 +274,7 @@ and some of which may have failed.  From that data loader can infer the right be
 ```java
         DataLoader<String, User> dataLoader = DataLoader.newDataLoaderWithTry(new BatchLoader<String, Try<User>>() {
             @Override
-            public CompletionStage<List<Try<User>>> load(List<String> keys) {
+            public CompletionStage<List<Try<User>>> load(List<String> keys, BatchLoaderEnvironment environment) {
                 return CompletableFuture.supplyAsync(() -> {
                     List<Try<User>> users = new ArrayList<>();
                     for (String key : keys) {
@@ -278,7 +285,6 @@ and some of which may have failed.  From that data loader can infer the right be
                 });
             }
         });
-
 ```
 
 On the above example if one of the `Try` objects represents a failure, then its `load()` promise will complete exceptionally and you can 

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

@@ -16,7 +16,9 @@
 
 package org.dataloader;
 
+import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
 /**
@@ -76,31 +78,32 @@ public interface BatchLoader<K, V> {
     /**
      * 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.
+     * If you need calling context then implement the {@link #load(java.util.List, BatchLoaderEnvironment)} method
+     * instead which is the preferred method.
      *
      * @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
+     * version can be given an environment 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
+     * @param keys        the collection of keys to load
+     * @param environment an environment 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) {
+    @SuppressWarnings({"unused", "deprecation"})
+    default CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
         return load(keys);
     }
 }

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

@@ -0,0 +1,40 @@
+package org.dataloader;
+
+/**
+ * This object is passed to a batch loader as calling context.  It could contain security credentials
+ * of the calling users say or database connection parameters that allow the data layer call to succeed.
+ */
+public class BatchLoaderEnvironment {
+
+    private final Object context;
+
+    private BatchLoaderEnvironment(Object context) {
+        this.context = context;
+    }
+
+    @SuppressWarnings("unchecked")
+    public <T> T getContext() {
+        return (T) context;
+    }
+
+    public static Builder newBatchLoaderEnvironment() {
+        return new Builder();
+    }
+
+    public static class Builder {
+        private Object context;
+
+        private Builder() {
+
+        }
+
+        public Builder context(Object context) {
+            this.context = context;
+            return this;
+        }
+
+        public BatchLoaderEnvironment build() {
+            return new BatchLoaderEnvironment(context);
+        }
+    }
+}

src/main/java/org/dataloader/BatchLoaderEnvironment.java

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

src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java

@@ -298,8 +298,7 @@ public CompletableFuture<V> load(K key) {
             } else {
                 stats.incrementBatchLoadCountBy(1);
                 // immediate execution of batch function
-                Object context = loaderOptions.getBatchContextProvider().get();
-                future = invokeLoaderImmediately(key, context);
+                future = invokeLoaderImmediately(key);
             }
             if (cachingEnabled) {
                 futureCache.set(cacheKey, future);
@@ -331,25 +330,6 @@ public CompletableFuture<List<V>> loadMany(List<K> keys) {
         }
     }
 
-    private CompletableFuture<V> invokeLoaderImmediately(K key, Object context) {
-        List<K> keys = singletonList(key);
-        CompletionStage<V> singleLoadCall;
-        if (isMapLoader()) {
-            singleLoadCall = mapBatchLoadFunction
-                    .load(keys, context)
-                    .thenApply(map -> map.get(key));
-        } else {
-            singleLoadCall = batchLoadFunction
-                    .load(keys, context)
-                    .thenApply(list -> list.get(0));
-        }
-        return singleLoadCall.toCompletableFuture();
-    }
-
-    private boolean isMapLoader() {
-        return mapBatchLoadFunction != null;
-    }
-
     /**
      * Dispatches the queued load requests to the batch execution function and returns a promise of the result.
      * <p>
@@ -420,7 +400,7 @@ private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<
     @SuppressWarnings("unchecked")
     private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<CompletableFuture<V>> queuedFutures) {
         stats.incrementBatchLoadCountBy(keys.size());
-        CompletionStage<List<V>> batchLoad = invokeBatchFunction(keys);
+        CompletionStage<List<V>> batchLoad = invokeLoader(keys);
         return batchLoad
                 .toCompletableFuture()
                 .thenApply(values -> {
@@ -463,31 +443,51 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Complet
                 });
     }
 
-    private CompletionStage<List<V>> invokeBatchFunction(List<K> keys) {
+    private boolean isMapLoader() {
+        return mapBatchLoadFunction != null;
+    }
+
+    private CompletableFuture<V> invokeLoaderImmediately(K key) {
+        BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get();
+        List<K> keys = singletonList(key);
+        CompletionStage<V> singleLoadCall;
+        if (isMapLoader()) {
+            singleLoadCall = mapBatchLoadFunction
+                    .load(keys, environment)
+                    .thenApply(map -> map.get(key));
+        } else {
+            singleLoadCall = batchLoadFunction
+                    .load(keys, environment)
+                    .thenApply(list -> list.get(0));
+        }
+        return singleLoadCall.toCompletableFuture();
+    }
+
+    private CompletionStage<List<V>> invokeLoader(List<K> keys) {
         CompletionStage<List<V>> batchLoad;
         try {
-            Object context = loaderOptions.getBatchContextProvider().get();
+            BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get();
             if (isMapLoader()) {
-                batchLoad = invokeMapBatchLoader(keys, context);
+                batchLoad = invokeMapBatchLoader(keys, environment);
             } else {
-                batchLoad = invokeListBatchLoader(keys, context);
+                batchLoad = invokeListBatchLoader(keys, environment);
             }
         } catch (Exception e) {
             batchLoad = CompletableFutureKit.failedFuture(e);
         }
         return batchLoad;
     }
 
-    private CompletionStage<List<V>> invokeListBatchLoader(List<K> keys, Object context) {
-        return nonNull(batchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise");
+    private CompletionStage<List<V>> invokeListBatchLoader(List<K> keys, BatchLoaderEnvironment environment) {
+        return nonNull(batchLoadFunction.load(keys, environment), "Your batch loader function MUST return a non null CompletionStage promise");
     }
 
     /*
      * Turns a map of results that MAY be smaller than the key list back into a list by mapping null
      * to missing elements.
      */
-    private CompletionStage<List<V>> invokeMapBatchLoader(List<K> keys, Object context) {
-        CompletionStage<Map<K, V>> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise");
+    private CompletionStage<List<V>> invokeMapBatchLoader(List<K> keys, BatchLoaderEnvironment environment) {
+        CompletionStage<Map<K, V>> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, environment), "Your batch loader function MUST return a non null CompletionStage promise");
         return mapBatchLoad.thenApply(map -> {
             List<V> values = new ArrayList<>();
             for (K key : keys) {

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

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

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

@@ -57,21 +57,19 @@
  * @param <K> type parameter indicating the type of keys to use for data load requests.
  * @param <V> type parameter indicating the type of values returned
  *
- * @author <a href="https://github.com/aschrijver/">Arnold Schrijver</a>
  * @author <a href="https://github.com/bbakerman/">Brad Baker</a>
  */
-@FunctionalInterface
 public interface MapBatchLoader<K, V> {
 
     /**
-     * Called to batch load the provided keys and return a promise to a map of values.  It can be given a context object to
+     * Called to batch load the provided keys and return a promise to a map of values.  It can be given an environment object to
      * that maybe be useful during the call.  A typical use case is passing in security credentials or database connection details say.
      *
-     * @param keys    the collection of keys to load
-     * @param context a context object that can help with the call
+     * @param keys        the collection of keys to load
+     * @param environment an environment object that can help with the call
      *
      * @return a promise to a map of values for those keys
      */
     @SuppressWarnings("unused")
-    CompletionStage<Map<K, V>> load(List<K> keys, Object context);
+    CompletionStage<Map<K, V>> load(List<K> keys, BatchLoaderEnvironment environment);
 }