Skip to content

Commit f0506c4

Browse files
bbakermanbbakerman
committed
Added the list of call context objects into the environment
1 parent 20c6e7a commit f0506c4

File tree

5 files changed

+124
-30
lines changed

5 files changed

+124
-30
lines changed

‎README.md

+30
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,36 @@ for the context object.
194194
The batch loading code will now receive this environment object and it can be used to get context perhaps allowing it
195195
to connect to other systems.
196196

197+
You can also pass in context objects per load call. This will be captured and passed to the batch loader function.
198+
199+
You can gain access to them as a map by key or as the original list of context objects.
200+
201+
```java
202+
DataLoaderOptions options = DataLoaderOptions.newOptions()
203+
.setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx());
204+
205+
BatchLoaderWithContext<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
206+
@Override
207+
public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironment environment) {
208+
SecurityCtx callCtx = environment.getContext();
209+
//
210+
// this is the load context objects in map form by key
211+
// in this case [ keyA : contextForA, keyB : contextForB ]
212+
//
213+
Map<Object, Object> keyContexts = environment.getKeyContexts();
214+
//
215+
// this is load context in list form
216+
//
217+
// in this case [ contextForA, contextForB ]
218+
return callDatabaseForResults(callCtx, keys);
219+
}
220+
};
221+
222+
DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
223+
loader.load("keyA", "contextForA");
224+
loader.load("keyB", "contextForB");
225+
```
226+
197227
### Returning a Map of results from your batch loader
198228

199229
Often there is not a 1:1 mapping of your batch loaded keys to the values returned.

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

+39-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package org.dataloader;
22

3+
import java.util.ArrayList;
34
import java.util.Collections;
5+
import java.util.HashMap;
6+
import java.util.List;
47
import java.util.Map;
58

6-
import static org.dataloader.impl.Assertions.nonNull;
9+
import static java.util.Objects.nonNull;
710

811
/**
912
* This object is passed to a batch loader as calling context. It could contain security credentials
@@ -13,10 +16,12 @@ public class BatchLoaderEnvironment {
1316

1417
private final Object context;
1518
private final Map<Object, Object> keyContexts;
19+
private final List<Object> keyContextsList;
1620

17-
private BatchLoaderEnvironment(Object context, Map<Object, Object> keyContexts) {
21+
private BatchLoaderEnvironment(Object context, List<Object> keyContextsList, Map<Object, Object> keyContexts) {
1822
this.context = context;
1923
this.keyContexts = keyContexts;
24+
this.keyContextsList = keyContextsList;
2025
}
2126

2227
/**
@@ -42,13 +47,25 @@ public Map<Object, Object> getKeyContexts() {
4247
return keyContexts;
4348
}
4449

50+
/**
51+
* Each call to {@link org.dataloader.DataLoader#load(Object, Object)} or
52+
* {@link org.dataloader.DataLoader#loadMany(java.util.List, java.util.List)} can be given
53+
* a context object when it is invoked. A list of them is present by this method.
54+
*
55+
* @return a list of key context objects in the order they where encountered
56+
*/
57+
public List<Object> getKeyContextsList() {
58+
return keyContextsList;
59+
}
60+
4561
public static Builder newBatchLoaderEnvironment() {
4662
return new Builder();
4763
}
4864

