From f0506c4c1ec1dcf9d4a6e609fb4b81cc82551a25 Mon Sep 17 00:00:00 2001 From: Brad Baker Date: Tue, 23 Oct 2018 09:16:34 +1100 Subject: [PATCH] Added the list of call context objects into the environment --- README.md | 30 +++++++++++++ .../dataloader/BatchLoaderEnvironment.java | 44 ++++++++++++++++--- .../java/org/dataloader/DataLoaderHelper.java | 22 +--------- src/test/java/ReadmeExamples.java | 26 +++++++++++ .../DataLoaderBatchLoaderEnvironmentTest.java | 32 +++++++++++--- 5 files changed, 124 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 0c9dc4e..acff2b4 100644 --- a/README.md +++ b/README.md @@ -194,6 +194,36 @@ for the context object. 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. +You can also pass in context objects per load call. This will be captured and passed to the batch loader function. + +You can gain access to them as a map by key or as the original list of context objects. + +```java + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx()); + + BatchLoaderWithContext batchLoader = new BatchLoaderWithContext() { + @Override + public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { + SecurityCtx callCtx = environment.getContext(); + // + // this is the load context objects in map form by key + // in this case [ keyA : contextForA, keyB : contextForB ] + // + Map keyContexts = environment.getKeyContexts(); + // + // this is load context in list form + // + // in this case [ contextForA, contextForB ] + return callDatabaseForResults(callCtx, keys); + } + }; + + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); + loader.load("keyA", "contextForA"); + loader.load("keyB", "contextForB"); +``` + ### 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. diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironment.java b/src/main/java/org/dataloader/BatchLoaderEnvironment.java index 8120673..5aa4a35 100644 --- a/src/main/java/org/dataloader/BatchLoaderEnvironment.java +++ b/src/main/java/org/dataloader/BatchLoaderEnvironment.java @@ -1,9 +1,12 @@ package org.dataloader; +import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; +import java.util.List; import java.util.Map; -import static org.dataloader.impl.Assertions.nonNull; +import static java.util.Objects.nonNull; /** * This object is passed to a batch loader as calling context. It could contain security credentials @@ -13,10 +16,12 @@ public class BatchLoaderEnvironment { private final Object context; private final Map keyContexts; + private final List keyContextsList; - private BatchLoaderEnvironment(Object context, Map keyContexts) { + private BatchLoaderEnvironment(Object context, List keyContextsList, Map keyContexts) { this.context = context; this.keyContexts = keyContexts; + this.keyContextsList = keyContextsList; } /** @@ -42,6 +47,17 @@ public Map getKeyContexts() { return keyContexts; } + /** + * 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 list of them is present by this method. + * + * @return a list of key context objects in the order they where encountered + */ + public List getKeyContextsList() { + return keyContextsList; + } + public static Builder newBatchLoaderEnvironment() { return new Builder(); } @@ -49,6 +65,7 @@ public static Builder newBatchLoaderEnvironment() { public static class Builder { private Object context; private Map keyContexts = Collections.emptyMap(); + private List keyContextsList = Collections.emptyList(); private Builder() { @@ -59,13 +76,30 @@ public Builder context(Object context) { return this; } - public Builder keyContexts(Map keyContexts) { - this.keyContexts = nonNull(keyContexts); + public Builder keyContexts(List keys, List keyContexts) { + nonNull(keys); + nonNull(keyContexts); + + Map map = new HashMap<>(); + List list = 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); + } + if (keyContext != null) { + map.put(key, keyContext); + } + list.add(keyContext); + } + this.keyContexts = map; + this.keyContextsList = list; return this; } public BatchLoaderEnvironment build() { - return new BatchLoaderEnvironment(context, keyContexts); + return new BatchLoaderEnvironment(context, keyContextsList, keyContexts); } } } diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java index d1e4be5..ea8693a 100644 --- a/src/main/java/org/dataloader/DataLoaderHelper.java +++ b/src/main/java/org/dataloader/DataLoaderHelper.java @@ -5,7 +5,6 @@ import java.util.ArrayList; import java.util.Collection; -import java.util.HashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Map; @@ -225,9 +224,8 @@ CompletableFuture invokeLoaderImmediately(K key, Object keyContext) { CompletionStage singleLoadCall; try { Object context = loaderOptions.getBatchLoaderContextProvider().getContext(); - Map keyContextMap = mkKeyContextMap(keys, singletonList(keyContext)); BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment() - .context(context).keyContexts(keyContextMap).build(); + .context(context).keyContexts(keys, singletonList(keyContext)).build(); if (isMapLoader()) { singleLoadCall = invokeMapBatchLoader(keys, environment).thenApply(list -> list.get(0)); } else { @@ -243,9 +241,8 @@ CompletionStage> invokeLoader(List keys, List keyContexts) { CompletionStage> batchLoad; try { Object context = loaderOptions.getBatchLoaderContextProvider().getContext(); - Map keyContextMap = mkKeyContextMap(keys, keyContexts); BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment() - .context(context).keyContexts(keyContextMap).build(); + .context(context).keyContexts(keys, keyContexts).build(); if (isMapLoader()) { batchLoad = invokeMapBatchLoader(keys, environment); } else { @@ -297,21 +294,6 @@ 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/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 0803728..523cb5a 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -99,6 +99,32 @@ public CompletionStage> load(List keys, BatchLoaderEnvironm DataLoader loader = DataLoader.newDataLoader(batchLoader, options); } + private void keyContextExample() { + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx()); + + BatchLoaderWithContext batchLoader = new BatchLoaderWithContext() { + @Override + public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { + SecurityCtx callCtx = environment.getContext(); + // + // this is the load context objects in map form by key + // in this case [ keyA : contextForA, keyB : contextForB ] + // + Map keyContexts = environment.getKeyContexts(); + // + // this is load context in list form + // + // in this case [ contextForA, contextForB ] + return callDatabaseForResults(callCtx, keys); + } + }; + + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); + loader.load("keyA", "contextForA"); + loader.load("keyB", "contextForB"); + } + private CompletionStage> callDatabaseForResults(SecurityCtx callCtx, List keys) { return null; } diff --git a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java index a727bf1..575fffd 100644 --- a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java +++ b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java @@ -7,6 +7,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static java.util.Arrays.asList; @@ -21,10 +22,13 @@ public class DataLoaderBatchLoaderEnvironmentTest { private BatchLoaderWithContext contextBatchLoader() { return (keys, environment) -> { + AtomicInteger index = new AtomicInteger(0); List list = keys.stream().map(k -> { + int i = index.getAndIncrement(); Object context = environment.getContext(); - Object keyContext = environment.getKeyContexts().get(k); - return k + "-" + context + "-" + keyContext; + Object keyContextM = environment.getKeyContexts().get(k); + Object keyContextL = environment.getKeyContextsList().get(i); + return k + "-" + context + "-m:" + keyContextM + "-l:" + keyContextL; }).collect(Collectors.toList()); return CompletableFuture.completedFuture(list); }; @@ -63,7 +67,7 @@ public void key_contexts_are_passed_to_batch_loader_function() throws Exception List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-bCtx", "C-ctx-cCtx", "D-ctx-dCtx"))); + assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx"))); } @Test @@ -81,7 +85,7 @@ public void key_contexts_are_passed_to_batch_loader_function_when_batching_disab 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"))); + assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx"))); } @Test @@ -97,7 +101,7 @@ public void missing_key_contexts_are_passed_to_batch_loader_function() throws Ex List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null"))); + assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null"))); } @Test @@ -158,4 +162,22 @@ public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() throws E assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null"))); } + + @Test + public void mmap_semantics_apply_to_batch_loader_context() throws Exception { + BatchLoaderWithContext batchLoader = contextBatchLoader(); + DataLoaderOptions options = DataLoaderOptions.newOptions() + .setBatchLoaderContextProvider(() -> "ctx") + .setCachingEnabled(false); + DataLoader loader = DataLoader.newDataLoader(batchLoader, options); + + loader.load("A", "aCtx"); + loader.load("B", "bCtx"); + loader.load("A", "overridesCtx"); + + List results = loader.dispatchAndJoin(); + + assertThat(results, equalTo(asList("A-ctx-m:overridesCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "A-ctx-m:overridesCtx-l:overridesCtx"))); + } + }