From bdecf8b15fb5edd6564a6f8e998beab4cb313673 Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 19 Aug 2018 12:25:28 +1000 Subject: [PATCH 01/11] Added the ability to have call context on batch loader calls --- README.md | 28 +++++++++ .../org/dataloader/BatchContextProvider.java | 13 ++++ src/main/java/org/dataloader/BatchLoader.java | 24 +++++++- src/main/java/org/dataloader/DataLoader.java | 6 +- .../org/dataloader/DataLoaderOptions.java | 22 +++++++ src/test/java/ReadmeExamples.java | 30 ++++++++++ .../org/dataloader/DataLoaderContextTest.java | 59 +++++++++++++++++++ 7 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 src/main/java/org/dataloader/BatchContextProvider.java create mode 100644 src/test/java/org/dataloader/DataLoaderContextTest.java 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..5c588bf --- /dev/null +++ b/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 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 = CompletableFuture::completedFuture; + 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", "B", "C", "D"))); + } +} From 4e3558d40ed70c2ce8f172eeb5e01503f7216f7a Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 19 Aug 2018 12:47:31 +1000 Subject: [PATCH 02/11] Better test of null context provider --- .../org/dataloader/DataLoaderContextTest.java | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/test/java/org/dataloader/DataLoaderContextTest.java b/src/test/java/org/dataloader/DataLoaderContextTest.java index 5c588bf..786bc97 100644 --- a/src/test/java/org/dataloader/DataLoaderContextTest.java +++ b/src/test/java/org/dataloader/DataLoaderContextTest.java @@ -45,7 +45,18 @@ public CompletionStage> load(List keys, Object context) { @Test public void null_is_passed_as_context_if_you_do_nothing() throws Exception { - BatchLoader batchLoader = CompletableFuture::completedFuture; + 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"); @@ -54,6 +65,6 @@ public void null_is_passed_as_context_if_you_do_nothing() throws Exception { List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A", "B", "C", "D"))); + assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null"))); } } From 5a0d364ad7dd5ea3c4fa26087bb32ce07c53133d Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 19 Aug 2018 18:26:33 +1000 Subject: [PATCH 03/11] Added map batch loaders so that you can be more natural in getting less results than asked for --- README.md | 36 ++++ src/main/java/org/dataloader/DataLoader.java | 183 ++++++++++++++++-- .../java/org/dataloader/MapBatchLoader.java | 77 ++++++++ src/test/java/ReadmeExamples.java | 26 ++- .../org/dataloader/DataLoaderContextTest.java | 40 ++++ .../DataLoaderMapBatchLoaderTest.java | 179 +++++++++++++++++ .../java/org/dataloader/DataLoaderTest.java | 72 +------ .../org/dataloader/DataLoaderWithTryTest.java | 103 ++++++++++ src/test/java/org/dataloader/TestKit.java | 23 +++ .../org/dataloader/fixtures/SecurityCtx.java | 8 + .../org/dataloader/fixtures/UserManager.java | 11 ++ 11 files changed, 664 insertions(+), 94 deletions(-) create mode 100644 src/main/java/org/dataloader/MapBatchLoader.java create mode 100644 src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java create mode 100644 src/test/java/org/dataloader/DataLoaderWithTryTest.java create mode 100644 src/test/java/org/dataloader/TestKit.java create mode 100644 src/test/java/org/dataloader/fixtures/SecurityCtx.java diff --git a/README.md b/README.md index c8b2305..092bb4b 100644 --- a/README.md +++ b/README.md @@ -194,6 +194,42 @@ credentials or the database connection parameters. You can do this by implement 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. +### Returning a Map of results from your batch loader + +Often there is not a 1:1 mapping of your batch loaded keys to the values returned. + +For example, let's assume you want to load users from a database, you could probably use a query that looks like this: + +```sql + SELECT * FROM User WHERE id IN (keys) +``` + + Given say 10 user id keys you might only get 7 results back. This can be more naturally represented in a map + than in an order list of values when returning values from the batch loader function. + + You can use `org.dataloader.MapBatchLoader` for this purpose. + + When the map is processed by the `DataLoader` code, any keys that are missing in the map + will be replaced with null values. The semantics that the number of `DataLoader.load` requests + are matched with values is kept. + + Your keys provided MUST be first class keys since they will be used to examine the returned map and + create the list of results, with nulls filling in for missing values. + +```java + MapBatchLoader mapBatchLoader = new MapBatchLoader() { + @Override + public CompletionStage> load(List userIds, Object context) { + SecurityCtx callCtx = (SecurityCtx) context; + return CompletableFuture.supplyAsync(() -> userManager.loadMapOfUsersById(callCtx, userIds)); + } + }; + + DataLoader userLoader = DataLoader.newDataLoader(mapBatchLoader); + + // ... +``` + ### Error object is not a thing in a type safe Java world In the reference JS implementation if the batch loader returns an `Error` object back from the `load()` promise is rejected diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 69b3178..3df85f3 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -24,6 +24,7 @@ import java.util.ArrayList; import java.util.Collection; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.stream.Collectors; @@ -63,6 +64,7 @@ public class DataLoader { private final BatchLoader batchLoadFunction; + private final MapBatchLoader mapBatchLoadFunction; private final DataLoaderOptions loaderOptions; private final CacheMap> futureCache; private final List>> loaderQueue; @@ -96,6 +98,35 @@ public static DataLoader newDataLoader(BatchLoader batchLoadF return new DataLoader<>(batchLoadFunction, options); } + /** + * Creates new DataLoader with the specified map batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param mapBatchLoaderFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newDataLoader(MapBatchLoader mapBatchLoaderFunction) { + return newDataLoader(mapBatchLoaderFunction, null); + } + + /** + * Creates new DataLoader with the specified map batch loader function with the provided options + * + * @param mapBatchLoaderFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newDataLoader(MapBatchLoader mapBatchLoaderFunction, DataLoaderOptions options) { + return new DataLoader<>(mapBatchLoaderFunction, options); + } + + /** * Creates new DataLoader with the specified batch loader function and default options * (batching, caching and unlimited batch size) where the batch loader function returns a list of @@ -134,6 +165,43 @@ public static DataLoader newDataLoaderWithTry(BatchLoader return new DataLoader<>((BatchLoader) batchLoadFunction, options); } + /** + * Creates new DataLoader with the specified map abatch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to + * know gather exact status of each item in a batch call and whether it threw exceptions when fetched then + * you can use this form to create the data loader. + * + * @param mapBatchLoaderFunction the map batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newDataLoaderWithTry(MapBatchLoader> mapBatchLoaderFunction) { + return newDataLoaderWithTry(mapBatchLoaderFunction, null); + } + + /** + * Creates new DataLoader with the specified map batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * @param mapBatchLoaderFunction the map batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + * + * @see #newDataLoaderWithTry(MapBatchLoader) + */ + @SuppressWarnings("unchecked") + public static DataLoader newDataLoaderWithTry(MapBatchLoader> mapBatchLoaderFunction, DataLoaderOptions options) { + return new DataLoader<>((MapBatchLoader) mapBatchLoaderFunction, options); + } /** * Creates a new data loader with the provided batch load function, and default options. @@ -144,6 +212,15 @@ public DataLoader(BatchLoader batchLoadFunction) { this(batchLoadFunction, null); } + /** + * Creates a new data loader with the provided map batch load function. + * + * @param mapBatchLoadFunction the map batch load function to use + */ + public DataLoader(MapBatchLoader mapBatchLoadFunction) { + this(mapBatchLoadFunction, null); + } + /** * Creates a new data loader with the provided batch load function and options. * @@ -151,12 +228,37 @@ public DataLoader(BatchLoader batchLoadFunction) { * @param options the batch load options */ public DataLoader(BatchLoader batchLoadFunction, DataLoaderOptions options) { - this.batchLoadFunction = nonNull(batchLoadFunction); - this.loaderOptions = options == null ? new DataLoaderOptions() : options; + this.batchLoadFunction = batchLoadFunction; + this.mapBatchLoadFunction = null; + this.loaderOptions = determineOptions(options); this.futureCache = determineCacheMap(loaderOptions); // order of keys matter in data loader this.loaderQueue = new ArrayList<>(); - this.stats = nonNull(this.loaderOptions.getStatisticsCollector()); + this.stats = determineCollector(this.loaderOptions); + } + + /** + * Creates a new data loader with the provided map batch load function and options. + * + * @param mapBatchLoadFunction the map batch load function to use + * @param options the batch load options + */ + public DataLoader(MapBatchLoader mapBatchLoadFunction, DataLoaderOptions options) { + this.batchLoadFunction = null; + this.mapBatchLoadFunction = mapBatchLoadFunction; + this.loaderOptions = determineOptions(options); + this.futureCache = determineCacheMap(loaderOptions); + // order of keys matter in data loader + this.loaderQueue = new ArrayList<>(); + this.stats = determineCollector(this.loaderOptions); + } + + private StatisticsCollector determineCollector(DataLoaderOptions loaderOptions) { + return nonNull(loaderOptions.getStatisticsCollector()); + } + + private DataLoaderOptions determineOptions(DataLoaderOptions options) { + return options == null ? new DataLoaderOptions() : options; } @SuppressWarnings("unchecked") @@ -197,11 +299,7 @@ public CompletableFuture load(K key) { stats.incrementBatchLoadCountBy(1); // immediate execution of batch function Object context = loaderOptions.getBatchContextProvider().get(); - CompletableFuture> batchedLoad = batchLoadFunction - .load(singletonList(key), context) - .toCompletableFuture(); - future = batchedLoad - .thenApply(list -> list.get(0)); + future = invokeLoaderImmediately(key, context); } if (cachingEnabled) { futureCache.set(cacheKey, future); @@ -210,6 +308,7 @@ public CompletableFuture load(K key) { } } + /** * Requests to load the list of data provided by the specified keys asynchronously, and returns a composite future * of the resulting values. @@ -232,6 +331,25 @@ public CompletableFuture> loadMany(List keys) { } } + private CompletableFuture invokeLoaderImmediately(K key, Object context) { + List keys = singletonList(key); + CompletionStage 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. *

