Skip to content

Commit

Permalink
Added the list of call context objects into the environment
Browse files Browse the repository at this point in the history
  • Loading branch information
@bbakerman
bbakerman committed Oct 23, 2018
1 parent 20c6e7a commit f0506c4
Show file tree
Hide file tree
Showing 5 changed files with 124 additions and 30 deletions.
30 changes: 30 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
@Override
public CompletionStage<List<String>> load(List<String> 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<Object, Object> keyContexts = environment.getKeyContexts();
//
// this is load context in list form
//
// in this case [ contextForA, contextForB ]
return callDatabaseForResults(callCtx, keys);
}
};

DataLoader<String, String> 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.
Expand Down
44 changes: 39 additions & 5 deletions src/main/java/org/dataloader/BatchLoaderEnvironment.java
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -13,10 +16,12 @@ public class BatchLoaderEnvironment {

private final Object context;
private final Map<Object, Object> keyContexts;
private final List<Object> keyContextsList;

private BatchLoaderEnvironment(Object context, Map<Object, Object> keyContexts) {
private BatchLoaderEnvironment(Object context, List<Object> keyContextsList, Map<Object, Object> keyContexts) {
this.context = context;
this.keyContexts = keyContexts;
this.keyContextsList = keyContextsList;
}

/**
Expand All @@ -42,13 +47,25 @@ public Map<Object, Object> 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<Object> getKeyContextsList() {
return keyContextsList;
}

public static Builder newBatchLoaderEnvironment() {
return new Builder();
}

public static class Builder {
private Object context;
private Map<Object, Object> keyContexts = Collections.emptyMap();
private List<Object> keyContextsList = Collections.emptyList();

private Builder() {

Expand All @@ -59,13 +76,30 @@ public Builder context(Object context) {
return this;
}

public Builder keyContexts(Map<Object, Object> keyContexts) {
this.keyContexts = nonNull(keyContexts);
public <K> Builder keyContexts(List<K> keys, List<Object> keyContexts) {
nonNull(keys);
nonNull(keyContexts);

Map<Object, Object> map = new HashMap<>();
List<Object> 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);
}
}
}
22 changes: 2 additions & 20 deletions src/main/java/org/dataloader/DataLoaderHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -225,9 +224,8 @@ CompletableFuture<V> invokeLoaderImmediately(K key, Object keyContext) {
CompletionStage<V> singleLoadCall;
try {
Object context = loaderOptions.getBatchLoaderContextProvider().getContext();
Map<Object, Object> 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 {
Expand All @@ -243,9 +241,8 @@ CompletionStage<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts) {
CompletionStage<List<V>> batchLoad;
try {
Object context = loaderOptions.getBatchLoaderContextProvider().getContext();
Map<Object, Object> 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 {
Expand Down Expand Up @@ -297,21 +294,6 @@ private boolean isMapLoader() {
return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext;
}

private Map<Object, Object> mkKeyContextMap(List<K> keys, List<Object> keyContexts) {
Map<Object, Object> 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();
Expand Down
26 changes: 26 additions & 0 deletions src/test/java/ReadmeExamples.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,32 @@ public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironm
DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
}

private void keyContextExample() {
DataLoaderOptions options = DataLoaderOptions.newOptions()
.setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx());

BatchLoaderWithContext<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
@Override
public CompletionStage<List<String>> load(List<String> 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<Object, Object> keyContexts = environment.getKeyContexts();
//
// this is load context in list form
//
// in this case [ contextForA, contextForB ]
return callDatabaseForResults(callCtx, keys);
}
};

DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
loader.load("keyA", "contextForA");
loader.load("keyB", "contextForB");
}

private CompletionStage<List<String>> callDatabaseForResults(SecurityCtx callCtx, List<String> keys) {
return null;
}
Expand Down
32 changes: 27 additions & 5 deletions src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -21,10 +22,13 @@ public class DataLoaderBatchLoaderEnvironmentTest {

private BatchLoaderWithContext<String, String> contextBatchLoader() {
return (keys, environment) -> {
AtomicInteger index = new AtomicInteger(0);
List<String> 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);
};
Expand Down Expand Up @@ -63,7 +67,7 @@ public void key_contexts_are_passed_to_batch_loader_function() throws Exception

List<String> 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
Expand All @@ -81,7 +85,7 @@ public void key_contexts_are_passed_to_batch_loader_function_when_batching_disab
List<String> 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
Expand All @@ -97,7 +101,7 @@ public void missing_key_contexts_are_passed_to_batch_loader_function() throws Ex

List<String> 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
Expand Down Expand Up @@ -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<String, String> batchLoader = contextBatchLoader();
DataLoaderOptions options = DataLoaderOptions.newOptions()
.setBatchLoaderContextProvider(() -> "ctx")
.setCachingEnabled(false);
DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);

loader.load("A", "aCtx");
loader.load("B", "bCtx");
loader.load("A", "overridesCtx");

List<String> 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")));
}

}

0 comments on commit f0506c4

Please sign in to comment.