diff --git a/README.md b/README.md index 293d0ba..c8b2305 100644 --- a/README.md +++ b/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 batchLoader = new BatchLoader() { + @Override + public CompletionStage> load(List keys) { + throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method"); + } + + @Override + public CompletionStage> load(List keys, Object context) { + SecurityCtx callCtx = (SecurityCtx) context; + return callDatabaseForResults(callCtx, keys); + } + + }; + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchContextProvider(() -> SecurityCtx.getCallingUserCtx()); + DataLoader 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 diff --git a/src/main/java/org/dataloader/BatchContextProvider.java b/src/main/java/org/dataloader/BatchContextProvider.java new file mode 100644 index 0000000..0a43664 --- /dev/null +++ b/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(); +} \ No newline at end of file diff --git a/src/main/java/org/dataloader/BatchLoader.java b/src/main/java/org/dataloader/BatchLoader.java index ed51330..8791f14 100644 --- a/src/main/java/org/dataloader/BatchLoader.java +++ b/src/main/java/org/dataloader/BatchLoader.java @@ -74,11 +74,33 @@ public interface BatchLoader { /** - * 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> load(List 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> load(List keys, Object context) { + return load(keys); + } } diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index a942468..69b3178 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -196,8 +196,9 @@ public CompletableFuture load(K key) { } else { stats.incrementBatchLoadCountBy(1); // immediate execution of batch function + Object context = loaderOptions.getBatchContextProvider().get(); CompletableFuture> batchedLoad = batchLoadFunction - .load(singletonList(key)) + .load(singletonList(key), context) .toCompletableFuture(); future = batchedLoad .thenApply(list -> list.get(0)); @@ -303,7 +304,8 @@ private CompletableFuture> dispatchQueueBatch(List keys, List> 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); } diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java index 97b19ad..63f1ed3 100644 --- a/src/main/java/org/dataloader/DataLoaderOptions.java +++ b/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; + 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 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; + } } diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 1fa722e..6280362 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -78,6 +78,36 @@ public CompletionStage> load(List userIds) { userLoader.dispatchAndJoin(); } + private static class SecurityCtx { + + public static Object getCallingUserCtx() { + return null; + } + } + + private void callContextExample() { + BatchLoader batchLoader = new BatchLoader() { + @Override + public CompletionStage> load(List keys) { + throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method"); + } + + @Override + public CompletionStage> load(List keys, Object context) { + SecurityCtx callCtx = (SecurityCtx) context; + return callDatabaseForResults(callCtx, keys); + } + + }; + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchContextProvider(() -> SecurityCtx.getCallingUserCtx()); + DataLoader loader = new DataLoader<>(batchLoader, options); + } + + private CompletionStage> callDatabaseForResults(SecurityCtx callCtx, List keys) { + return null; + } + private void tryExample() { Try tryS = Try.tryCall(() -> { diff --git a/src/test/java/org/dataloader/DataLoaderContextTest.java b/src/test/java/org/dataloader/DataLoaderContextTest.java new file mode 100644 index 0000000..786bc97 --- /dev/null +++ b/src/test/java/org/dataloader/DataLoaderContextTest.java @@ -0,0 +1,70 @@ +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 batchLoader = new BatchLoader() { + @Override + public CompletionStage> load(List keys) { + throw new UnsupportedOperationException("this wont be called"); + } + + @Override + public CompletionStage> load(List keys, Object context) { + List list = keys.stream().map(k -> k + "-" + context).collect(Collectors.toList()); + return CompletableFuture.completedFuture(list); + } + }; + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchContextProvider(() -> "ctx"); + DataLoader loader = new DataLoader<>(batchLoader, options); + + loader.load("A"); + loader.load("B"); + loader.loadMany(asList("C", "D")); + + List 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 batchLoader = new BatchLoader() { + @Override + public CompletionStage> load(List keys) { + throw new UnsupportedOperationException("this wont be called"); + } + + @Override + public CompletionStage> load(List keys, Object context) { + List list = keys.stream().map(k -> k + "-" + context).collect(Collectors.toList()); + return CompletableFuture.completedFuture(list); + } + }; + DataLoader loader = new DataLoader<>(batchLoader); + + loader.load("A"); + loader.load("B"); + loader.loadMany(asList("C", "D")); + + List results = loader.dispatchAndJoin(); + + assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null"))); + } +}