4965
public static class Builder {
5066
private Object context;
5167
private Map<Object, Object> keyContexts = Collections.emptyMap();
68+
private List<Object> keyContextsList = Collections.emptyList();
5269

5370
private Builder() {
5471

@@ -59,13 +76,30 @@ public Builder context(Object context) {
5976
return this;
6077
}
6178

62-
public Builder keyContexts(Map<Object, Object> keyContexts) {
63-
this.keyContexts = nonNull(keyContexts);
79+
public <K> Builder keyContexts(List<K> keys, List<Object> keyContexts) {
80+
nonNull(keys);
81+
nonNull(keyContexts);
82+
83+
Map<Object, Object> map = new HashMap<>();
84+
List<Object> list = new ArrayList<>();
85+
for (int i = 0; i < keys.size(); i++) {
86+
K key = keys.get(i);
87+
Object keyContext = null;
88+
if (i < keyContexts.size()) {
89+
keyContext = keyContexts.get(i);
90+
}
91+
if (keyContext != null) {
92+
map.put(key, keyContext);
93+
}
94+
list.add(keyContext);
95+
}
96+
this.keyContexts = map;
97+
this.keyContextsList = list;
6498
return this;
6599
}
66100

67101
public BatchLoaderEnvironment build() {
68-
return new BatchLoaderEnvironment(context, keyContexts);
102+
return new BatchLoaderEnvironment(context, keyContextsList, keyContexts);
69103
}
70104
}
71105
}

‎src/main/java/org/dataloader/DataLoaderHelper.java

+2-20
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
import java.util.ArrayList;
77
import java.util.Collection;
8-
import java.util.HashMap;
98
import java.util.LinkedHashSet;
109
import java.util.List;
1110
import java.util.Map;
@@ -225,9 +224,8 @@ CompletableFuture<V> invokeLoaderImmediately(K key, Object keyContext) {
225224
CompletionStage<V> singleLoadCall;
226225
try {
227226
Object context = loaderOptions.getBatchLoaderContextProvider().getContext();
228-
Map<Object, Object> keyContextMap = mkKeyContextMap(keys, singletonList(keyContext));
229227
BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment()
230-
.context(context).keyContexts(keyContextMap).build();
228+
.context(context).keyContexts(keys, singletonList(keyContext)).build();
231229
if (isMapLoader()) {
232230
singleLoadCall = invokeMapBatchLoader(keys, environment).thenApply(list -> list.get(0));
233231
} else {
@@ -243,9 +241,8 @@ CompletionStage<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts) {
243241
CompletionStage<List<V>> batchLoad;
244242
try {
245243
Object context = loaderOptions.getBatchLoaderContextProvider().getContext();
246-
Map<Object, Object> keyContextMap = mkKeyContextMap(keys, keyContexts);
247244
BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment()
248-
.context(context).keyContexts(keyContextMap).build();
245+
.context(context).keyContexts(keys, keyContexts).build();
249246
if (isMapLoader()) {
250247
batchLoad = invokeMapBatchLoader(keys, environment);
251248
} else {
@@ -297,21 +294,6 @@ private boolean isMapLoader() {
297294
return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext;
298295
}
299296

300-
private Map<Object, Object> mkKeyContextMap(List<K> keys, List<Object> keyContexts) {
301-
Map<Object, Object> map = new HashMap<>();
302-
for (int i = 0; i < keys.size(); i++) {
303-
K key = keys.get(i);
304-
Object keyContext = null;
305-
if (i < keyContexts.size()) {
306-
keyContext = keyContexts.get(i);
307-
}
308-
if (keyContext != null) {
309-
map.put(key, keyContext);
310-
}
311-
}
312-
return map;
313-
}
314-
315297
int dispatchDepth() {
316298
synchronized (dataLoader) {
317299
return loaderQueue.size();

‎src/test/java/ReadmeExamples.java

+26
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,32 @@ public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironm
9999
DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
100100
}
101101

102+
private void keyContextExample() {
103+
DataLoaderOptions options = DataLoaderOptions.newOptions()
104+
.setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx());
105+
106+
BatchLoaderWithContext<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
107+
@Override
108+
public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironment environment) {
109+
SecurityCtx callCtx = environment.getContext();
110+
//
111+
// this is the load context objects in map form by key
112+
// in this case [ keyA : contextForA, keyB : contextForB ]
113+
//
114+
Map<Object, Object> keyContexts = environment.getKeyContexts();
115+
//
116+
// this is load context in list form
117+
//
118+
// in this case [ contextForA, contextForB ]
119+
return callDatabaseForResults(callCtx, keys);
120+
}
121+
};
122+
123+
DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
124+
loader.load("keyA", "contextForA");
125+
loader.load("keyB", "contextForB");
126+
}
127+
102128
private CompletionStage<List<String>> callDatabaseForResults(SecurityCtx callCtx, List<String> keys) {
103129
return null;
104130
}

‎src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java

+27-5
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.util.List;
88
import java.util.Map;
99
import java.util.concurrent.CompletableFuture;
10+
import java.util.concurrent.atomic.AtomicInteger;
1011
import java.util.stream.Collectors;
1112

1213
import static java.util.Arrays.asList;
@@ -21,10 +22,13 @@ public class DataLoaderBatchLoaderEnvironmentTest {
2122

2223
private BatchLoaderWithContext<String, String> contextBatchLoader() {
2324
return (keys, environment) -> {
25+
AtomicInteger index = new AtomicInteger(0);
2426
List<String> list = keys.stream().map(k -> {
27+
int i = index.getAndIncrement();
2528
Object context = environment.getContext();
26-
Object keyContext = environment.getKeyContexts().get(k);
27-
return k + "-" + context + "-" + keyContext;
29+
Object keyContextM = environment.getKeyContexts().get(k);
30+
Object keyContextL = environment.getKeyContextsList().get(i);
31+
return k + "-" + context + "-m:" + keyContextM + "-l:" + keyContextL;
2832
}).collect(Collectors.toList());
2933
return CompletableFuture.completedFuture(list);
3034
};
@@ -63,7 +67,7 @@ public void key_contexts_are_passed_to_batch_loader_function() throws Exception
6367

6468
List<String> results = loader.dispatchAndJoin();
6569

66-
assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-bCtx", "C-ctx-cCtx", "D-ctx-dCtx")));
70+
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")));
6771
}
6872

6973
@Test
@@ -81,7 +85,7 @@ public void key_contexts_are_passed_to_batch_loader_function_when_batching_disab
8185
List<String> results = new ArrayList<>(asList(aLoad.join(), bLoad.join()));
8286
results.addAll(canDLoad.join());
8387

84-
assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-bCtx", "C-ctx-cCtx", "D-ctx-dCtx")));
88+
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")));
8589
}
8690

8791
@Test
@@ -97,7 +101,7 @@ public void missing_key_contexts_are_passed_to_batch_loader_function() throws Ex
97101

98102
List<String> results = loader.dispatchAndJoin();
99103

100-
assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null")));
104+
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")));
101105
}
102106

103107
@Test
@@ -158,4 +162,22 @@ public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() throws E
158162

159163
assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null")));
160164
}
165+
166+
@Test
167+
public void mmap_semantics_apply_to_batch_loader_context() throws Exception {
168+
BatchLoaderWithContext<String, String> batchLoader = contextBatchLoader();
169+
DataLoaderOptions options = DataLoaderOptions.newOptions()
170+
.setBatchLoaderContextProvider(() -> "ctx")
171+
.setCachingEnabled(false);
172+
DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
173+
174+
loader.load("A", "aCtx");
175+
loader.load("B", "bCtx");
176+
loader.load("A", "overridesCtx");
177+
178+
List<String> results = loader.dispatchAndJoin();
179+
180+
assertThat(results, equalTo(asList("A-ctx-m:overridesCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "A-ctx-m:overridesCtx-l:overridesCtx")));
181+
}
182+
161183
}

0 commit comments

Comments
 (0)