@@ -302,17 +420,11 @@ private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List< @SuppressWarnings("unchecked") private CompletableFuture> dispatchQueueBatch(List keys, List> queuedFutures) { stats.incrementBatchLoadCountBy(keys.size()); - CompletionStage> batchLoad; - try { - 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); - } + CompletionStage> batchLoad = invokeBatchFunction(keys); return batchLoad .toCompletableFuture() .thenApply(values -> { - assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list"); + assertResultSize(keys, values); for (int idx = 0; idx < queuedFutures.size(); idx++) { Object value = values.get(idx); @@ -351,6 +463,45 @@ private CompletableFuture> dispatchQueueBatch(List keys, List> invokeBatchFunction(List keys) { + CompletionStage> batchLoad; + try { + Object context = loaderOptions.getBatchContextProvider().get(); + if (isMapLoader()) { + batchLoad = invokeMapBatchLoader(keys, context); + } else { + batchLoad = invokeListBatchLoader(keys, context); + } + } catch (Exception e) { + batchLoad = CompletableFutureKit.failedFuture(e); + } + return batchLoad; + } + + private CompletionStage> invokeListBatchLoader(List keys, Object context) { + return nonNull(batchLoadFunction.load(keys, context), "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> invokeMapBatchLoader(List keys, Object context) { + CompletionStage> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise"); + return mapBatchLoad.thenApply(map -> { + List values = new ArrayList<>(); + for (K key : keys) { + V value = map.get(key); + values.add(value); + } + return values; + }); + } + + private void assertResultSize(List keys, List values) { + assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list"); + } + /** * Normally {@link #dispatch()} is an asynchronous operation but this version will 'join' on the * results if dispatch and wait for them to complete. If the {@link CompletableFuture} callbacks make more diff --git a/src/main/java/org/dataloader/MapBatchLoader.java b/src/main/java/org/dataloader/MapBatchLoader.java new file mode 100644 index 0000000..66df3d3 --- /dev/null +++ b/src/main/java/org/dataloader/MapBatchLoader.java @@ -0,0 +1,77 @@ +/* + * Copyright (c) 2016 The original author or authors + * + * All rights reserved. This program and the accompanying materials + * are made available under the terms of the Eclipse Public License v1.0 + * and Apache License v2.0 which accompanies this distribution. + * + * The Eclipse Public License is available at + * http://www.eclipse.org/legal/epl-v10.html + * + * The Apache License v2.0 is available at + * http://www.opensource.org/licenses/apache2.0.php + * + * You may elect to redistribute this code under either of these licenses. + */ + +package org.dataloader; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletionStage; + +/** + * A function that is invoked for batch loading a map of of data values indicated by the provided list of keys. The + * function returns a promise of a map of results of individual load requests. + * + * There are a few constraints that must be upheld: + *

    + *
  • The keys MUST be able to be first class keys in a Java map. Get your equals() and hashCode() methods in order
    • + *
  • The caller of the {@link org.dataloader.DataLoader} that uses this batch loader function MUSt be able to cope with + * null values coming back as results + *
    • + *
  • The function MUST be resilient to the same key being presented twice.
    • + *
+ * + * This form is useful when you don't have a 1:1 mapping of keys to values or when null is an acceptable value for a missing value. + * + * For example, let's assume you want to load users from a database, you could probably use a query that looks like this: + * + * 
+ *    SELECT * FROM User WHERE id IN (keys)
+ *
+ * + * Given say 10 user id keys you might only get 7 results back. This can be more naturally represented in a map + * than in an order list of values when returning values from the batch loader function. + * + * When the map is processed by the {@link org.dataloader.DataLoader} code, any keys that are missing in the map + * will be replaced with null values. The semantics that the number of {@link org.dataloader.DataLoader#load(Object)} requests + * are matched with values is kept. + * + * This means that if 10 keys are asked for then {@link DataLoader#dispatch()} will return a promise 10 value results and each + * of the {@link org.dataloader.DataLoader#load(Object)} will complete with a value, null or an exception. + * + * When caching is disabled, its possible for the same key to be presented in the list of keys more than once. You map + * batch loader function needs to be resilient to this situation. + * + * @param type parameter indicating the type of keys to use for data load requests. + * @param type parameter indicating the type of values returned + * + * @author Arnold Schrijver + * @author Brad Baker + */ +@FunctionalInterface +public interface MapBatchLoader { + + /** + * Called to batch load the provided keys and return a promise to a map of values. It 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. + * + * @param keys the collection of keys to load + * @param context a context object that can help with the call + * + * @return a promise to a map of values for those keys + */ + @SuppressWarnings("unused") + CompletionStage> load(List keys, Object context); +} diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 6280362..378f2a6 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -2,7 +2,9 @@ import org.dataloader.CacheMap; import org.dataloader.DataLoader; import org.dataloader.DataLoaderOptions; +import org.dataloader.MapBatchLoader; import org.dataloader.Try; +import org.dataloader.fixtures.SecurityCtx; import org.dataloader.fixtures.User; import org.dataloader.fixtures.UserManager; import org.dataloader.stats.Statistics; @@ -10,6 +12,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.stream.Collectors; @@ -78,13 +81,6 @@ 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 @@ -108,6 +104,20 @@ private CompletionStage> callDatabaseForResults(SecurityCtx callCtx return null; } + private void mapBatchLoader() { + MapBatchLoader mapBatchLoader = new MapBatchLoader() { + @Override + public CompletionStage> load(List userIds, Object context) { + SecurityCtx callCtx = (SecurityCtx) context; + return CompletableFuture.supplyAsync(() -> userManager.loadMapOfUsersById(callCtx, userIds)); + } + }; + + DataLoader userLoader = DataLoader.newDataLoader(mapBatchLoader); + + // ... + } + private void tryExample() { Try tryS = Try.tryCall(() -> { @@ -230,7 +240,7 @@ private void statsExample() { private void statsConfigExample() { DataLoaderOptions options = DataLoaderOptions.newOptions().setStatisticsCollector(() -> new ThreadLocalStatisticsCollector()); - DataLoader userDataLoader = DataLoader.newDataLoader(userBatchLoader,options); + DataLoader userDataLoader = DataLoader.newDataLoader(userBatchLoader, options); } } diff --git a/src/test/java/org/dataloader/DataLoaderContextTest.java b/src/test/java/org/dataloader/DataLoaderContextTest.java index 786bc97..0fde65a 100644 --- a/src/test/java/org/dataloader/DataLoaderContextTest.java +++ b/src/test/java/org/dataloader/DataLoaderContextTest.java @@ -2,7 +2,9 @@ import org.junit.Test; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.stream.Collectors; @@ -43,6 +45,26 @@ public CompletionStage> load(List keys, Object context) { assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx"))); } + @Test + public void context_is_passed_to_map_batch_loader_function() throws Exception { + MapBatchLoader mapBatchLoader = (keys, context) -> { + Map map = new HashMap<>(); + keys.forEach(k -> map.put(k, k + "-" + context)); + return CompletableFuture.completedFuture(map); + }; + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchContextProvider(() -> "ctx"); + DataLoader loader = new DataLoader<>(mapBatchLoader, 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() { @@ -67,4 +89,22 @@ public CompletionStage> load(List keys, Object context) { assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null"))); } + + @Test + public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() throws Exception { + MapBatchLoader mapBatchLoader = (keys, context) -> { + Map map = new HashMap<>(); + keys.forEach(k -> map.put(k, k + "-" + context)); + return CompletableFuture.completedFuture(map); + }; + DataLoader loader = new DataLoader<>(mapBatchLoader); + + 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"))); + } } diff --git a/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java b/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java new file mode 100644 index 0000000..8223095 --- /dev/null +++ b/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java @@ -0,0 +1,179 @@ +package org.dataloader; + +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; + +import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; +import static org.awaitility.Awaitility.await; +import static org.dataloader.DataLoaderOptions.newOptions; +import static org.dataloader.TestKit.futureError; +import static org.dataloader.TestKit.listFrom; +import static org.dataloader.impl.CompletableFutureKit.cause; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.is; +import static org.junit.Assert.assertThat; + +/** + * Much of the tests that related to {@link org.dataloader.MapBatchLoader} also related to + * {@link org.dataloader.BatchLoader}. This is white box testing somewhat because we could have repeated + * ALL the tests in {@link org.dataloader.DataLoaderTest} here as well but chose not to because we KNOW that + * DataLoader differs only a little in how it handles the 2 types of loader functions. We choose to grab some + * common functionality for repeat testing and otherwise rely on the very complete other tests. + */ +public class DataLoaderMapBatchLoaderTest { + + MapBatchLoader evensOnlyMapBatchLoader = (keys, context) -> { + Map mapOfResults = new HashMap<>(); + for (int i = 0; i < keys.size(); i++) { + String k = keys.get(i); + if (i % 2 == 0) { + mapOfResults.put(k, k); + } + } + return CompletableFuture.completedFuture(mapOfResults); + }; + + private static DataLoader idMapLoader(DataLoaderOptions options, List> loadCalls) { + MapBatchLoader kvBatchLoader = (keys, context) -> { + loadCalls.add(new ArrayList<>(keys)); + Map map = new HashMap<>(); + //noinspection unchecked + keys.forEach(k -> map.put(k, (V) k)); + return CompletableFuture.completedFuture(map); + }; + return DataLoader.newDataLoader(kvBatchLoader, options); + } + + private static DataLoader idMapLoaderBlowsUps( + DataLoaderOptions options, List> loadCalls) { + return new DataLoader<>((keys, context) -> { + loadCalls.add(new ArrayList<>(keys)); + return futureError(); + }, options); + } + + + @Test + public void basic_map_batch_loading() throws Exception { + DataLoader loader = new DataLoader<>(evensOnlyMapBatchLoader); + + loader.load("A"); + loader.load("B"); + loader.loadMany(asList("C", "D")); + + List results = loader.dispatchAndJoin(); + + assertThat(results.size(), equalTo(4)); + assertThat(results, equalTo(asList("A", null, "C", null))); + } + + + @Test + public void should_map_Batch_multiple_requests() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idMapLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load(1); + CompletableFuture future2 = identityLoader.load(2); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo(1)); + assertThat(future2.get(), equalTo(2)); + assertThat(loadCalls, equalTo(singletonList(asList(1, 2)))); + } + + @Test + public void can_split_max_batch_sizes_correctly() throws Exception { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idMapLoader(newOptions().setMaxBatchSize(5), loadCalls); + + for (int i = 0; i < 21; i++) { + identityLoader.load(i); + } + List> expectedCalls = new ArrayList<>(); + expectedCalls.add(listFrom(0, 5)); + expectedCalls.add(listFrom(5, 10)); + expectedCalls.add(listFrom(10, 15)); + expectedCalls.add(listFrom(15, 20)); + expectedCalls.add(listFrom(20, 21)); + + List result = identityLoader.dispatch().join(); + + assertThat(result, equalTo(listFrom(0, 21))); + assertThat(loadCalls, equalTo(expectedCalls)); + } + + @Test + public void should_Propagate_error_to_all_loads() { + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idMapLoaderBlowsUps(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + CompletableFuture future2 = errorLoader.load(2); + errorLoader.dispatch(); + + await().until(future1::isDone); + + assertThat(future1.isCompletedExceptionally(), is(true)); + Throwable cause = cause(future1); + assert cause != null; + assertThat(cause, instanceOf(IllegalStateException.class)); + assertThat(cause.getMessage(), equalTo("Error")); + + await().until(future2::isDone); + cause = cause(future2); + assert cause != null; + assertThat(cause.getMessage(), equalTo(cause.getMessage())); + + assertThat(loadCalls, equalTo(singletonList(asList(1, 2)))); + } + + @Test + public void should_work_with_duplicate_keys_when_caching_disabled() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idMapLoader(newOptions().setCachingEnabled(false), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture future3 = identityLoader.load("A"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(future3.get(), equalTo("A")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B", "A")))); + } + + @Test + public void should_work_with_duplicate_keys_when_caching_enabled() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idMapLoader(newOptions().setCachingEnabled(true), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture future3 = identityLoader.load("A"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(future3.get(), equalTo("A")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + } + + + +} diff --git a/src/test/java/org/dataloader/DataLoaderTest.java b/src/test/java/org/dataloader/DataLoaderTest.java index 56b743b..77bcc3e 100644 --- a/src/test/java/org/dataloader/DataLoaderTest.java +++ b/src/test/java/org/dataloader/DataLoaderTest.java @@ -19,7 +19,6 @@ import org.dataloader.fixtures.User; import org.dataloader.fixtures.UserManager; import org.dataloader.impl.CompletableFutureKit; -import org.hamcrest.Matchers; import org.junit.Test; import java.util.ArrayList; @@ -39,8 +38,8 @@ import static java.util.Collections.singletonList; import static org.awaitility.Awaitility.await; import static org.dataloader.DataLoaderOptions.newOptions; +import static org.dataloader.TestKit.listFrom; import static org.dataloader.impl.CompletableFutureKit.cause; -import static org.dataloader.impl.CompletableFutureKit.failedFuture; import static org.hamcrest.Matchers.empty; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; @@ -855,14 +854,6 @@ public void can_split_max_batch_sizes_correctly() throws Exception { } - private Collection listFrom(int i, int max) { - List ints = new ArrayList<>(); - for (int j = i; j < max; j++) { - ints.add(j); - } - return ints; - } - @Test public void should_Batch_loads_occurring_within_futures() { List> loadCalls = new ArrayList<>(); @@ -978,61 +969,6 @@ public void should_allow_composition_of_data_loader_calls() throws Exception { assertThat(allResults.size(), equalTo(4)); } - @Test - public void should_handle_Trys_coming_back_from_batchLoader() throws Exception { - - List> batchKeyCalls = new ArrayList<>(); - BatchLoader> batchLoader = keys -> { - batchKeyCalls.add(keys); - - List> result = new ArrayList<>(); - for (String key : keys) { - if ("bang".equalsIgnoreCase(key)) { - result.add(Try.failed(new RuntimeException(key))); - } else { - result.add(Try.succeeded(key)); - } - } - return CompletableFuture.completedFuture(result); - }; - - DataLoader dataLoader = DataLoader.newDataLoaderWithTry(batchLoader); - - CompletableFuture a = dataLoader.load("A"); - CompletableFuture b = dataLoader.load("B"); - CompletableFuture bang = dataLoader.load("bang"); - - dataLoader.dispatch(); - - assertThat(a.get(), equalTo("A")); - assertThat(b.get(), equalTo("B")); - assertThat(bang.isCompletedExceptionally(), equalTo(true)); - bang.whenComplete((s, throwable) -> { - assertThat(s, nullValue()); - assertThat(throwable, Matchers.instanceOf(RuntimeException.class)); - assertThat(throwable.getMessage(), equalTo("Bang")); - }); - - assertThat(batchKeyCalls, equalTo(singletonList(asList("A", "B", "bang")))); - - a = dataLoader.load("A"); - b = dataLoader.load("B"); - bang = dataLoader.load("bang"); - - dataLoader.dispatch(); - - assertThat(a.get(), equalTo("A")); - assertThat(b.get(), equalTo("B")); - assertThat(bang.isCompletedExceptionally(), equalTo(true)); - bang.whenComplete((s, throwable) -> { - assertThat(s, nullValue()); - assertThat(throwable, Matchers.instanceOf(RuntimeException.class)); - assertThat(throwable.getMessage(), equalTo("Bang")); - }); - - // the failed value should have been cached as per Facebook DL behaviour - assertThat(batchKeyCalls, equalTo(singletonList(asList("A", "B", "bang")))); - } private static CacheKey getJsonObjectCacheMapFn() { @@ -1057,7 +993,7 @@ private static DataLoader idLoaderBlowsUps( DataLoaderOptions options, List> loadCalls) { return new DataLoader<>(keys -> { loadCalls.add(new ArrayList<>(keys)); - return futureError(); + return TestKit.futureError(); }, options); } @@ -1088,10 +1024,6 @@ private static DataLoader idLoaderOddEvenExceptions( }, options); } - private static CompletableFuture futureError() { - return failedFuture(new IllegalStateException("Error")); - } - private BatchLoader keysAsValues() { return CompletableFuture::completedFuture; diff --git a/src/test/java/org/dataloader/DataLoaderWithTryTest.java b/src/test/java/org/dataloader/DataLoaderWithTryTest.java new file mode 100644 index 0000000..08fdda2 --- /dev/null +++ b/src/test/java/org/dataloader/DataLoaderWithTryTest.java @@ -0,0 +1,103 @@ +package org.dataloader; + +import org.hamcrest.Matchers; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.nullValue; +import static org.junit.Assert.assertThat; + +public class DataLoaderWithTryTest { + + @Test + public void should_handle_Trys_coming_back_from_batchLoader() throws Exception { + + List> batchKeyCalls = new ArrayList<>(); + BatchLoader> batchLoader = keys -> { + batchKeyCalls.add(keys); + + List> result = new ArrayList<>(); + for (String key : keys) { + if ("bang".equalsIgnoreCase(key)) { + result.add(Try.failed(new RuntimeException(key))); + } else { + result.add(Try.succeeded(key)); + } + } + return CompletableFuture.completedFuture(result); + }; + + DataLoader dataLoader = DataLoader.newDataLoaderWithTry(batchLoader); + + commonTryAsserts(batchKeyCalls, dataLoader); + } + + @Test + public void should_handle_Trys_coming_back_from_map_batchLoader() throws Exception { + + List> batchKeyCalls = new ArrayList<>(); + MapBatchLoader> batchLoader = (keys, context) -> { + batchKeyCalls.add(keys); + + Map> result = new HashMap<>(); + for (String key : keys) { + if ("bang".equalsIgnoreCase(key)) { + result.put(key, Try.failed(new RuntimeException(key))); + } else { + result.put(key, Try.succeeded(key)); + } + } + return CompletableFuture.completedFuture(result); + }; + + DataLoader dataLoader = DataLoader.newDataLoaderWithTry(batchLoader); + + commonTryAsserts(batchKeyCalls, dataLoader); + } + + private void commonTryAsserts(List> batchKeyCalls, DataLoader dataLoader) throws InterruptedException, java.util.concurrent.ExecutionException { + CompletableFuture a = dataLoader.load("A"); + CompletableFuture b = dataLoader.load("B"); + CompletableFuture bang = dataLoader.load("bang"); + + dataLoader.dispatch(); + + assertThat(a.get(), equalTo("A")); + assertThat(b.get(), equalTo("B")); + assertThat(bang.isCompletedExceptionally(), equalTo(true)); + bang.whenComplete((s, throwable) -> { + assertThat(s, nullValue()); + assertThat(throwable, Matchers.instanceOf(RuntimeException.class)); + assertThat(throwable.getMessage(), equalTo("Bang")); + }); + + assertThat(batchKeyCalls, equalTo(singletonList(asList("A", "B", "bang")))); + + a = dataLoader.load("A"); + b = dataLoader.load("B"); + bang = dataLoader.load("bang"); + + dataLoader.dispatch(); + + assertThat(a.get(), equalTo("A")); + assertThat(b.get(), equalTo("B")); + assertThat(bang.isCompletedExceptionally(), equalTo(true)); + bang.whenComplete((s, throwable) -> { + assertThat(s, nullValue()); + assertThat(throwable, Matchers.instanceOf(RuntimeException.class)); + assertThat(throwable.getMessage(), equalTo("Bang")); + }); + + // the failed value should have been cached as per Facebook DL behaviour + assertThat(batchKeyCalls, equalTo(singletonList(asList("A", "B", "bang")))); + } + +} diff --git a/src/test/java/org/dataloader/TestKit.java b/src/test/java/org/dataloader/TestKit.java new file mode 100644 index 0000000..82f73d6 --- /dev/null +++ b/src/test/java/org/dataloader/TestKit.java @@ -0,0 +1,23 @@ +package org.dataloader; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static org.dataloader.impl.CompletableFutureKit.failedFuture; + +public class TestKit { + + public static Collection listFrom(int i, int max) { + List ints = new ArrayList<>(); + for (int j = i; j < max; j++) { + ints.add(j); + } + return ints; + } + + static CompletableFuture futureError() { + return failedFuture(new IllegalStateException("Error")); + } +} diff --git a/src/test/java/org/dataloader/fixtures/SecurityCtx.java b/src/test/java/org/dataloader/fixtures/SecurityCtx.java new file mode 100644 index 0000000..79bdfb0 --- /dev/null +++ b/src/test/java/org/dataloader/fixtures/SecurityCtx.java @@ -0,0 +1,8 @@ +package org.dataloader.fixtures; + +public class SecurityCtx { + + public static Object getCallingUserCtx() { + return null; + } +} diff --git a/src/test/java/org/dataloader/fixtures/UserManager.java b/src/test/java/org/dataloader/fixtures/UserManager.java index 4fdc5bd..95b712e 100644 --- a/src/test/java/org/dataloader/fixtures/UserManager.java +++ b/src/test/java/org/dataloader/fixtures/UserManager.java @@ -1,10 +1,12 @@ package org.dataloader.fixtures; +import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; +@SuppressWarnings({"unused", "SpellCheckingInspection", "NonAsciiCharacters"}) public class UserManager { public static final User ILÚVATAR = new User(-1L, -1L, "Ilúvatar"); @@ -48,4 +50,13 @@ public User loadUserById(Long userId) { public List loadUsersById(List userIds) { return userIds.stream().map(this::loadUserById).collect(Collectors.toList()); } + + public Map loadMapOfUsersById(SecurityCtx callCtx, List userIds) { + Map map = new HashMap<>(); + userIds.forEach(userId -> { + User user = loadUserById(userId); + map.put(userId, user); + }); + return map; + } } From f263dbb65f0130ff7231750d4ccecf6df3e6092b Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Mon, 20 Aug 2018 18:11:43 +1000 Subject: [PATCH 04/11] Added the BatchLoadingEnvironment as a way to hang future objects into the call context --- README.md | 34 ++++++----- .../org/dataloader/BatchContextProvider.java | 13 ---- src/main/java/org/dataloader/BatchLoader.java | 17 +++--- .../dataloader/BatchLoaderEnvironment.java | 40 +++++++++++++ .../BatchLoaderEnvironmentProvider.java | 14 +++++ src/main/java/org/dataloader/DataLoader.java | 60 +++++++++---------- .../org/dataloader/DataLoaderOptions.java | 22 +++---- .../java/org/dataloader/MapBatchLoader.java | 10 ++-- src/test/java/ReadmeExamples.java | 23 ++++--- ...DataLoaderBatchLoaderEnvironmentTest.java} | 25 ++++---- .../DataLoaderMapBatchLoaderTest.java | 6 +- .../org/dataloader/DataLoaderWithTryTest.java | 2 +- 12 files changed, 162 insertions(+), 104 deletions(-) delete mode 100644 src/main/java/org/dataloader/BatchContextProvider.java create mode 100644 src/main/java/org/dataloader/BatchLoaderEnvironment.java create mode 100644 src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java rename src/test/java/org/dataloader/{DataLoaderContextTest.java => DataLoaderBatchLoaderEnvironmentTest.java} (78%) diff --git a/README.md b/README.md index 092bb4b..2921770 100644 --- a/README.md +++ b/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 batchLoader = new BatchLoader() { + // + // 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> load(List keys) { - throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method"); + throw new UnsupportedOperationException(); } @Override - public CompletionStage> load(List keys, Object context) { - SecurityCtx callCtx = (SecurityCtx) context; + public CompletionStage> load(List 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 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 mapBatchLoader = new MapBatchLoader() { @Override - public CompletionStage> load(List userIds, Object context) { - SecurityCtx callCtx = (SecurityCtx) context; + public CompletionStage> load(List 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 dataLoader = DataLoader.newDataLoaderWithTry(new BatchLoader>() { @Override - public CompletionStage>> load(List keys) { + public CompletionStage>> load(List keys, BatchLoaderEnvironment environment) { return CompletableFuture.supplyAsync(() -> { List> 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 diff --git a/src/main/java/org/dataloader/BatchContextProvider.java b/src/main/java/org/dataloader/BatchContextProvider.java deleted file mode 100644 index 0a43664..0000000 --- a/src/main/java/org/dataloader/BatchContextProvider.java +++ /dev/null @@ -1,13 +0,0 @@ -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 8791f14..871b6f5 100644 --- a/src/main/java/org/dataloader/BatchLoader.java +++ b/src/main/java/org/dataloader/BatchLoader.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 { /** * 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> 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 + * 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> load(List keys, Object context) { + @SuppressWarnings({"unused", "deprecation"}) + default CompletionStage> load(List keys, BatchLoaderEnvironment environment) { return load(keys); } } diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironment.java b/src/main/java/org/dataloader/BatchLoaderEnvironment.java new file mode 100644 index 0000000..969c7c6 --- /dev/null +++ b/src/main/java/org/dataloader/BatchLoaderEnvironment.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 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); + } + } +} diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java b/src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java new file mode 100644 index 0000000..b614e24 --- /dev/null +++ b/src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.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(); +} \ No newline at end of file diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 3df85f3..404c48d 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -298,8 +298,7 @@ public CompletableFuture 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> loadMany(List keys) { } } - private CompletableFuture invokeLoaderImmediately(K key, Object context) { - List keys = singletonList(key); - CompletionStage 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. *

@@ -420,7 +400,7 @@ private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List< @SuppressWarnings("unchecked") private CompletableFuture> dispatchQueueBatch(List keys, List> queuedFutures) { stats.incrementBatchLoadCountBy(keys.size()); - CompletionStage> batchLoad = invokeBatchFunction(keys); + CompletionStage> batchLoad = invokeLoader(keys); return batchLoad .toCompletableFuture() .thenApply(values -> { @@ -463,14 +443,34 @@ private CompletableFuture> dispatchQueueBatch(List keys, List> invokeBatchFunction(List keys) { + private boolean isMapLoader() { + return mapBatchLoadFunction != null; + } + + private CompletableFuture invokeLoaderImmediately(K key) { + BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); + List keys = singletonList(key); + CompletionStage 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> invokeLoader(List keys) { CompletionStage> 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); @@ -478,16 +478,16 @@ private CompletionStage> invokeBatchFunction(List keys) { return batchLoad; } - private CompletionStage> invokeListBatchLoader(List keys, Object context) { - return nonNull(batchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise"); + private CompletionStage> invokeListBatchLoader(List 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> invokeMapBatchLoader(List keys, Object context) { - CompletionStage> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise"); + private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderEnvironment environment) { + CompletionStage> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, environment), "Your batch loader function MUST return a non null CompletionStage promise"); return mapBatchLoad.thenApply(map -> { List values = new ArrayList<>(); for (K key : keys) { diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java index 63f1ed3..8cab094 100644 --- a/src/main/java/org/dataloader/DataLoaderOptions.java +++ b/src/main/java/org/dataloader/DataLoaderOptions.java @@ -31,7 +31,7 @@ */ 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; @@ -39,7 +39,7 @@ public class DataLoaderOptions { private CacheMap cacheMap; private int maxBatchSize; private Supplier 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 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; } } diff --git a/src/main/java/org/dataloader/MapBatchLoader.java b/src/main/java/org/dataloader/MapBatchLoader.java index 66df3d3..5513664 100644 --- a/src/main/java/org/dataloader/MapBatchLoader.java +++ b/src/main/java/org/dataloader/MapBatchLoader.java @@ -57,21 +57,19 @@ * @param type parameter indicating the type of keys to use for data load requests. * @param type parameter indicating the type of values returned * - * @author Arnold Schrijver * @author Brad Baker */ -@FunctionalInterface public interface MapBatchLoader { /** - * 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> load(List keys, Object context); + CompletionStage> load(List keys, BatchLoaderEnvironment environment); } diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 378f2a6..4f8db88 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -1,4 +1,5 @@ import org.dataloader.BatchLoader; +import org.dataloader.BatchLoaderEnvironment; import org.dataloader.CacheMap; import org.dataloader.DataLoader; import org.dataloader.DataLoaderOptions; @@ -83,20 +84,28 @@ public CompletionStage> load(List userIds) { private void callContextExample() { BatchLoader batchLoader = new BatchLoader() { + // + // 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> load(List keys) { - throw new UnsupportedOperationException("This wont be called if you implement the other defaulted method"); + throw new UnsupportedOperationException(); } @Override - public CompletionStage> load(List keys, Object context) { - SecurityCtx callCtx = (SecurityCtx) context; + public CompletionStage> load(List 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 loader = new DataLoader<>(batchLoader, options); } @@ -107,8 +116,8 @@ private CompletionStage> callDatabaseForResults(SecurityCtx callCtx private void mapBatchLoader() { MapBatchLoader mapBatchLoader = new MapBatchLoader() { @Override - public CompletionStage> load(List userIds, Object context) { - SecurityCtx callCtx = (SecurityCtx) context; + public CompletionStage> load(List userIds, BatchLoaderEnvironment environment) { + SecurityCtx callCtx = environment.getContext(); return CompletableFuture.supplyAsync(() -> userManager.loadMapOfUsersById(callCtx, userIds)); } }; diff --git a/src/test/java/org/dataloader/DataLoaderContextTest.java b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java similarity index 78% rename from src/test/java/org/dataloader/DataLoaderContextTest.java rename to src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java index 0fde65a..236541c 100644 --- a/src/test/java/org/dataloader/DataLoaderContextTest.java +++ b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java @@ -10,30 +10,31 @@ import java.util.stream.Collectors; import static java.util.Arrays.asList; +import static org.dataloader.BatchLoaderEnvironment.newBatchLoaderEnvironment; 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 { +public class DataLoaderBatchLoaderEnvironmentTest { @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"); + throw new UnsupportedOperationException(); } @Override - public CompletionStage> load(List keys, Object context) { - List list = keys.stream().map(k -> k + "-" + context).collect(Collectors.toList()); + public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { + List list = keys.stream().map(k -> k + "-" + environment.getContext()).collect(Collectors.toList()); return CompletableFuture.completedFuture(list); } }; DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchContextProvider(() -> "ctx"); + .setBatchLoaderEnvironmentProvider(() -> newBatchLoaderEnvironment().context("ctx").build()); DataLoader loader = new DataLoader<>(batchLoader, options); loader.load("A"); @@ -47,13 +48,13 @@ public CompletionStage> load(List keys, Object context) { @Test public void context_is_passed_to_map_batch_loader_function() throws Exception { - MapBatchLoader mapBatchLoader = (keys, context) -> { + MapBatchLoader mapBatchLoader = (keys, environment) -> { Map map = new HashMap<>(); - keys.forEach(k -> map.put(k, k + "-" + context)); + keys.forEach(k -> map.put(k, k + "-" + environment.getContext())); return CompletableFuture.completedFuture(map); }; DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchContextProvider(() -> "ctx"); + .setBatchLoaderEnvironmentProvider(() -> newBatchLoaderEnvironment().context("ctx").build()); DataLoader loader = new DataLoader<>(mapBatchLoader, options); loader.load("A"); @@ -74,8 +75,8 @@ public CompletionStage> load(List keys) { } @Override - public CompletionStage> load(List keys, Object context) { - List list = keys.stream().map(k -> k + "-" + context).collect(Collectors.toList()); + public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { + List list = keys.stream().map(k -> k + "-" + environment.getContext()).collect(Collectors.toList()); return CompletableFuture.completedFuture(list); } }; @@ -92,9 +93,9 @@ public CompletionStage> load(List keys, Object context) { @Test public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() throws Exception { - MapBatchLoader mapBatchLoader = (keys, context) -> { + MapBatchLoader mapBatchLoader = (keys, environment) -> { Map map = new HashMap<>(); - keys.forEach(k -> map.put(k, k + "-" + context)); + keys.forEach(k -> map.put(k, k + "-" + environment.getContext())); return CompletableFuture.completedFuture(map); }; DataLoader loader = new DataLoader<>(mapBatchLoader); diff --git a/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java b/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java index 8223095..f6aa397 100644 --- a/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java +++ b/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java @@ -31,7 +31,7 @@ */ public class DataLoaderMapBatchLoaderTest { - MapBatchLoader evensOnlyMapBatchLoader = (keys, context) -> { + MapBatchLoader evensOnlyMapBatchLoader = (keys, environment) -> { Map mapOfResults = new HashMap<>(); for (int i = 0; i < keys.size(); i++) { String k = keys.get(i); @@ -43,7 +43,7 @@ public class DataLoaderMapBatchLoaderTest { }; private static DataLoader idMapLoader(DataLoaderOptions options, List> loadCalls) { - MapBatchLoader kvBatchLoader = (keys, context) -> { + MapBatchLoader kvBatchLoader = (keys, environment) -> { loadCalls.add(new ArrayList<>(keys)); Map map = new HashMap<>(); //noinspection unchecked @@ -55,7 +55,7 @@ private static DataLoader idMapLoader(DataLoaderOptions options, Li private static DataLoader idMapLoaderBlowsUps( DataLoaderOptions options, List> loadCalls) { - return new DataLoader<>((keys, context) -> { + return new DataLoader<>((keys, environment) -> { loadCalls.add(new ArrayList<>(keys)); return futureError(); }, options); diff --git a/src/test/java/org/dataloader/DataLoaderWithTryTest.java b/src/test/java/org/dataloader/DataLoaderWithTryTest.java index 08fdda2..ee6c459 100644 --- a/src/test/java/org/dataloader/DataLoaderWithTryTest.java +++ b/src/test/java/org/dataloader/DataLoaderWithTryTest.java @@ -44,7 +44,7 @@ public void should_handle_Trys_coming_back_from_batchLoader() throws Exception { public void should_handle_Trys_coming_back_from_map_batchLoader() throws Exception { List> batchKeyCalls = new ArrayList<>(); - MapBatchLoader> batchLoader = (keys, context) -> { + MapBatchLoader> batchLoader = (keys, environment) -> { batchKeyCalls.add(keys); Map> result = new HashMap<>(); From dc3e3b7d9f3ce06f897af431c26511f377094178 Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sat, 25 Aug 2018 19:42:09 +1000 Subject: [PATCH 05/11] Now with repeated factory methods of the new different loader interfaces --- README.md | 34 +-- src/main/java/org/dataloader/BatchLoader.java | 24 +- .../dataloader/BatchLoaderWithContext.java | 26 ++ src/main/java/org/dataloader/DataLoader.java | 285 ++++++++++++------ ...atchLoader.java => MappedBatchLoader.java} | 9 +- .../MappedBatchLoaderWithContext.java | 42 +++ src/test/java/ReadmeExamples.java | 36 +-- src/test/java/TestFI.java | 63 ++++ .../java/org/dataloader/CustomCacheMap.java | 41 +++ .../DataLoaderBatchLoaderEnvironmentTest.java | 41 +-- .../DataLoaderMapBatchLoaderTest.java | 13 +- .../java/org/dataloader/DataLoaderTest.java | 38 --- .../org/dataloader/DataLoaderWithTryTest.java | 6 +- 13 files changed, 424 insertions(+), 234 deletions(-) create mode 100644 src/main/java/org/dataloader/BatchLoaderWithContext.java rename src/main/java/org/dataloader/{MapBatchLoader.java => MappedBatchLoader.java} (86%) create mode 100644 src/main/java/org/dataloader/MappedBatchLoaderWithContext.java create mode 100644 src/test/java/TestFI.java create mode 100644 src/test/java/org/dataloader/CustomCacheMap.java diff --git a/README.md b/README.md index 2921770..d3ce756 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ a list of keys } }; - DataLoader userLoader = new DataLoader<>(userBatchLoader); + DataLoader userLoader = DataLoader.newDataLoader(userBatchLoader); ``` @@ -169,20 +169,17 @@ That said, with key caching turn on (the default), it will still be more efficie 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`. +`org.dataloader.BatchLoaderEnvironmentProvider` and using one of the `xxxWithContext` batch loading interfaces +such as `org.dataloader.BatchLoaderWithContext`. ```java - BatchLoader batchLoader = new BatchLoader() { - // - // 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> load(List keys) { - throw new UnsupportedOperationException(); - } + BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment() + .context(SecurityCtx.getCallingUserCtx()).build(); + + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment); + BatchLoaderWithContext batchLoader = new BatchLoaderWithContext() { @Override public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { SecurityCtx callCtx = environment.getContext(); @@ -190,12 +187,7 @@ credentials or the database connection parameters. You can do this by implement } }; - BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment() - .context(SecurityCtx.getCallingUserCtx()).build(); - - DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment); - DataLoader loader = new DataLoader<>(batchLoader, options); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); ``` The batch loading code will now receive this environment object and it can be used to get context perhaps allowing it @@ -214,7 +206,7 @@ For example, let's assume you want to load users from a database, you could prob Given say 10 user id keys you might only get 7 results back. This can be more naturally represented in a map than in an order list of values when returning values from the batch loader function. - You can use `org.dataloader.MapBatchLoader` for this purpose. + You can use `org.dataloader.MappedBatchLoader` for this purpose. When the map is processed by the `DataLoader` code, any keys that are missing in the map will be replaced with null values. The semantics that the number of `DataLoader.load` requests @@ -297,7 +289,7 @@ react to that, in a type safe manner. In certain uncommon cases, a DataLoader which does not cache may be desirable. ```java - new DataLoader(userBatchLoader, DataLoaderOptions.newOptions().setCachingEnabled(false)); + DataLoader.newDataLoader(userBatchLoader, DataLoaderOptions.newOptions().setCachingEnabled(false)); ``` Calling the above will ensure that every call to `.load()` will produce a new promise, and requested keys will not be saved in memory. @@ -391,7 +383,7 @@ However you can create your own custom cache and supply it to the data loader on ```java MyCustomCache customCache = new MyCustomCache(); DataLoaderOptions options = DataLoaderOptions.newOptions().setCacheMap(customCache); - new DataLoader(userBatchLoader, options); + DataLoader.newDataLoader(userBatchLoader, options); ``` You could choose to use one of the fancy cache implementations from Guava or Kaffeine and wrap it in a `CacheMap` wrapper ready diff --git a/src/main/java/org/dataloader/BatchLoader.java b/src/main/java/org/dataloader/BatchLoader.java index 871b6f5..463cac6 100644 --- a/src/main/java/org/dataloader/BatchLoader.java +++ b/src/main/java/org/dataloader/BatchLoader.java @@ -78,32 +78,12 @@ public interface BatchLoader { /** * 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, BatchLoaderEnvironment)} method - * instead which is the preferred method. + * If you need calling context then implement {@link org.dataloader.BatchLoaderWithContext} * * @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 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 environment an environment object that can help with the call - * - * @return a promise of the values for those keys - */ - @SuppressWarnings({"unused", "deprecation"}) - default CompletionStage> load(List keys, BatchLoaderEnvironment environment) { - return load(keys); - } } + diff --git a/src/main/java/org/dataloader/BatchLoaderWithContext.java b/src/main/java/org/dataloader/BatchLoaderWithContext.java new file mode 100644 index 0000000..1ee8f66 --- /dev/null +++ b/src/main/java/org/dataloader/BatchLoaderWithContext.java @@ -0,0 +1,26 @@ +package org.dataloader; + +import java.util.List; +import java.util.concurrent.CompletionStage; + +/** + * This form of {@link org.dataloader.BatchLoader} is given a {@link org.dataloader.BatchLoaderEnvironment} object + * that encapsulates the calling context. A typical use case is passing in security credentials or database connection details + * say. + * + * See {@link org.dataloader.BatchLoader} for more details on the design invariants that you must implement in order to + * use this interface. + */ +public interface BatchLoaderWithContext { + /** + * Called to batch load the provided keys and return a promise to a list of values. This default + * 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. + * + * @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 + */ + CompletionStage> load(List keys, BatchLoaderEnvironment environment); +} diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 404c48d..d97bcbd 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -63,8 +63,7 @@ */ public class DataLoader { - private final BatchLoader batchLoadFunction; - private final MapBatchLoader mapBatchLoadFunction; + private final Object batchLoadFunction; private final DataLoaderOptions loaderOptions; private final CacheMap> futureCache; private final List>> loaderQueue; @@ -99,33 +98,70 @@ public static DataLoader newDataLoader(BatchLoader batchLoadF } /** - * Creates new DataLoader with the specified map batch loader function and default options - * (batching, caching and unlimited batch size). + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to + * know the exact status of each item in a batch call and whether it threw exceptions when fetched then + * you can use this form to create the data loader. * - * @param mapBatchLoaderFunction the batch load function to use - * @param the key type - * @param the value type + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type * * @return a new DataLoader */ - public static DataLoader newDataLoader(MapBatchLoader mapBatchLoaderFunction) { - return newDataLoader(mapBatchLoaderFunction, null); + public static DataLoader newDataLoaderWithTry(BatchLoader> batchLoadFunction) { + return newDataLoaderWithTry(batchLoadFunction, null); } /** - * Creates new DataLoader with the specified map batch loader function with the provided options + * Creates new DataLoader with the specified batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. * - * @param mapBatchLoaderFunction the batch load function to use - * @param options the options to use - * @param the key type - * @param the value type + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type * * @return a new DataLoader + * + * @see #newDataLoaderWithTry(BatchLoader) */ - public static DataLoader newDataLoader(MapBatchLoader mapBatchLoaderFunction, DataLoaderOptions options) { - return new DataLoader<>(mapBatchLoaderFunction, options); + @SuppressWarnings("unchecked") + public static DataLoader newDataLoaderWithTry(BatchLoader> batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>((BatchLoader) batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newDataLoader(BatchLoaderWithContext batchLoadFunction) { + return newDataLoader(batchLoadFunction, null); } + /** + * Creates new DataLoader with the specified batch loader function with the provided options + * + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newDataLoader(BatchLoaderWithContext batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>(batchLoadFunction, options); + } /** * Creates new DataLoader with the specified batch loader function and default options @@ -133,7 +169,7 @@ public static DataLoader newDataLoader(MapBatchLoader mapBatc * {@link org.dataloader.Try} objects. * * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to - * know gther exact status of each item in a batch call and whether it threw exceptions when fetched then + * know the exact status of each item in a batch call and whether it threw exceptions when fetched then * you can use this form to create the data loader. * * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects @@ -142,7 +178,7 @@ public static DataLoader newDataLoader(MapBatchLoader mapBatc * * @return a new DataLoader */ - public static DataLoader newDataLoaderWithTry(BatchLoader> batchLoadFunction) { + public static DataLoader newDataLoaderWithTry(BatchLoaderWithContext> batchLoadFunction) { return newDataLoaderWithTry(batchLoadFunction, null); } @@ -160,105 +196,168 @@ public static DataLoader newDataLoaderWithTry(BatchLoader * * @see #newDataLoaderWithTry(BatchLoader) */ - @SuppressWarnings("unchecked") - public static DataLoader newDataLoaderWithTry(BatchLoader> batchLoadFunction, DataLoaderOptions options) { - return new DataLoader<>((BatchLoader) batchLoadFunction, options); + public static DataLoader newDataLoaderWithTry(BatchLoaderWithContext> batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>(batchLoadFunction, options); } /** - * Creates new DataLoader with the specified map abatch loader function and default options + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedDataLoader(MappedBatchLoader batchLoadFunction) { + return newMappedDataLoader(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function with the provided options + * + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedDataLoader(MappedBatchLoader batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. * * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to - * know gather exact status of each item in a batch call and whether it threw exceptions when fetched then + * know the exact status of each item in a batch call and whether it threw exceptions when fetched then * you can use this form to create the data loader. * - * @param mapBatchLoaderFunction the map batch load function to use that uses {@link org.dataloader.Try} objects - * @param the key type - * @param the value type + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type * * @return a new DataLoader */ - public static DataLoader newDataLoaderWithTry(MapBatchLoader> mapBatchLoaderFunction) { - return newDataLoaderWithTry(mapBatchLoaderFunction, null); + public static DataLoader newMappedDataLoaderWithTry(MappedBatchLoader> batchLoadFunction) { + return newMappedDataLoaderWithTry(batchLoadFunction, null); } /** - * Creates new DataLoader with the specified map batch loader function and with the provided options + * Creates new DataLoader with the specified batch loader function and with the provided options * where the batch loader function returns a list of * {@link org.dataloader.Try} objects. * - * @param mapBatchLoaderFunction the map batch load function to use that uses {@link org.dataloader.Try} objects - * @param options the options to use - * @param the key type - * @param the value type + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type * * @return a new DataLoader * - * @see #newDataLoaderWithTry(MapBatchLoader) + * @see #newDataLoaderWithTry(BatchLoader) */ @SuppressWarnings("unchecked") - public static DataLoader newDataLoaderWithTry(MapBatchLoader> mapBatchLoaderFunction, DataLoaderOptions options) { - return new DataLoader<>((MapBatchLoader) mapBatchLoaderFunction, options); + public static DataLoader newMappedDataLoaderWithTry(MappedBatchLoader> batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>(batchLoadFunction, options); } /** - * Creates a new data loader with the provided batch load function, and default options. + * Creates new DataLoader with the specified mapped batch loader function and default options + * (batching, caching and unlimited batch size). * * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader */ - public DataLoader(BatchLoader batchLoadFunction) { - this(batchLoadFunction, null); + public static DataLoader newMappedDataLoader(MappedBatchLoaderWithContext batchLoadFunction) { + return newMappedDataLoader(batchLoadFunction, null); } /** - * Creates a new data loader with the provided map batch load function. + * Creates new DataLoader with the specified batch loader function with the provided options * - * @param mapBatchLoadFunction the map batch load function to use + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader */ - public DataLoader(MapBatchLoader mapBatchLoadFunction) { - this(mapBatchLoadFunction, null); + public static DataLoader newMappedDataLoader(MappedBatchLoaderWithContext batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>(batchLoadFunction, options); } /** - * Creates a new data loader with the provided batch load function and options. + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. * - * @param batchLoadFunction the batch load function to use - * @param options the batch load options + * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to + * know the exact status of each item in a batch call and whether it threw exceptions when fetched then + * you can use this form to create the data loader. + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type + * + * @return a new DataLoader */ - public DataLoader(BatchLoader batchLoadFunction, DataLoaderOptions options) { - this.batchLoadFunction = batchLoadFunction; - this.mapBatchLoadFunction = null; - this.loaderOptions = determineOptions(options); - this.futureCache = determineCacheMap(loaderOptions); - // order of keys matter in data loader - this.loaderQueue = new ArrayList<>(); - this.stats = determineCollector(this.loaderOptions); + public static DataLoader newMappedDataLoaderWithTry(MappedBatchLoaderWithContext> batchLoadFunction) { + return newMappedDataLoaderWithTry(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + * + * @see #newDataLoaderWithTry(BatchLoader) + */ + @SuppressWarnings("unchecked") + public static DataLoader newMappedDataLoaderWithTry(MappedBatchLoaderWithContext> batchLoadFunction, DataLoaderOptions options) { + return new DataLoader<>(batchLoadFunction, options); } /** - * Creates a new data loader with the provided map batch load function and options. + * Creates a new data loader with the provided batch load function, and default options. * - * @param mapBatchLoadFunction the map batch load function to use - * @param options the batch load options + * @param batchLoadFunction the batch load function to use */ - public DataLoader(MapBatchLoader mapBatchLoadFunction, DataLoaderOptions options) { - this.batchLoadFunction = null; - this.mapBatchLoadFunction = mapBatchLoadFunction; - this.loaderOptions = determineOptions(options); - this.futureCache = determineCacheMap(loaderOptions); - // order of keys matter in data loader - this.loaderQueue = new ArrayList<>(); - this.stats = determineCollector(this.loaderOptions); + public DataLoader(BatchLoader batchLoadFunction) { + this(batchLoadFunction, null); } - private StatisticsCollector determineCollector(DataLoaderOptions loaderOptions) { - return nonNull(loaderOptions.getStatisticsCollector()); + /** + * Creates a new data loader with the provided batch load function and options. + * + * @param batchLoadFunction the batch load function to use + * @param options the batch load options + */ + public DataLoader(BatchLoader batchLoadFunction, DataLoaderOptions options) { + this((Object) batchLoadFunction, options); } - private DataLoaderOptions determineOptions(DataLoaderOptions options) { - return options == null ? new DataLoaderOptions() : options; + private DataLoader(Object batchLoadFunction, DataLoaderOptions options) { + this.batchLoadFunction = nonNull(batchLoadFunction); + 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()); } @SuppressWarnings("unchecked") @@ -307,7 +406,6 @@ public CompletableFuture load(K key) { } } - /** * Requests to load the list of data provided by the specified keys asynchronously, and returns a composite future * of the resulting values. @@ -443,24 +541,20 @@ private CompletableFuture> dispatchQueueBatch(List keys, List invokeLoaderImmediately(K key) { - BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); List keys = singletonList(key); CompletionStage singleLoadCall; - if (isMapLoader()) { - singleLoadCall = mapBatchLoadFunction - .load(keys, environment) - .thenApply(map -> map.get(key)); - } else { - singleLoadCall = batchLoadFunction - .load(keys, environment) - .thenApply(list -> list.get(0)); + try { + BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); + if (isMapLoader()) { + singleLoadCall = invokeMapBatchLoader(keys, environment).thenApply(list -> list.get(0)); + } else { + singleLoadCall = invokeListBatchLoader(keys, environment).thenApply(list -> list.get(0)); + } + return singleLoadCall.toCompletableFuture(); + } catch (Exception e) { + return CompletableFutureKit.failedFuture(e); } - return singleLoadCall.toCompletableFuture(); } private CompletionStage> invokeLoader(List keys) { @@ -478,16 +572,31 @@ private CompletionStage> invokeLoader(List keys) { return batchLoad; } + @SuppressWarnings("unchecked") private CompletionStage> invokeListBatchLoader(List keys, BatchLoaderEnvironment environment) { - return nonNull(batchLoadFunction.load(keys, environment), "Your batch loader function MUST return a non null CompletionStage promise"); + CompletionStage> loadResult; + if (batchLoadFunction instanceof BatchLoaderWithContext) { + loadResult = ((BatchLoaderWithContext) batchLoadFunction).load(keys, environment); + } else { + loadResult = ((BatchLoader) batchLoadFunction).load(keys); + } + return nonNull(loadResult, "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. */ + + @SuppressWarnings("unchecked") private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderEnvironment environment) { - CompletionStage> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, environment), "Your batch loader function MUST return a non null CompletionStage promise"); + CompletionStage> loadResult; + if (batchLoadFunction instanceof MappedBatchLoaderWithContext) { + loadResult = ((MappedBatchLoaderWithContext) batchLoadFunction).load(keys, environment); + } else { + loadResult = ((MappedBatchLoader) batchLoadFunction).load(keys); + } + CompletionStage> mapBatchLoad = nonNull(loadResult, "Your batch loader function MUST return a non null CompletionStage promise"); return mapBatchLoad.thenApply(map -> { List values = new ArrayList<>(); for (K key : keys) { @@ -498,6 +607,10 @@ private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderE }); } + private boolean isMapLoader() { + return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext; + } + private void assertResultSize(List keys, List values) { assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list"); } diff --git a/src/main/java/org/dataloader/MapBatchLoader.java b/src/main/java/org/dataloader/MappedBatchLoader.java similarity index 86% rename from src/main/java/org/dataloader/MapBatchLoader.java rename to src/main/java/org/dataloader/MappedBatchLoader.java index 5513664..ee567c1 100644 --- a/src/main/java/org/dataloader/MapBatchLoader.java +++ b/src/main/java/org/dataloader/MappedBatchLoader.java @@ -57,19 +57,16 @@ * @param type parameter indicating the type of keys to use for data load requests. * @param type parameter indicating the type of values returned * - * @author Brad Baker */ -public interface MapBatchLoader { +public interface MappedBatchLoader { /** - * 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. + * Called to batch load the provided keys and return a promise to a map of values. * * @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> load(List keys, BatchLoaderEnvironment environment); + CompletionStage> load(List keys); } diff --git a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java new file mode 100644 index 0000000..d750a34 --- /dev/null +++ b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java @@ -0,0 +1,42 @@ +/* + * Copyright (c) 2016 The original author or authors + * + * All rights reserved. This program and the accompanying materials + * are made available under the terms of the Eclipse Public License v1.0 + * and Apache License v2.0 which accompanies this distribution. + * + * The Eclipse Public License is available at + * http://www.eclipse.org/legal/epl-v10.html + * + * The Apache License v2.0 is available at + * http://www.opensource.org/licenses/apache2.0.php + * + * You may elect to redistribute this code under either of these licenses. + */ + +package org.dataloader; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletionStage; + +/** + * This form of {@link MappedBatchLoader} is given a {@link org.dataloader.BatchLoaderEnvironment} object + * that encapsulates the calling context. A typical use case is passing in security credentials or database connection details + * say. + * + * See {@link MappedBatchLoader} for more details on the design invariants that you must implement in order to + * use this interface. + */ +public interface MappedBatchLoaderWithContext { + /** + * Called to batch load the provided keys and return a promise to a map of values. + * + * @param keys the collection of keys to load + * @param environment the calling environment + * + * @return a promise to a map of values for those keys + */ + @SuppressWarnings("unused") + CompletionStage> load(List keys, BatchLoaderEnvironment environment); +} diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 4f8db88..878943b 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -1,9 +1,10 @@ import org.dataloader.BatchLoader; import org.dataloader.BatchLoaderEnvironment; +import org.dataloader.BatchLoaderWithContext; import org.dataloader.CacheMap; import org.dataloader.DataLoader; import org.dataloader.DataLoaderOptions; -import org.dataloader.MapBatchLoader; +import org.dataloader.MappedBatchLoaderWithContext; import org.dataloader.Try; import org.dataloader.fixtures.SecurityCtx; import org.dataloader.fixtures.User; @@ -57,7 +58,7 @@ public CompletionStage> load(List userIds) { } }; - DataLoader userLoader = new DataLoader<>(userBatchLoader); + DataLoader userLoader = DataLoader.newDataLoader(userBatchLoader); CompletionStage load1 = userLoader.load(1L); @@ -83,17 +84,13 @@ public CompletionStage> load(List userIds) { } private void callContextExample() { - BatchLoader batchLoader = new BatchLoader() { - // - // 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> load(List keys) { - throw new UnsupportedOperationException(); - } + BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment() + .context(SecurityCtx.getCallingUserCtx()).build(); + + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment); + BatchLoaderWithContext batchLoader = new BatchLoaderWithContext() { @Override public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { SecurityCtx callCtx = environment.getContext(); @@ -101,12 +98,7 @@ public CompletionStage> load(List keys, BatchLoaderEnvironm } }; - BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment() - .context(SecurityCtx.getCallingUserCtx()).build(); - - DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment); - DataLoader loader = new DataLoader<>(batchLoader, options); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); } private CompletionStage> callDatabaseForResults(SecurityCtx callCtx, List keys) { @@ -114,7 +106,7 @@ private CompletionStage> callDatabaseForResults(SecurityCtx callCtx } private void mapBatchLoader() { - MapBatchLoader mapBatchLoader = new MapBatchLoader() { + MappedBatchLoaderWithContext mapBatchLoader = new MappedBatchLoaderWithContext() { @Override public CompletionStage> load(List userIds, BatchLoaderEnvironment environment) { SecurityCtx callCtx = environment.getContext(); @@ -122,7 +114,7 @@ public CompletionStage> load(List userIds, BatchLoaderEnvi } }; - DataLoader userLoader = DataLoader.newDataLoader(mapBatchLoader); + DataLoader userLoader = DataLoader.newMappedDataLoader(mapBatchLoader); // ... } @@ -178,7 +170,7 @@ private void clearCacheOnError() { BatchLoader userBatchLoader; private void disableCache() { - new DataLoader(userBatchLoader, DataLoaderOptions.newOptions().setCachingEnabled(false)); + DataLoader.newDataLoader(userBatchLoader, DataLoaderOptions.newOptions().setCachingEnabled(false)); userDataLoader.load("A"); @@ -221,7 +213,7 @@ private void customCache() { MyCustomCache customCache = new MyCustomCache(); DataLoaderOptions options = DataLoaderOptions.newOptions().setCacheMap(customCache); - new DataLoader(userBatchLoader, options); + DataLoader.newDataLoader(userBatchLoader, options); } private void processUser(User user) { diff --git a/src/test/java/TestFI.java b/src/test/java/TestFI.java new file mode 100644 index 0000000..eeafadb --- /dev/null +++ b/src/test/java/TestFI.java @@ -0,0 +1,63 @@ +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static java.util.Collections.emptyMap; + +public class TestFI { + + interface BL { + List load(List keys); + } + + interface MapBL { + Map load(List keys); + } + + interface BLwithContext { + List load(List keys, Object context); + } + + interface MapBLwithContext { + Map load(List keys, Object context); + } + + + static class DL { + public DL(BL loader) { + } + + public DL(BLwithContext loader) { + } + + public DL(MapBLwithContext loader) { + } + + + public static DL newDLWithMap(MapBLwithContext loader) { + return null; + } + + public static DL newDL(BLwithContext loader) { + return null; + } + } + + + public static void main(String[] args) { + DL dl = new DL<>(keys -> keys.stream() + .map(Object::toString) + .collect(Collectors.toList())); + + BLwithContext integerStringBL = (keys, env) -> keys.stream() + .map(Object::toString) + .collect(Collectors.toList()); + DL dl2 = new DL(integerStringBL); + + MapBLwithContext mapBLwithContext = (keys, env) -> emptyMap(); + dl = new DL<>(mapBLwithContext); + + DL.newDLWithMap(((keys, context) -> emptyMap())); + } +} diff --git a/src/test/java/org/dataloader/CustomCacheMap.java b/src/test/java/org/dataloader/CustomCacheMap.java new file mode 100644 index 0000000..505148d --- /dev/null +++ b/src/test/java/org/dataloader/CustomCacheMap.java @@ -0,0 +1,41 @@ +package org.dataloader; + +import java.util.LinkedHashMap; +import java.util.Map; + +public class CustomCacheMap implements CacheMap { + + public Map stash; + + public CustomCacheMap() { + stash = new LinkedHashMap<>(); + } + + @Override + public boolean containsKey(String key) { + return stash.containsKey(key); + } + + @Override + public Object get(String key) { + return stash.get(key); + } + + @Override + public CacheMap set(String key, Object value) { + stash.put(key, value); + return this; + } + + @Override + public CacheMap delete(String key) { + stash.remove(key); + return this; + } + + @Override + public CacheMap clear() { + stash.clear(); + return this; + } +} diff --git a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java index 236541c..88c2a4e 100644 --- a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java +++ b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java @@ -6,7 +6,6 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionStage; import java.util.stream.Collectors; import static java.util.Arrays.asList; @@ -21,21 +20,13 @@ public class DataLoaderBatchLoaderEnvironmentTest { @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(); - } - - @Override - public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { - List list = keys.stream().map(k -> k + "-" + environment.getContext()).collect(Collectors.toList()); - return CompletableFuture.completedFuture(list); - } + BatchLoaderWithContext batchLoader = (keys, environment) -> { + List list = keys.stream().map(k -> k + "-" + environment.getContext()).collect(Collectors.toList()); + return CompletableFuture.completedFuture(list); }; DataLoaderOptions options = DataLoaderOptions.newOptions() .setBatchLoaderEnvironmentProvider(() -> newBatchLoaderEnvironment().context("ctx").build()); - DataLoader loader = new DataLoader<>(batchLoader, options); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); loader.load("A"); loader.load("B"); @@ -48,14 +39,14 @@ public CompletionStage> load(List keys, BatchLoaderEnvironm @Test public void context_is_passed_to_map_batch_loader_function() throws Exception { - MapBatchLoader mapBatchLoader = (keys, environment) -> { + MappedBatchLoaderWithContext mapBatchLoader = (keys, environment) -> { Map map = new HashMap<>(); keys.forEach(k -> map.put(k, k + "-" + environment.getContext())); return CompletableFuture.completedFuture(map); }; DataLoaderOptions options = DataLoaderOptions.newOptions() .setBatchLoaderEnvironmentProvider(() -> newBatchLoaderEnvironment().context("ctx").build()); - DataLoader loader = new DataLoader<>(mapBatchLoader, options); + DataLoader loader = DataLoader.newMappedDataLoader(mapBatchLoader, options); loader.load("A"); loader.load("B"); @@ -68,19 +59,11 @@ public void context_is_passed_to_map_batch_loader_function() throws Exception { @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, BatchLoaderEnvironment environment) { - List list = keys.stream().map(k -> k + "-" + environment.getContext()).collect(Collectors.toList()); - return CompletableFuture.completedFuture(list); - } + BatchLoaderWithContext batchLoader = (keys, environment) -> { + List list = keys.stream().map(k -> k + "-" + environment.getContext()).collect(Collectors.toList()); + return CompletableFuture.completedFuture(list); }; - DataLoader loader = new DataLoader<>(batchLoader); + DataLoader loader = DataLoader.newDataLoader(batchLoader); loader.load("A"); loader.load("B"); @@ -93,12 +76,12 @@ public CompletionStage> load(List keys, BatchLoaderEnvironm @Test public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() throws Exception { - MapBatchLoader mapBatchLoader = (keys, environment) -> { + MappedBatchLoaderWithContext mapBatchLoader = (keys, environment) -> { Map map = new HashMap<>(); keys.forEach(k -> map.put(k, k + "-" + environment.getContext())); return CompletableFuture.completedFuture(map); }; - DataLoader loader = new DataLoader<>(mapBatchLoader); + DataLoader loader = DataLoader.newMappedDataLoader(mapBatchLoader); loader.load("A"); loader.load("B"); diff --git a/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java b/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java index f6aa397..8e83f82 100644 --- a/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java +++ b/src/test/java/org/dataloader/DataLoaderMapBatchLoaderTest.java @@ -23,7 +23,7 @@ import static org.junit.Assert.assertThat; /** - * Much of the tests that related to {@link org.dataloader.MapBatchLoader} also related to + * Much of the tests that related to {@link MappedBatchLoader} also related to * {@link org.dataloader.BatchLoader}. This is white box testing somewhat because we could have repeated * ALL the tests in {@link org.dataloader.DataLoaderTest} here as well but chose not to because we KNOW that * DataLoader differs only a little in how it handles the 2 types of loader functions. We choose to grab some @@ -31,7 +31,7 @@ */ public class DataLoaderMapBatchLoaderTest { - MapBatchLoader evensOnlyMapBatchLoader = (keys, environment) -> { + MappedBatchLoader evensOnlyMappedBatchLoader = (keys) -> { Map mapOfResults = new HashMap<>(); for (int i = 0; i < keys.size(); i++) { String k = keys.get(i); @@ -43,19 +43,19 @@ public class DataLoaderMapBatchLoaderTest { }; private static DataLoader idMapLoader(DataLoaderOptions options, List> loadCalls) { - MapBatchLoader kvBatchLoader = (keys, environment) -> { + MappedBatchLoader kvBatchLoader = (keys) -> { loadCalls.add(new ArrayList<>(keys)); Map map = new HashMap<>(); //noinspection unchecked keys.forEach(k -> map.put(k, (V) k)); return CompletableFuture.completedFuture(map); }; - return DataLoader.newDataLoader(kvBatchLoader, options); + return DataLoader.newMappedDataLoader(kvBatchLoader, options); } private static DataLoader idMapLoaderBlowsUps( DataLoaderOptions options, List> loadCalls) { - return new DataLoader<>((keys, environment) -> { + return new DataLoader<>((keys) -> { loadCalls.add(new ArrayList<>(keys)); return futureError(); }, options); @@ -64,7 +64,7 @@ private static DataLoader idMapLoaderBlowsUps( @Test public void basic_map_batch_loading() throws Exception { - DataLoader loader = new DataLoader<>(evensOnlyMapBatchLoader); + DataLoader loader = DataLoader.newMappedDataLoader(evensOnlyMappedBatchLoader); loader.load("A"); loader.load("B"); @@ -175,5 +175,4 @@ public void should_work_with_duplicate_keys_when_caching_enabled() throws Execut } - } diff --git a/src/test/java/org/dataloader/DataLoaderTest.java b/src/test/java/org/dataloader/DataLoaderTest.java index 77bcc3e..91f8f02 100644 --- a/src/test/java/org/dataloader/DataLoaderTest.java +++ b/src/test/java/org/dataloader/DataLoaderTest.java @@ -23,9 +23,7 @@ import java.util.ArrayList; import java.util.Collection; -import java.util.LinkedHashMap; import java.util.List; -import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.concurrent.ExecutionException; @@ -1029,41 +1027,5 @@ private BatchLoader keysAsValues() { return CompletableFuture::completedFuture; } - public class CustomCacheMap implements CacheMap { - - public Map stash; - - public CustomCacheMap() { - stash = new LinkedHashMap<>(); - } - - @Override - public boolean containsKey(String key) { - return stash.containsKey(key); - } - - @Override - public Object get(String key) { - return stash.get(key); - } - - @Override - public CacheMap set(String key, Object value) { - stash.put(key, value); - return this; - } - - @Override - public CacheMap delete(String key) { - stash.remove(key); - return this; - } - - @Override - public CacheMap clear() { - stash.clear(); - return this; - } - } } diff --git a/src/test/java/org/dataloader/DataLoaderWithTryTest.java b/src/test/java/org/dataloader/DataLoaderWithTryTest.java index ee6c459..b926841 100644 --- a/src/test/java/org/dataloader/DataLoaderWithTryTest.java +++ b/src/test/java/org/dataloader/DataLoaderWithTryTest.java @@ -41,10 +41,10 @@ public void should_handle_Trys_coming_back_from_batchLoader() throws Exception { } @Test - public void should_handle_Trys_coming_back_from_map_batchLoader() throws Exception { + public void should_handle_Trys_coming_back_from_mapped_batchLoader() throws Exception { List> batchKeyCalls = new ArrayList<>(); - MapBatchLoader> batchLoader = (keys, environment) -> { + MappedBatchLoaderWithContext> batchLoader = (keys, environment) -> { batchKeyCalls.add(keys); Map> result = new HashMap<>(); @@ -58,7 +58,7 @@ public void should_handle_Trys_coming_back_from_map_batchLoader() throws Excepti return CompletableFuture.completedFuture(result); }; - DataLoader dataLoader = DataLoader.newDataLoaderWithTry(batchLoader); + DataLoader dataLoader = DataLoader.newMappedDataLoaderWithTry(batchLoader); commonTryAsserts(batchKeyCalls, dataLoader); } From a8a18f7c6c709eb3e56fbcc2522e0a8eff926f50 Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sat, 25 Aug 2018 20:09:45 +1000 Subject: [PATCH 06/11] Split dataloader into a helper implementation for most of the GUTS of the code --- src/main/java/org/dataloader/DataLoader.java | 221 +-------------- .../java/org/dataloader/DataLoaderHelper.java | 262 ++++++++++++++++++ 2 files changed, 268 insertions(+), 215 deletions(-) create mode 100644 src/main/java/org/dataloader/DataLoaderHelper.java diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index d97bcbd..24edd76 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -22,16 +22,10 @@ import java.util.AbstractMap.SimpleImmutableEntry; import java.util.ArrayList; -import java.util.Collection; import java.util.List; -import java.util.Map; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionStage; import java.util.stream.Collectors; -import static java.util.Collections.emptyList; -import static java.util.Collections.singletonList; -import static org.dataloader.impl.Assertions.assertState; import static org.dataloader.impl.Assertions.nonNull; /** @@ -63,7 +57,7 @@ */ public class DataLoader { - private final Object batchLoadFunction; + private final DataLoaderHelper helper; private final DataLoaderOptions loaderOptions; private final CacheMap> futureCache; private final List>> loaderQueue; @@ -352,12 +346,13 @@ public DataLoader(BatchLoader batchLoadFunction, DataLoaderOptions options } private DataLoader(Object batchLoadFunction, DataLoaderOptions options) { - this.batchLoadFunction = nonNull(batchLoadFunction); 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, futureCache, loaderQueue, stats); } @SuppressWarnings("unchecked") @@ -377,33 +372,7 @@ private CacheMap> determineCacheMap(DataLoaderOptio * @return the future of the value */ public CompletableFuture load(K key) { - synchronized (this) { - Object cacheKey = getCacheKey(nonNull(key)); - stats.incrementLoadCount(); - - boolean batchingEnabled = loaderOptions.batchingEnabled(); - boolean cachingEnabled = loaderOptions.cachingEnabled(); - - if (cachingEnabled) { - if (futureCache.containsKey(cacheKey)) { - stats.incrementCacheHitCount(); - return futureCache.get(cacheKey); - } - } - - CompletableFuture future = new CompletableFuture<>(); - if (batchingEnabled) { - loaderQueue.add(new SimpleImmutableEntry<>(key, future)); - } else { - stats.incrementBatchLoadCountBy(1); - // immediate execution of batch function - future = invokeLoaderImmediately(key); - } - if (cachingEnabled) { - futureCache.set(cacheKey, future); - } - return future; - } + return helper.load(key); } /** @@ -436,183 +405,7 @@ public CompletableFuture> loadMany(List keys) { * @return the promise of the queued load requests */ public CompletableFuture> dispatch() { - boolean batchingEnabled = loaderOptions.batchingEnabled(); - // - // we copy the pre-loaded set of futures ready for dispatch - final List keys = new ArrayList<>(); - final List> queuedFutures = new ArrayList<>(); - synchronized (this) { - loaderQueue.forEach(entry -> { - keys.add(entry.getKey()); - queuedFutures.add(entry.getValue()); - }); - loaderQueue.clear(); - } - if (!batchingEnabled || keys.size() == 0) { - return CompletableFuture.completedFuture(emptyList()); - } - // - // order of keys -> values matter in data loader hence the use of linked hash map - // - // See https://github.com/facebook/dataloader/blob/master/README.md for more details - // - - // - // when the promised list of values completes, we transfer the values into - // the previously cached future objects that the client already has been given - // via calls to load("foo") and loadMany(["foo","bar"]) - // - int maxBatchSize = loaderOptions.maxBatchSize(); - if (maxBatchSize > 0 && maxBatchSize < keys.size()) { - return sliceIntoBatchesOfBatches(keys, queuedFutures, maxBatchSize); - } else { - return dispatchQueueBatch(keys, queuedFutures); - } - } - - private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List> queuedFutures, int maxBatchSize) { - // the number of keys is > than what the batch loader function can accept - // so make multiple calls to the loader - List>> allBatches = new ArrayList<>(); - int len = keys.size(); - int batchCount = (int) Math.ceil(len / (double) maxBatchSize); - for (int i = 0; i < batchCount; i++) { - - int fromIndex = i * maxBatchSize; - int toIndex = Math.min((i + 1) * maxBatchSize, len); - - List subKeys = keys.subList(fromIndex, toIndex); - List> subFutures = queuedFutures.subList(fromIndex, toIndex); - - allBatches.add(dispatchQueueBatch(subKeys, subFutures)); - } - // - // now reassemble all the futures into one that is the complete set of results - return CompletableFuture.allOf(allBatches.toArray(new CompletableFuture[allBatches.size()])) - .thenApply(v -> allBatches.stream() - .map(CompletableFuture::join) - .flatMap(Collection::stream) - .collect(Collectors.toList())); - } - - @SuppressWarnings("unchecked") - private CompletableFuture> dispatchQueueBatch(List keys, List> queuedFutures) { - stats.incrementBatchLoadCountBy(keys.size()); - CompletionStage> batchLoad = invokeLoader(keys); - return batchLoad - .toCompletableFuture() - .thenApply(values -> { - assertResultSize(keys, values); - - for (int idx = 0; idx < queuedFutures.size(); idx++) { - Object value = values.get(idx); - CompletableFuture future = queuedFutures.get(idx); - if (value instanceof Throwable) { - stats.incrementLoadErrorCount(); - future.completeExceptionally((Throwable) value); - // we don't clear the cached view of this entry to avoid - // frequently loading the same error - } else if (value instanceof Try) { - // we allow the batch loader to return a Try so we can better represent a computation - // that might have worked or not. - Try tryValue = (Try) value; - if (tryValue.isSuccess()) { - future.complete(tryValue.get()); - } else { - stats.incrementLoadErrorCount(); - future.completeExceptionally(tryValue.getThrowable()); - } - } else { - V val = (V) value; - future.complete(val); - } - } - return values; - }).exceptionally(ex -> { - stats.incrementBatchLoadExceptionCount(); - for (int idx = 0; idx < queuedFutures.size(); idx++) { - K key = keys.get(idx); - CompletableFuture future = queuedFutures.get(idx); - future.completeExceptionally(ex); - // clear any cached view of this key because they all failed - clear(key); - } - return emptyList(); - }); - } - - private CompletableFuture invokeLoaderImmediately(K key) { - List keys = singletonList(key); - CompletionStage singleLoadCall; - try { - BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); - if (isMapLoader()) { - singleLoadCall = invokeMapBatchLoader(keys, environment).thenApply(list -> list.get(0)); - } else { - singleLoadCall = invokeListBatchLoader(keys, environment).thenApply(list -> list.get(0)); - } - return singleLoadCall.toCompletableFuture(); - } catch (Exception e) { - return CompletableFutureKit.failedFuture(e); - } - } - - private CompletionStage> invokeLoader(List keys) { - CompletionStage> batchLoad; - try { - BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); - if (isMapLoader()) { - batchLoad = invokeMapBatchLoader(keys, environment); - } else { - batchLoad = invokeListBatchLoader(keys, environment); - } - } catch (Exception e) { - batchLoad = CompletableFutureKit.failedFuture(e); - } - return batchLoad; - } - - @SuppressWarnings("unchecked") - private CompletionStage> invokeListBatchLoader(List keys, BatchLoaderEnvironment environment) { - CompletionStage> loadResult; - if (batchLoadFunction instanceof BatchLoaderWithContext) { - loadResult = ((BatchLoaderWithContext) batchLoadFunction).load(keys, environment); - } else { - loadResult = ((BatchLoader) batchLoadFunction).load(keys); - } - return nonNull(loadResult, "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. - */ - - @SuppressWarnings("unchecked") - private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderEnvironment environment) { - CompletionStage> loadResult; - if (batchLoadFunction instanceof MappedBatchLoaderWithContext) { - loadResult = ((MappedBatchLoaderWithContext) batchLoadFunction).load(keys, environment); - } else { - loadResult = ((MappedBatchLoader) batchLoadFunction).load(keys); - } - CompletionStage> mapBatchLoad = nonNull(loadResult, "Your batch loader function MUST return a non null CompletionStage promise"); - return mapBatchLoad.thenApply(map -> { - List values = new ArrayList<>(); - for (K key : keys) { - V value = map.get(key); - values.add(value); - } - return values; - }); - } - - private boolean isMapLoader() { - return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext; - } - - private void assertResultSize(List keys, List values) { - assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list"); + return helper.dispatch(); } /** @@ -718,10 +511,8 @@ public DataLoader prime(K key, Exception error) { * * @return the cache key after the input is transformed with the cache key function */ - @SuppressWarnings("unchecked") public Object getCacheKey(K key) { - return loaderOptions.cacheKeyFunction().isPresent() ? - loaderOptions.cacheKeyFunction().get().getKey(key) : key; + return helper.getCacheKey(key); } /** diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java new file mode 100644 index 0000000..87e1ebc --- /dev/null +++ b/src/main/java/org/dataloader/DataLoaderHelper.java @@ -0,0 +1,262 @@ +package org.dataloader; + +import org.dataloader.impl.CompletableFutureKit; +import org.dataloader.stats.StatisticsCollector; + +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.stream.Collectors; + +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; +import static org.dataloader.impl.Assertions.assertState; +import static org.dataloader.impl.Assertions.nonNull; + +/** + * This helps break up the large DataLoader class functionality and it contains the logic to dispatch the + * promises on behalf of its peer dataloader + * + * @param the type of keys + * @param the type of values + */ +class DataLoaderHelper { + + private final DataLoader dataLoader; + private final Object batchLoadFunction; + private final DataLoaderOptions loaderOptions; + private final CacheMap> futureCache; + private final List>> loaderQueue; + private final StatisticsCollector stats; + + DataLoaderHelper(DataLoader dataLoader, Object batchLoadFunction, DataLoaderOptions loaderOptions, CacheMap> futureCache, List>> loaderQueue, StatisticsCollector stats) { + this.dataLoader = dataLoader; + this.batchLoadFunction = batchLoadFunction; + this.loaderOptions = loaderOptions; + this.futureCache = futureCache; + this.loaderQueue = loaderQueue; + this.stats = stats; + } + + CompletableFuture load(K key) { + synchronized (dataLoader) { + Object cacheKey = getCacheKey(nonNull(key)); + stats.incrementLoadCount(); + + boolean batchingEnabled = loaderOptions.batchingEnabled(); + boolean cachingEnabled = loaderOptions.cachingEnabled(); + + if (cachingEnabled) { + if (futureCache.containsKey(cacheKey)) { + stats.incrementCacheHitCount(); + return futureCache.get(cacheKey); + } + } + + CompletableFuture future = new CompletableFuture<>(); + if (batchingEnabled) { + loaderQueue.add(new AbstractMap.SimpleImmutableEntry<>(key, future)); + } else { + stats.incrementBatchLoadCountBy(1); + // immediate execution of batch function + future = invokeLoaderImmediately(key); + } + if (cachingEnabled) { + futureCache.set(cacheKey, future); + } + return future; + } + } + + @SuppressWarnings("unchecked") + Object getCacheKey(K key) { + return loaderOptions.cacheKeyFunction().isPresent() ? + loaderOptions.cacheKeyFunction().get().getKey(key) : key; + } + + CompletableFuture> dispatch() { + boolean batchingEnabled = loaderOptions.batchingEnabled(); + // + // we copy the pre-loaded set of futures ready for dispatch + final List keys = new ArrayList<>(); + final List> queuedFutures = new ArrayList<>(); + synchronized (dataLoader) { + loaderQueue.forEach(entry -> { + keys.add(entry.getKey()); + queuedFutures.add(entry.getValue()); + }); + loaderQueue.clear(); + } + if (!batchingEnabled || keys.size() == 0) { + return CompletableFuture.completedFuture(emptyList()); + } + // + // order of keys -> values matter in data loader hence the use of linked hash map + // + // See https://github.com/facebook/dataloader/blob/master/README.md for more details + // + + // + // when the promised list of values completes, we transfer the values into + // the previously cached future objects that the client already has been given + // via calls to load("foo") and loadMany(["foo","bar"]) + // + int maxBatchSize = loaderOptions.maxBatchSize(); + if (maxBatchSize > 0 && maxBatchSize < keys.size()) { + return sliceIntoBatchesOfBatches(keys, queuedFutures, maxBatchSize); + } else { + return dispatchQueueBatch(keys, queuedFutures); + } + } + + private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List> queuedFutures, int maxBatchSize) { + // the number of keys is > than what the batch loader function can accept + // so make multiple calls to the loader + List>> allBatches = new ArrayList<>(); + int len = keys.size(); + int batchCount = (int) Math.ceil(len / (double) maxBatchSize); + for (int i = 0; i < batchCount; i++) { + + int fromIndex = i * maxBatchSize; + int toIndex = Math.min((i + 1) * maxBatchSize, len); + + List subKeys = keys.subList(fromIndex, toIndex); + List> subFutures = queuedFutures.subList(fromIndex, toIndex); + + allBatches.add(dispatchQueueBatch(subKeys, subFutures)); + } + // + // now reassemble all the futures into one that is the complete set of results + return CompletableFuture.allOf(allBatches.toArray(new CompletableFuture[allBatches.size()])) + .thenApply(v -> allBatches.stream() + .map(CompletableFuture::join) + .flatMap(Collection::stream) + .collect(Collectors.toList())); + } + + @SuppressWarnings("unchecked") + private CompletableFuture> dispatchQueueBatch(List keys, List> queuedFutures) { + stats.incrementBatchLoadCountBy(keys.size()); + CompletionStage> batchLoad = invokeLoader(keys); + return batchLoad + .toCompletableFuture() + .thenApply(values -> { + assertResultSize(keys, values); + + for (int idx = 0; idx < queuedFutures.size(); idx++) { + Object value = values.get(idx); + CompletableFuture future = queuedFutures.get(idx); + if (value instanceof Throwable) { + stats.incrementLoadErrorCount(); + future.completeExceptionally((Throwable) value); + // we don't clear the cached view of this entry to avoid + // frequently loading the same error + } else if (value instanceof Try) { + // we allow the batch loader to return a Try so we can better represent a computation + // that might have worked or not. + Try tryValue = (Try) value; + if (tryValue.isSuccess()) { + future.complete(tryValue.get()); + } else { + stats.incrementLoadErrorCount(); + future.completeExceptionally(tryValue.getThrowable()); + } + } else { + V val = (V) value; + future.complete(val); + } + } + return values; + }).exceptionally(ex -> { + stats.incrementBatchLoadExceptionCount(); + for (int idx = 0; idx < queuedFutures.size(); idx++) { + K key = keys.get(idx); + CompletableFuture future = queuedFutures.get(idx); + future.completeExceptionally(ex); + // clear any cached view of this key because they all failed + dataLoader.clear(key); + } + return emptyList(); + }); + } + + + private void assertResultSize(List keys, List values) { + assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list"); + } + + + CompletableFuture invokeLoaderImmediately(K key) { + List keys = singletonList(key); + CompletionStage singleLoadCall; + try { + BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); + if (isMapLoader()) { + singleLoadCall = invokeMapBatchLoader(keys, environment).thenApply(list -> list.get(0)); + } else { + singleLoadCall = invokeListBatchLoader(keys, environment).thenApply(list -> list.get(0)); + } + return singleLoadCall.toCompletableFuture(); + } catch (Exception e) { + return CompletableFutureKit.failedFuture(e); + } + } + + CompletionStage> invokeLoader(List keys) { + CompletionStage> batchLoad; + try { + BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); + if (isMapLoader()) { + batchLoad = invokeMapBatchLoader(keys, environment); + } else { + batchLoad = invokeListBatchLoader(keys, environment); + } + } catch (Exception e) { + batchLoad = CompletableFutureKit.failedFuture(e); + } + return batchLoad; + } + + @SuppressWarnings("unchecked") + private CompletionStage> invokeListBatchLoader(List keys, BatchLoaderEnvironment environment) { + CompletionStage> loadResult; + if (batchLoadFunction instanceof BatchLoaderWithContext) { + loadResult = ((BatchLoaderWithContext) batchLoadFunction).load(keys, environment); + } else { + loadResult = ((BatchLoader) batchLoadFunction).load(keys); + } + return nonNull(loadResult, "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. + */ + + @SuppressWarnings("unchecked") + private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderEnvironment environment) { + CompletionStage> loadResult; + if (batchLoadFunction instanceof MappedBatchLoaderWithContext) { + loadResult = ((MappedBatchLoaderWithContext) batchLoadFunction).load(keys, environment); + } else { + loadResult = ((MappedBatchLoader) batchLoadFunction).load(keys); + } + CompletionStage> mapBatchLoad = nonNull(loadResult, "Your batch loader function MUST return a non null CompletionStage promise"); + return mapBatchLoad.thenApply(map -> { + List values = new ArrayList<>(); + for (K key : keys) { + V value = map.get(key); + values.add(value); + } + return values; + }); + } + + private boolean isMapLoader() { + return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext; + } +} From d5fbf52d3b33803d198a419364343ce42ff38a5a Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 26 Aug 2018 12:07:52 +1000 Subject: [PATCH 07/11] Spilling misteaks and gradma changes in doco --- README.md | 8 +++--- .../dataloader/BatchLoaderEnvironment.java | 2 +- .../dataloader/BatchLoaderWithContext.java | 2 +- src/main/java/org/dataloader/DataLoader.java | 28 ++++++++++++------- .../org/dataloader/MappedBatchLoader.java | 10 +++---- .../MappedBatchLoaderWithContext.java | 4 +-- 6 files changed, 31 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index d3ce756..6e9d268 100644 --- a/README.md +++ b/README.md @@ -204,15 +204,15 @@ For example, let's assume you want to load users from a database, you could prob ``` Given say 10 user id keys you might only get 7 results back. This can be more naturally represented in a map - than in an order list of values when returning values from the batch loader function. + than in an ordered list of values from the batch loader function. You can use `org.dataloader.MappedBatchLoader` for this purpose. When the map is processed by the `DataLoader` code, any keys that are missing in the map - will be replaced with null values. The semantics that the number of `DataLoader.load` requests - are matched with values is kept. + will be replaced with null values. The semantic that the number of `DataLoader.load` requests + are matched with an equal number of values is kept. - Your keys provided MUST be first class keys since they will be used to examine the returned map and + The keys provided MUST be first class keys since they will be used to examine the returned map and create the list of results, with nulls filling in for missing values. ```java diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironment.java b/src/main/java/org/dataloader/BatchLoaderEnvironment.java index 969c7c6..c6d0126 100644 --- a/src/main/java/org/dataloader/BatchLoaderEnvironment.java +++ b/src/main/java/org/dataloader/BatchLoaderEnvironment.java @@ -2,7 +2,7 @@ /** * 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. + * of the calling users for example or database parameters that allow the data layer call to succeed. */ public class BatchLoaderEnvironment { diff --git a/src/main/java/org/dataloader/BatchLoaderWithContext.java b/src/main/java/org/dataloader/BatchLoaderWithContext.java index 1ee8f66..e287d90 100644 --- a/src/main/java/org/dataloader/BatchLoaderWithContext.java +++ b/src/main/java/org/dataloader/BatchLoaderWithContext.java @@ -15,7 +15,7 @@ public interface BatchLoaderWithContext { /** * Called to batch load the provided keys and return a promise to a list of values. This default * 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. + * is passing in security credentials or database details for example. * * @param keys the collection of keys to load * @param environment an environment object that can help with the call diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 24edd76..5420bb4 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -57,7 +57,7 @@ */ public class DataLoader { - private final DataLoaderHelper helper; + private final DataLoaderHelper helper; private final DataLoaderOptions loaderOptions; private final CacheMap> futureCache; private final List>> loaderQueue; @@ -96,10 +96,12 @@ public static DataLoader newDataLoader(BatchLoader batchLoadF * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. * - * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to - * know the exact status of each item in a batch call and whether it threw exceptions when fetched then + * If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. * + * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects * @param the key type * @param the value type @@ -162,10 +164,12 @@ public static DataLoader newDataLoader(BatchLoaderWithContext * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. * - * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to - * know the exact status of each item in a batch call and whether it threw exceptions when fetched then + * If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. * + * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects * @param the key type * @param the value type @@ -227,10 +231,12 @@ public static DataLoader newMappedDataLoader(MappedBatchLoader the key type * @param the value type @@ -293,10 +299,12 @@ public static DataLoader newMappedDataLoader(MappedBatchLoaderWithC * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. * - * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value. If its important you to - * know the exact status of each item in a batch call and whether it threw exceptions when fetched then + * If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. * + * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects * @param the key type * @param the value type @@ -352,7 +360,7 @@ private DataLoader(Object batchLoadFunction, DataLoaderOptions options) { this.loaderQueue = new ArrayList<>(); this.stats = nonNull(this.loaderOptions.getStatisticsCollector()); - this.helper = new DataLoaderHelper<>(this, batchLoadFunction,this.loaderOptions, futureCache, loaderQueue, stats); + this.helper = new DataLoaderHelper<>(this, batchLoadFunction, this.loaderOptions, futureCache, loaderQueue, stats); } @SuppressWarnings("unchecked") diff --git a/src/main/java/org/dataloader/MappedBatchLoader.java b/src/main/java/org/dataloader/MappedBatchLoader.java index ee567c1..16206b6 100644 --- a/src/main/java/org/dataloader/MappedBatchLoader.java +++ b/src/main/java/org/dataloader/MappedBatchLoader.java @@ -42,16 +42,16 @@ * * * Given say 10 user id keys you might only get 7 results back. This can be more naturally represented in a map - * than in an order list of values when returning values from the batch loader function. + * than in an ordered list of values from the batch loader function. * * When the map is processed by the {@link org.dataloader.DataLoader} code, any keys that are missing in the map - * will be replaced with null values. The semantics that the number of {@link org.dataloader.DataLoader#load(Object)} requests - * are matched with values is kept. + * will be replaced with null values. The semantic that the number of {@link org.dataloader.DataLoader#load(Object)} requests + * are matched with and equal number of values is kept. * - * This means that if 10 keys are asked for then {@link DataLoader#dispatch()} will return a promise 10 value results and each + * This means that if 10 keys are asked for then {@link DataLoader#dispatch()} will return a promise of 10 value results and each * of the {@link org.dataloader.DataLoader#load(Object)} will complete with a value, null or an exception. * - * When caching is disabled, its possible for the same key to be presented in the list of keys more than once. You map + * When caching is disabled, its possible for the same key to be presented in the list of keys more than once. Your map * batch loader function needs to be resilient to this situation. * * @param type parameter indicating the type of keys to use for data load requests. diff --git a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java index d750a34..0092563 100644 --- a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java +++ b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java @@ -22,8 +22,8 @@ /** * This form of {@link MappedBatchLoader} is given a {@link org.dataloader.BatchLoaderEnvironment} object - * that encapsulates the calling context. A typical use case is passing in security credentials or database connection details - * say. + * that encapsulates the calling context. A typical use case is passing in security credentials or database details + * for example. * * See {@link MappedBatchLoader} for more details on the design invariants that you must implement in order to * use this interface. From 376ae55cec36bb0d852b4d716493e1062ad9e63c Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 26 Aug 2018 12:11:57 +1000 Subject: [PATCH 08/11] Javadoc p fix ups --- src/main/java/org/dataloader/BatchLoader.java | 12 +++++------- src/main/java/org/dataloader/DataLoader.java | 16 ++++++++-------- .../java/org/dataloader/MappedBatchLoader.java | 14 +++++++------- .../dataloader/MappedBatchLoaderWithContext.java | 2 +- .../java/org/dataloader/impl/PromisedValues.java | 4 ++-- .../stats/ThreadLocalStatisticsCollector.java | 4 ++-- 6 files changed, 25 insertions(+), 27 deletions(-) diff --git a/src/main/java/org/dataloader/BatchLoader.java b/src/main/java/org/dataloader/BatchLoader.java index 463cac6..6780f56 100644 --- a/src/main/java/org/dataloader/BatchLoader.java +++ b/src/main/java/org/dataloader/BatchLoader.java @@ -16,21 +16,19 @@ package org.dataloader; -import java.util.Collections; import java.util.List; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; /** * A function that is invoked for batch loading a list of data values indicated by the provided list of keys. The * function returns a promise of a list of results of individual load requests. - * + *

* There are a few constraints that must be upheld: *

    *
  • The list of values must be the same size as the list of keys.
    • *
  • Each index in the list of values must correspond to the same index in the list of keys.
    • *
- * + *

* For example, if your batch function was provided the list of keys: * * 

@@ -50,10 +48,10 @@
  *
* * then the batch loader function contract has been broken. - * + *

* The back-end service returned results in a different order than we requested, likely because it was more efficient for it to * do so. Also, it omitted a result for key 6, which we may interpret as no value existing for that key. - * + *

* To uphold the constraints of the batch function, it must return an List of values the same length as * the List of keys, and re-order them to ensure each index aligns with the original keys [ 2, 9, 6, 1 ]: * @@ -77,7 +75,7 @@ public interface BatchLoader { /** * Called to batch load the provided keys and return a promise to a list of values. - * + *

* If you need calling context then implement {@link org.dataloader.BatchLoaderWithContext} * * @param keys the collection of keys to load diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 5420bb4..4cc7c57 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -95,10 +95,10 @@ public static DataLoader newDataLoader(BatchLoader batchLoadF * Creates new DataLoader with the specified batch loader function and default options * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. - * + *

* If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. - * + *

* Using Try objects allows you to capture a value returned or an exception that might * have occurred trying to get a value. . * @@ -163,10 +163,10 @@ public static DataLoader newDataLoader(BatchLoaderWithContext * Creates new DataLoader with the specified batch loader function and default options * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. - * + *

* If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. - * + *

* Using Try objects allows you to capture a value returned or an exception that might * have occurred trying to get a value. . * @@ -230,13 +230,13 @@ public static DataLoader newMappedDataLoader(MappedBatchLoader * If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. * * Using Try objects allows you to capture a value returned or an exception that might * have occurred trying to get a value. . - * + *

* @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects * @param the key type * @param the value type @@ -298,10 +298,10 @@ public static DataLoader newMappedDataLoader(MappedBatchLoaderWithC * Creates new DataLoader with the specified batch loader function and default options * (batching, caching and unlimited batch size) where the batch loader function returns a list of * {@link org.dataloader.Try} objects. - * + *

* If its important you to know the exact status of each item in a batch call and whether it threw exceptions then * you can use this form to create the data loader. - * + *

* Using Try objects allows you to capture a value returned or an exception that might * have occurred trying to get a value. . * diff --git a/src/main/java/org/dataloader/MappedBatchLoader.java b/src/main/java/org/dataloader/MappedBatchLoader.java index 16206b6..b230e7f 100644 --- a/src/main/java/org/dataloader/MappedBatchLoader.java +++ b/src/main/java/org/dataloader/MappedBatchLoader.java @@ -23,7 +23,7 @@ /** * A function that is invoked for batch loading a map of of data values indicated by the provided list of keys. The * function returns a promise of a map of results of individual load requests. - * + *

* There are a few constraints that must be upheld: *

    *
  • The keys MUST be able to be first class keys in a Java map. Get your equals() and hashCode() methods in order
    • @@ -32,25 +32,25 @@ * *
  • The function MUST be resilient to the same key being presented twice.
    • *
- * + *

* This form is useful when you don't have a 1:1 mapping of keys to values or when null is an acceptable value for a missing value. - * + *

* For example, let's assume you want to load users from a database, you could probably use a query that looks like this: * * 

  *    SELECT * FROM User WHERE id IN (keys)
  *
- * + *

* Given say 10 user id keys you might only get 7 results back. This can be more naturally represented in a map * than in an ordered list of values from the batch loader function. - * + *

* When the map is processed by the {@link org.dataloader.DataLoader} code, any keys that are missing in the map * will be replaced with null values. The semantic that the number of {@link org.dataloader.DataLoader#load(Object)} requests * are matched with and equal number of values is kept. - * + *

* This means that if 10 keys are asked for then {@link DataLoader#dispatch()} will return a promise of 10 value results and each * of the {@link org.dataloader.DataLoader#load(Object)} will complete with a value, null or an exception. - * + *

* When caching is disabled, its possible for the same key to be presented in the list of keys more than once. Your map * batch loader function needs to be resilient to this situation. * diff --git a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java index 0092563..a899fae 100644 --- a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java +++ b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java @@ -24,7 +24,7 @@ * This form of {@link MappedBatchLoader} is given a {@link org.dataloader.BatchLoaderEnvironment} object * that encapsulates the calling context. A typical use case is passing in security credentials or database details * for example. - * + *

* See {@link MappedBatchLoader} for more details on the design invariants that you must implement in order to * use this interface. */ diff --git a/src/main/java/org/dataloader/impl/PromisedValues.java b/src/main/java/org/dataloader/impl/PromisedValues.java index df0d374..840cf83 100644 --- a/src/main/java/org/dataloader/impl/PromisedValues.java +++ b/src/main/java/org/dataloader/impl/PromisedValues.java @@ -13,10 +13,10 @@ * This allows multiple {@link CompletionStage}s to be combined together and completed * as one and should something go wrong, instead of throwing {@link CompletionException}s it captures the cause and returns null for that * data value, other wise it allows you to access them as a list of values. - * + *

* This class really encapsulate a list of promised values. It is considered finished when all of the underlying futures * are finished. - * + *

* You can get that list of values via {@link #toList()}. You can also compose a {@link CompletableFuture} of that * list of values via {@link #toCompletableFuture()} ()} * diff --git a/src/main/java/org/dataloader/stats/ThreadLocalStatisticsCollector.java b/src/main/java/org/dataloader/stats/ThreadLocalStatisticsCollector.java index 7093584..ab7b51e 100644 --- a/src/main/java/org/dataloader/stats/ThreadLocalStatisticsCollector.java +++ b/src/main/java/org/dataloader/stats/ThreadLocalStatisticsCollector.java @@ -3,9 +3,9 @@ /** * This can collect statistics per thread as well as in an overall sense. This allows you to snapshot stats for a web request say * as well as all requests. - * + *

* You will want to call {@link #resetThread()} to clean up the thread local aspects of this object per request thread. - * + *

* ThreadLocals have their place in the Java world but be careful on how you use them. If you don't clean them up on "request boundaries" * then you WILL have misleading statistics. * From 51a1c7cbe4ec7385d16f62528ce8d5a03aac749a Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 26 Aug 2018 13:14:43 +1000 Subject: [PATCH 09/11] Readme correction --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 6e9d268..1d2b4e5 100644 --- a/README.md +++ b/README.md @@ -266,7 +266,7 @@ and some of which may have failed. From that data loader can infer the right be ```java DataLoader dataLoader = DataLoader.newDataLoaderWithTry(new BatchLoader>() { @Override - public CompletionStage>> load(List keys, BatchLoaderEnvironment environment) { + public CompletionStage>> load(List keys) { return CompletableFuture.supplyAsync(() -> { List> users = new ArrayList<>(); for (String key : keys) { From 70faf723970c625be8989bb690a129cdc65b0830 Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Sun, 26 Aug 2018 13:18:37 +1000 Subject: [PATCH 10/11] Opps - didnt want that --- src/main/java/org/dataloader/DataLoader.java | 2 +- src/test/java/TestFI.java | 63 -------------------- 2 files changed, 1 insertion(+), 64 deletions(-) delete mode 100644 src/test/java/TestFI.java diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 4cc7c57..fa86d2b 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -360,7 +360,7 @@ private DataLoader(Object batchLoadFunction, DataLoaderOptions options) { this.loaderQueue = new ArrayList<>(); this.stats = nonNull(this.loaderOptions.getStatisticsCollector()); - this.helper = new DataLoaderHelper<>(this, batchLoadFunction, this.loaderOptions, futureCache, loaderQueue, stats); + this.helper = new DataLoaderHelper<>(this, batchLoadFunction, this.loaderOptions, this.futureCache, this.loaderQueue, this.stats); } @SuppressWarnings("unchecked") diff --git a/src/test/java/TestFI.java b/src/test/java/TestFI.java deleted file mode 100644 index eeafadb..0000000 --- a/src/test/java/TestFI.java +++ /dev/null @@ -1,63 +0,0 @@ -import java.util.Collections; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; - -import static java.util.Collections.emptyMap; - -public class TestFI { - - interface BL { - List load(List keys); - } - - interface MapBL { - Map load(List keys); - } - - interface BLwithContext { - List load(List keys, Object context); - } - - interface MapBLwithContext { - Map load(List keys, Object context); - } - - - static class DL { - public DL(BL loader) { - } - - public DL(BLwithContext loader) { - } - - public DL(MapBLwithContext loader) { - } - - - public static DL newDLWithMap(MapBLwithContext loader) { - return null; - } - - public static DL newDL(BLwithContext loader) { - return null; - } - } - - - public static void main(String[] args) { - DL dl = new DL<>(keys -> keys.stream() - .map(Object::toString) - .collect(Collectors.toList())); - - BLwithContext integerStringBL = (keys, env) -> keys.stream() - .map(Object::toString) - .collect(Collectors.toList()); - DL dl2 = new DL(integerStringBL); - - MapBLwithContext mapBLwithContext = (keys, env) -> emptyMap(); - dl = new DL<>(mapBLwithContext); - - DL.newDLWithMap(((keys, context) -> emptyMap())); - } -} From bc307e05e9cbae0802f803708bc487d83ab30ee7 Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Mon, 27 Aug 2018 15:39:13 +1000 Subject: [PATCH 11/11] Added per load context objects and changed the context provider interface to be just an object --- README.md | 15 +-- .../BatchLoaderContextProvider.java | 13 +++ .../dataloader/BatchLoaderEnvironment.java | 35 ++++++- .../BatchLoaderEnvironmentProvider.java | 14 --- src/main/java/org/dataloader/DataLoader.java | 69 +++++++++++--- .../java/org/dataloader/DataLoaderHelper.java | 93 +++++++++++++++---- .../org/dataloader/DataLoaderOptions.java | 8 +- src/test/java/ReadmeExamples.java | 5 +- .../DataLoaderBatchLoaderEnvironmentTest.java | 82 ++++++++++++++-- 9 files changed, 264 insertions(+), 70 deletions(-) create mode 100644 src/main/java/org/dataloader/BatchLoaderContextProvider.java delete mode 100644 src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java diff --git a/README.md b/README.md index 1d2b4e5..ab80f6e 100644 --- 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 batchLoader = new BatchLoaderWithContext() { @Override diff --git a/src/main/java/org/dataloader/BatchLoaderContextProvider.java b/src/main/java/org/dataloader/BatchLoaderContextProvider.java new file mode 100644 index 0000000..b32c3f0 --- /dev/null +++ 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(); +} \ No newline at end of file diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironment.java b/src/main/java/org/dataloader/BatchLoaderEnvironment.java index c6d0126..8120673 100644 --- a/src/main/java/org/dataloader/BatchLoaderEnvironment.java +++ b/src/main/java/org/dataloader/BatchLoaderEnvironment.java @@ -1,5 +1,10 @@ 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. @@ -7,22 +12,43 @@ public class BatchLoaderEnvironment { private final Object context; + private final Map keyContexts; - private BatchLoaderEnvironment(Object context) { + private BatchLoaderEnvironment(Object context, Map keyContexts) { this.context = context; + this.keyContexts = keyContexts; } + /** + * Returns the overall context object provided by {@link org.dataloader.BatchLoaderContextProvider} + * + * @param the type you would like the object to be + * + * @return a context object or null if there isn't one + */ @SuppressWarnings("unchecked") public 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 getKeyContexts() { + return keyContexts; + } + public static Builder newBatchLoaderEnvironment() { return new Builder(); } public static class Builder { private Object context; + private Map keyContexts = Collections.emptyMap(); private Builder() { @@ -33,8 +59,13 @@ public Builder context(Object context) { return this; } + public Builder keyContexts(Map 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/BatchLoaderEnvironmentProvider.java b/src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java deleted file mode 100644 index b614e24..0000000 --- a/src/main/java/org/dataloader/BatchLoaderEnvironmentProvider.java +++ /dev/null @@ -1,14 +0,0 @@ -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(); -} \ No newline at end of file diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index fa86d2b..14508e2 100644 --- 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 { private final DataLoaderHelper helper; private final DataLoaderOptions loaderOptions; private final CacheMap> futureCache; - private final List>> loaderQueue; private final StatisticsCollector stats; /** @@ -237,6 +235,7 @@ public static DataLoader newMappedDataLoader(MappedBatchLoader + * * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects * @param the key type * @param 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> determineCacheMap(DataLoaderOptio * @return the future of the value */ public CompletableFuture 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. + *

+ * 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). + *

+ * 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 load(K key, Object keyContext) { + return helper.load(key, keyContext); } /** @@ -396,11 +413,39 @@ public CompletableFuture load(K key) { * @return the composite future of the list of values */ public CompletableFuture> loadMany(List keys) { - synchronized (this) { - List> 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. + *

+ * 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). + *

+ * 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> loadMany(List keys, List keyContexts) { + nonNull(keys); + nonNull(keyContexts); + + synchronized (this) { + List> 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 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(); } diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java index 87e1ebc..56a97c7 100644 --- a/src/main/java/org/dataloader/DataLoaderHelper.java +++ b/src/main/java/org/dataloader/DataLoaderHelper.java @@ -3,9 +3,9 @@ import org.dataloader.impl.CompletableFutureKit; import org.dataloader.stats.StatisticsCollector; -import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collection; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; @@ -26,23 +26,48 @@ */ class DataLoaderHelper { - private final DataLoader dataLoader; + class LoaderQueueEntry { + + final K key; + final V value; + final Object callContext; + + public LoaderQueueEntry(K key, V value, Object callContext) { + this.key = key; + this.value = value; + this.callContext = callContext; + } + + K getKey() { + return key; + } + + V getValue() { + return value; + } + + Object getCallContext() { + return callContext; + } + } + + private final DataLoader dataLoader; private final Object batchLoadFunction; private final DataLoaderOptions loaderOptions; private final CacheMap> futureCache; - private final List>> loaderQueue; + private final List>> loaderQueue; private final StatisticsCollector stats; - DataLoaderHelper(DataLoader dataLoader, Object batchLoadFunction, DataLoaderOptions loaderOptions, CacheMap> futureCache, List>> loaderQueue, StatisticsCollector stats) { + DataLoaderHelper(DataLoader dataLoader, Object batchLoadFunction, DataLoaderOptions loaderOptions, CacheMap> futureCache, StatisticsCollector stats) { this.dataLoader = dataLoader; this.batchLoadFunction = batchLoadFunction; this.loaderOptions = loaderOptions; this.futureCache = futureCache; - this.loaderQueue = loaderQueue; + this.loaderQueue = new ArrayList<>(); this.stats = stats; } - CompletableFuture load(K key) { + CompletableFuture load(K key, Object loadContext) { synchronized (dataLoader) { Object cacheKey = getCacheKey(nonNull(key)); stats.incrementLoadCount(); @@ -59,11 +84,11 @@ CompletableFuture load(K key) { CompletableFuture future = new CompletableFuture<>(); if (batchingEnabled) { - loaderQueue.add(new AbstractMap.SimpleImmutableEntry<>(key, future)); + loaderQueue.add(new LoaderQueueEntry<>(key, future, loadContext)); } else { stats.incrementBatchLoadCountBy(1); // immediate execution of batch function - future = invokeLoaderImmediately(key); + future = invokeLoaderImmediately(key, loadContext); } if (cachingEnabled) { futureCache.set(cacheKey, future); @@ -83,11 +108,13 @@ CompletableFuture> dispatch() { // // we copy the pre-loaded set of futures ready for dispatch final List keys = new ArrayList<>(); + final List callContexts = new ArrayList<>(); final List> queuedFutures = new ArrayList<>(); synchronized (dataLoader) { loaderQueue.forEach(entry -> { keys.add(entry.getKey()); queuedFutures.add(entry.getValue()); + callContexts.add(entry.getCallContext()); }); loaderQueue.clear(); } @@ -107,13 +134,13 @@ CompletableFuture> dispatch() { // int maxBatchSize = loaderOptions.maxBatchSize(); if (maxBatchSize > 0 && maxBatchSize < keys.size()) { - return sliceIntoBatchesOfBatches(keys, queuedFutures, maxBatchSize); + return sliceIntoBatchesOfBatches(keys, queuedFutures, callContexts, maxBatchSize); } else { - return dispatchQueueBatch(keys, queuedFutures); + return dispatchQueueBatch(keys, callContexts, queuedFutures); } } - private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List> queuedFutures, int maxBatchSize) { + private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List> queuedFutures, List callContexts, int maxBatchSize) { // the number of keys is > than what the batch loader function can accept // so make multiple calls to the loader List>> allBatches = new ArrayList<>(); @@ -126,8 +153,9 @@ private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List< List subKeys = keys.subList(fromIndex, toIndex); List> subFutures = queuedFutures.subList(fromIndex, toIndex); + List subCallContexts = callContexts.subList(fromIndex, toIndex); - allBatches.add(dispatchQueueBatch(subKeys, subFutures)); + allBatches.add(dispatchQueueBatch(subKeys, subCallContexts, subFutures)); } // // now reassemble all the futures into one that is the complete set of results @@ -139,9 +167,9 @@ private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List< } @SuppressWarnings("unchecked") - private CompletableFuture> dispatchQueueBatch(List keys, List> queuedFutures) { + private CompletableFuture> dispatchQueueBatch(List keys, List callContexts, List> queuedFutures) { stats.incrementBatchLoadCountBy(keys.size()); - CompletionStage> batchLoad = invokeLoader(keys); + CompletionStage> batchLoad = invokeLoader(keys, callContexts); return batchLoad .toCompletableFuture() .thenApply(values -> { @@ -190,11 +218,14 @@ private void assertResultSize(List keys, List values) { } - CompletableFuture invokeLoaderImmediately(K key) { + CompletableFuture invokeLoaderImmediately(K key, Object keyContext) { List keys = singletonList(key); CompletionStage singleLoadCall; try { - BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); + Object context = loaderOptions.getBatchLoaderEnvironmentProvider().getContext(); + Map keyContextMap = mkKeyContextMap(keys, singletonList(keyContext)); + BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment() + .context(context).keyContexts(keyContextMap).build(); if (isMapLoader()) { singleLoadCall = invokeMapBatchLoader(keys, environment).thenApply(list -> list.get(0)); } else { @@ -206,10 +237,13 @@ CompletableFuture invokeLoaderImmediately(K key) { } } - CompletionStage> invokeLoader(List keys) { + CompletionStage> invokeLoader(List keys, List keyContexts) { CompletionStage> batchLoad; try { - BatchLoaderEnvironment environment = loaderOptions.getBatchLoaderEnvironmentProvider().get(); + Object context = loaderOptions.getBatchLoaderEnvironmentProvider().getContext(); + Map keyContextMap = mkKeyContextMap(keys, keyContexts); + BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment() + .context(context).keyContexts(keyContextMap).build(); if (isMapLoader()) { batchLoad = invokeMapBatchLoader(keys, environment); } else { @@ -232,11 +266,11 @@ private CompletionStage> invokeListBatchLoader(List keys, BatchLoader return nonNull(loadResult, "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. */ - @SuppressWarnings("unchecked") private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderEnvironment environment) { CompletionStage> loadResult; @@ -259,4 +293,25 @@ private CompletionStage> invokeMapBatchLoader(List keys, BatchLoaderE private boolean isMapLoader() { return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext; } + + private Map mkKeyContextMap(List keys, List keyContexts) { + Map map = new HashMap<>(); + for (int i = 0; i < keys.size(); i++) { + K key = keys.get(i); + Object keyContext = null; + if (i < keyContexts.size()) { + keyContext = keyContexts.get(i); + } + if (keyContext != null) { + map.put(key, keyContext); + } + } + return map; + } + + int dispatchDepth() { + synchronized (dataLoader) { + return loaderQueue.size(); + } + } } diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java index 8cab094..c8710b7 100644 --- a/src/main/java/org/dataloader/DataLoaderOptions.java +++ b/src/main/java/org/dataloader/DataLoaderOptions.java @@ -31,7 +31,7 @@ */ public class DataLoaderOptions { - private static final BatchLoaderEnvironmentProvider NULL_PROVIDER = () -> BatchLoaderEnvironment.newBatchLoaderEnvironment().build(); + private static final BatchLoaderContextProvider NULL_PROVIDER = () -> null; private boolean batchingEnabled; private boolean cachingEnabled; @@ -39,7 +39,7 @@ public class DataLoaderOptions { private CacheMap cacheMap; private int maxBatchSize; private Supplier statisticsCollector; - private BatchLoaderEnvironmentProvider environmentProvider; + private BatchLoaderContextProvider environmentProvider; /** * Creates a new data loader options with default settings. @@ -210,7 +210,7 @@ public DataLoaderOptions setStatisticsCollector(Supplier st /** * @return the batch environment provider that will be used to give context to batch load functions */ - public BatchLoaderEnvironmentProvider getBatchLoaderEnvironmentProvider() { + public BatchLoaderContextProvider getBatchLoaderEnvironmentProvider() { return environmentProvider; } @@ -221,7 +221,7 @@ public BatchLoaderEnvironmentProvider getBatchLoaderEnvironmentProvider() { * * @return the data loader options for fluent coding */ - public DataLoaderOptions setBatchLoaderEnvironmentProvider(BatchLoaderEnvironmentProvider environmentProvider) { + public DataLoaderOptions setBatchLoaderEnvironmentProvider(BatchLoaderContextProvider environmentProvider) { this.environmentProvider = nonNull(environmentProvider); return this; } diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 878943b..a3345e6 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -84,11 +84,8 @@ public CompletionStage> load(List userIds) { } private void callContextExample() { - BatchLoaderEnvironment batchLoaderEnvironment = BatchLoaderEnvironment.newBatchLoaderEnvironment() - .context(SecurityCtx.getCallingUserCtx()).build(); - DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchLoaderEnvironmentProvider(() -> batchLoaderEnvironment); + .setBatchLoaderEnvironmentProvider(() -> SecurityCtx.getCallingUserCtx()); BatchLoaderWithContext batchLoader = new BatchLoaderWithContext() { @Override diff --git a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java index 88c2a4e..ac171bc 100644 --- a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java +++ b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java @@ -2,6 +2,8 @@ import org.junit.Test; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -9,7 +11,7 @@ import java.util.stream.Collectors; import static java.util.Arrays.asList; -import static org.dataloader.BatchLoaderEnvironment.newBatchLoaderEnvironment; +import static java.util.Collections.singletonList; import static org.hamcrest.Matchers.equalTo; import static org.junit.Assert.assertThat; @@ -18,6 +20,18 @@ */ public class DataLoaderBatchLoaderEnvironmentTest { + private BatchLoaderWithContext contextBatchLoader() { + return (keys, environment) -> { + List list = keys.stream().map(k -> { + Object context = environment.getContext(); + Object keyContext = environment.getKeyContexts().get(k); + return k + "-" + context + "-" + keyContext; + }).collect(Collectors.toList()); + return CompletableFuture.completedFuture(list); + }; + } + + @Test public void context_is_passed_to_batch_loader_function() throws Exception { BatchLoaderWithContext batchLoader = (keys, environment) -> { @@ -25,7 +39,7 @@ public void context_is_passed_to_batch_loader_function() throws Exception { return CompletableFuture.completedFuture(list); }; DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchLoaderEnvironmentProvider(() -> newBatchLoaderEnvironment().context("ctx").build()); + .setBatchLoaderEnvironmentProvider(() -> "ctx"); DataLoader loader = DataLoader.newDataLoader(batchLoader, options); loader.load("A"); @@ -37,24 +51,78 @@ public void context_is_passed_to_batch_loader_function() throws Exception { assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx"))); } + @Test + public void key_contexts_are_passed_to_batch_loader_function() throws Exception { + BatchLoaderWithContext batchLoader = contextBatchLoader(); + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderEnvironmentProvider(() -> "ctx"); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); + + loader.load("A", "aCtx"); + loader.load("B", "bCtx"); + loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx")); + + List results = loader.dispatchAndJoin(); + + assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-bCtx", "C-ctx-cCtx", "D-ctx-dCtx"))); + } + + @Test + public void key_contexts_are_passed_to_batch_loader_function_when_batching_disabled() throws Exception { + BatchLoaderWithContext batchLoader = contextBatchLoader(); + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchingEnabled(false) + .setBatchLoaderEnvironmentProvider(() -> "ctx"); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); + + CompletableFuture aLoad = loader.load("A", "aCtx"); + CompletableFuture bLoad = loader.load("B", "bCtx"); + CompletableFuture> canDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx")); + + List results = new ArrayList<>(asList(aLoad.join(), bLoad.join())); + results.addAll(canDLoad.join()); + + assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-bCtx", "C-ctx-cCtx", "D-ctx-dCtx"))); + } + + @Test + public void missing_key_contexts_are_passed_to_batch_loader_function() throws Exception { + BatchLoaderWithContext batchLoader = contextBatchLoader(); + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderEnvironmentProvider(() -> "ctx"); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); + + loader.load("A", "aCtx"); + loader.load("B"); + loader.loadMany(asList("C", "D"), singletonList("cCtx")); + + List results = loader.dispatchAndJoin(); + + assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null"))); + } + @Test public void context_is_passed_to_map_batch_loader_function() throws Exception { MappedBatchLoaderWithContext mapBatchLoader = (keys, environment) -> { Map map = new HashMap<>(); - keys.forEach(k -> map.put(k, k + "-" + environment.getContext())); + keys.forEach(k -> { + Object context = environment.getContext(); + Object keyContext = environment.getKeyContexts().get(k); + map.put(k, k + "-" + context + "-" + keyContext); + }); return CompletableFuture.completedFuture(map); }; DataLoaderOptions options = DataLoaderOptions.newOptions() - .setBatchLoaderEnvironmentProvider(() -> newBatchLoaderEnvironment().context("ctx").build()); + .setBatchLoaderEnvironmentProvider(() -> "ctx"); DataLoader loader = DataLoader.newMappedDataLoader(mapBatchLoader, options); - loader.load("A"); + loader.load("A", "aCtx"); loader.load("B"); - loader.loadMany(asList("C", "D")); + loader.loadMany(asList("C", "D"), singletonList("cCtx")); List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx"))); + assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null"))); } @Test