Spring/spring-graphql
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Context for batch loading functions

Browse Source 
Batch loading functions registered via BatchLoaderRegistry now have
Reactor Context propagated to them and also have access to the
GraphQLContext from the ExecutionInput.

Closes gh-173
This commit is contained in:
Rossen Stoyanchev
2021-11-03 17:08:19 +00:00
parent 6f16e5e6ea
commit 5aaf5b1df4
9 changed files with 128 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
spring-graphql-docs/src/docs/asciidoc/index.adoc
View File

@@ -353,14 +353,14 @@ defers until it is ready to batch load all related entities as one.
- `DataLoader` additionally maintains a cache of previously loaded entities that can
further improve efficiency when the same entity is in multiple places of the response.
Spring GraphQL provides a `BatchLoaderRegistry` that accepts and stores registrations of
batch loading functions. The `ExecutionGraphQlService` accepts the registry as input and
uses it to make per request `DataLoader` registrations. A `DataFetcher` then looks up the
`DataLoader` for an entity and uses it to load instances, or in an annotated controller,
simply declare a <<controllers-schema-mapping-data-loader,DataLoader argument>> to access the
registered loader. Annotated controllers also support a
<<controllers-batch-mapping,@BatchMapping>> that avoids the need to use `DataLoader`
directly.
Spring GraphQL provides:
- `BatchLoaderRegistry` that accepts and stores registrations of batch loading functions;
This is used in `ExecutionGraphQlService` to make `DataLoader` registrations per request.
- <<controllers-schema-mapping-data-loader,DataLoader argument>> for `@SchemaMapping`
methods to access the `DataLoader` for the field type.
- <<controllers-batch-mapping,@BatchMapping>> data controller methods that provide a
shortcut and avoid the need to use `DataLoader` directly.
The Spring Boot starter declares a
<<boot-graphql-batch-loader-registry,BatchLoaderRegistry bean>>, so that applications can

14
spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java
View File

@@ -21,6 +21,8 @@ import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import graphql.ExecutionInput;
import org.dataloader.BatchLoaderContextProvider;
import org.dataloader.BatchLoaderEnvironment;
import org.dataloader.DataLoaderOptions;
import reactor.core.publisher.Flux;
@@ -99,6 +101,11 @@ public interface BatchLoaderRegistry extends DataLoaderRegistrar {
/**
* Customize the {@link DataLoaderOptions} to use to create the
* {@link org.dataloader.DataLoader} via {@link org.dataloader.DataLoaderFactory}.
* <p><strong>Note:</strong> Do not set
* {@link DataLoaderOptions#setBatchLoaderContextProvider(BatchLoaderContextProvider)}
* as this will be set later to a provider that returns the context from
* {@link ExecutionInput#getGraphQLContext()}, so that batch loading
* functions and data fetchers can rely on access to the same context.
* @param optionsConsumer callback to customize the options, invoked
* immediately and given access to the options instance
* @return a spec to complete the registration
@@ -106,8 +113,13 @@ public interface BatchLoaderRegistry extends DataLoaderRegistrar {
RegistrationSpec<K, V> withOptions(Consumer<DataLoaderOptions> optionsConsumer);
/**
* Replace the {@link DataLoaderOptions} to use to create the
* Set the {@link DataLoaderOptions} to use to create the
* {@link org.dataloader.DataLoader} via {@link org.dataloader.DataLoaderFactory}.
* <p><strong>Note:</strong> Do not set
* {@link DataLoaderOptions#setBatchLoaderContextProvider(BatchLoaderContextProvider)}
* as this will be set later to a provider that returns the context from
* {@link ExecutionInput#getGraphQLContext()}, so that batch loading
* functions and data fetchers can rely on access to the same context.
* @param options the options to use
* @return a spec to complete the registration
*/

13
spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java
View File

@@ -15,14 +15,18 @@
*/
package org.springframework.graphql.execution;
import graphql.ExecutionInput;
import graphql.GraphQLContext;
import org.dataloader.BatchLoaderContextProvider;
import org.dataloader.DataLoaderRegistry;
/**
* Contract for callback access to the {@link DataLoaderRegistry} as it is
* initialized for each request.
* Contract for access to the {@link DataLoaderRegistry} for each request for
* the purpose of registering {@link org.dataloader.DataLoader} instances.
*
* @author Rossen Stoyanchev
* @since 1.0.0
* @see ExecutionInput#getDataLoaderRegistry()
*/
public interface DataLoaderRegistrar {
@@ -30,7 +34,10 @@ public interface DataLoaderRegistrar {
* Callback that provides access to the {@link DataLoaderRegistry} from the
* the {@link graphql.ExecutionInput}.
* @param registry the registry to make registrations against
* @param context the GraphQLContext from the ExecutionInput that registrars
* should set in the {@link org.dataloader.DataLoaderOptions} so that batch
* loaders can access it via {@link org.dataloader.BatchLoaderEnvironment}.
*/
void registerDataLoaders(DataLoaderRegistry registry);
void registerDataLoaders(DataLoaderRegistry registry, GraphQLContext context);
}

28
spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java
View File

@@ -23,6 +23,8 @@ import java.util.concurrent.CompletionStage;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import graphql.GraphQLContext;
import org.dataloader.BatchLoaderContextProvider;
import org.dataloader.BatchLoaderEnvironment;
import org.dataloader.BatchLoaderWithContext;
import org.dataloader.DataLoader;
@@ -32,6 +34,7 @@ import org.dataloader.DataLoaderRegistry;
import org.dataloader.MappedBatchLoaderWithContext;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -63,13 +66,16 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
}
@Override
public void registerDataLoaders(DataLoaderRegistry registry) {
public void registerDataLoaders(DataLoaderRegistry registry, GraphQLContext context) {
BatchLoaderContextProvider contextProvider = () -> context;
for (ReactorBatchLoader<?, ?> loader : this.loaders) {
DataLoader<?, ?> dataLoader = DataLoaderFactory.newDataLoader(loader, loader.getOptions());
DataLoaderOptions options = loader.getOptions().setBatchLoaderContextProvider(contextProvider);
DataLoader<?, ?> dataLoader = DataLoaderFactory.newDataLoader(loader, options);
registerDataLoader(loader.getName(), dataLoader, registry);
}
for (ReactorMappedBatchLoader<?, ?> loader : this.mappedLoaders) {
DataLoader<?, ?> dataLoader = DataLoaderFactory.newMappedDataLoader(loader, loader.getOptions());
DataLoaderOptions options = loader.getOptions().setBatchLoaderContextProvider(contextProvider);
DataLoader<?, ?> dataLoader = DataLoaderFactory.newMappedDataLoader(loader, options);
registerDataLoader(loader.getName(), dataLoader, registry);
}
}
@@ -141,6 +147,10 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
}
/**
* {@link BatchLoaderWithContext} that delegates to a {@link Flux} batch
* loading function and exposes Reactor context to it.
*/
private static class ReactorBatchLoader<K, V> implements BatchLoaderWithContext<K, V> {
private final String name;
@@ -168,11 +178,17 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
@Override
public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
return this.loader.apply(keys, environment).collectList().toFuture();
ContextView contextView = ReactorContextManager.getReactorContext(environment);
return this.loader.apply(keys, environment).collectList().contextWrite(contextView).toFuture();
}
}
/**
* {@link MappedBatchLoaderWithContext} that delegates to a {@link Mono}
* batch loading function and exposes Reactor context to it.
*/
private static class ReactorMappedBatchLoader<K, V> implements MappedBatchLoaderWithContext<K, V> {
private final String name;
@@ -200,8 +216,10 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
@Override
public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
return this.loader.apply(keys, environment).toFuture();
ContextView contextView = ReactorContextManager.getReactorContext(environment);
return this.loader.apply(keys, environment).contextWrite(contextView).toFuture();
}
}
}

20
spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java
View File

@@ -22,6 +22,7 @@ import java.util.List;
import graphql.ExecutionInput;
import graphql.ExecutionResult;
import graphql.GraphQL;
import graphql.GraphQLContext;
import org.dataloader.DataLoaderRegistry;
import reactor.core.publisher.Mono;
@@ -59,24 +60,23 @@ public class ExecutionGraphQlService implements GraphQlService {
@Override
public final Mono<ExecutionResult> execute(RequestInput requestInput) {
ExecutionInput executionInput = initExecutionInput(requestInput);
GraphQL graphQl = this.graphQlSource.graphQl();
return Mono.deferContextual((contextView) -> {
ExecutionInput executionInput = requestInput.toExecutionInput();
ReactorContextManager.setReactorContext(contextView, executionInput);
return Mono.fromFuture(graphQl.executeAsync(executionInput));
executionInput = regsterDataLoaders(executionInput);
return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(executionInput));
});
}
private ExecutionInput initExecutionInput(RequestInput requestInput) {
ExecutionInput input = requestInput.toExecutionInput();
private ExecutionInput regsterDataLoaders(ExecutionInput executionInput) {
if (!this.dataLoaderRegistrars.isEmpty()) {
DataLoaderRegistry previousRegistry = input.getDataLoaderRegistry();
GraphQLContext graphQLContext = executionInput.getGraphQLContext();
DataLoaderRegistry previousRegistry = executionInput.getDataLoaderRegistry();
DataLoaderRegistry newRegistry = DataLoaderRegistry.newRegistry().registerAll(previousRegistry).build();
this.dataLoaderRegistrars.forEach(registrar -> registrar.registerDataLoaders(newRegistry));
input = input.transform(builder -> builder.dataLoaderRegistry(newRegistry));
this.dataLoaderRegistrars.forEach(registrar -> registrar.registerDataLoaders(newRegistry, graphQLContext));
executionInput = executionInput.transform(builder -> builder.dataLoaderRegistry(newRegistry));
}
return input;
return executionInput;
}
}

13
spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java
View File

@@ -22,10 +22,12 @@ import java.util.Map;
import graphql.ExecutionInput;
import graphql.GraphQLContext;
import graphql.schema.DataFetchingEnvironment;
import org.dataloader.BatchLoaderEnvironment;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
/**
* Provides helper methods to save Reactor context in the {@link ExecutionInput}
@@ -68,6 +70,17 @@ public abstract class ReactorContextManager {
return graphQlContext.getOrDefault(CONTEXT_VIEW_KEY, Context.empty());
}
/**
* Return the Reactor {@link ContextView} saved in the given BatchLoaderEnvironment.
* @param environment the BatchLoaderEnvironment
* @return the reactor {@link ContextView}
*/
static ContextView getReactorContext(BatchLoaderEnvironment environment) {
Object context = environment.getContext();
Assert.isTrue(context instanceof GraphQLContext, "Expected GraphQLContext in BatchLoaderEnvironment");
return ((GraphQLContext) context).getOrDefault(CONTEXT_VIEW_KEY, Context.empty());
}
/**
* Use the given accessor to extract ThreadLocal values and save them in a
* sub-map in the given {@link Context}, so those can be restored later

3
spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java
View File

@@ -18,6 +18,7 @@ package org.springframework.graphql.data.method.annotation.support;
import java.util.List;
import java.util.Map;
import graphql.GraphQLContext;
import graphql.schema.DataFetcher;
import graphql.schema.idl.RuntimeWiring;
import org.dataloader.BatchLoaderEnvironment;
@@ -61,7 +62,7 @@ public class BatchMappingDetectionTests {
"authorFlux", "authorList", "authorMonoMap", "authorMap", "authorEnvironment");
DataLoaderRegistry registry = new DataLoaderRegistry();
this.batchLoaderRegistry.registerDataLoaders(registry);
this.batchLoaderRegistry.registerDataLoaders(registry, GraphQLContext.newContext().build());
assertThat(registry.getDataLoadersMap()).containsOnlyKeys(
"Book.authorFlux", "Book.authorList", "Book.authorMonoMap", "Book.authorMap", "Book.authorEnvironment");
}

3
spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataLoaderArgumentResolverTests.java
View File

@@ -18,6 +18,7 @@ package org.springframework.graphql.data.method.annotation.support;
import java.lang.reflect.Method;
import java.util.function.Consumer;
import graphql.GraphQLContext;
import graphql.schema.DataFetchingEnvironment;
import graphql.schema.DataFetchingEnvironmentImpl;
import org.dataloader.DataLoader;
@@ -117,7 +118,7 @@ public class DataLoaderArgumentResolverTests {
registryConsumer.accept(batchLoaderRegistry);
DataLoaderRegistry registry = DataLoaderRegistry.newRegistry().build();
batchLoaderRegistry.registerDataLoaders(registry);
batchLoaderRegistry.registerDataLoaders(registry, GraphQLContext.newContext().build());
return DataFetchingEnvironmentImpl.newDataFetchingEnvironment().dataLoaderRegistry(registry).build();
}

56
spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java
View File

@@ -16,7 +16,11 @@
package org.springframework.graphql.execution;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import graphql.ExecutionInput;
import graphql.GraphQLContext;
import org.dataloader.DataLoader;
import org.dataloader.DataLoaderRegistry;
import org.dataloader.stats.NoOpStatisticsCollector;
@@ -24,8 +28,11 @@ import org.dataloader.stats.StatisticsCollector;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.graphql.Book;
import org.springframework.graphql.BookSource;
import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat;
@@ -40,24 +47,49 @@ public class DefaultBatchLoaderRegistryTests {
@Test
void batchLoader() {
this.batchLoaderRegistry.forTypePair(String.class, Book.class).registerBatchLoader((keys, environment) -> Flux.empty());
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry);
void batchLoader() throws Exception {
AtomicReference<String> valueRef = new AtomicReference<>();
this.batchLoaderRegistry.forTypePair(Long.class, Book.class)
.withOptions(options -> options.setBatchingEnabled(false)) // DataLoader invoked immediately
.registerBatchLoader((ids, environment) ->
Flux.deferContextual(contextView -> {
valueRef.set(contextView.get("key"));
return Flux.fromIterable(ids).map(BookSource::getBook);
}));
GraphQLContext graphQLContext = initGraphQLContext(Context.of("key", "value"));
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry, graphQLContext);
Map<String, DataLoader<?, ?>> map = this.dataLoaderRegistry.getDataLoadersMap();
assertThat(map).hasSize(1).containsKey(Book.class.getName());
// Invoke DataLoader to check the context
((DataLoader<Long, Book>) map.get(Book.class.getName())).load(1L).get();
assertThat(valueRef.get()).isEqualTo("value");
}
@Test
void mappedBatchLoader() {
this.batchLoaderRegistry
.forTypePair(String.class, Book.class)
.registerMappedBatchLoader((keys, environment) -> Mono.empty());
void mappedBatchLoader() throws Exception {
AtomicReference<String> valueRef = new AtomicReference<>();
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry);
this.batchLoaderRegistry.forTypePair(Long.class, Book.class)
.withOptions(options -> options.setBatchingEnabled(false)) // DataLoader invoked immediately
.registerMappedBatchLoader((ids, environment) ->
Mono.deferContextual(contextView -> {
valueRef.set(contextView.get("key"));
return Flux.fromIterable(ids).map(BookSource::getBook).collectMap(Book::getId, Function.identity());
}));
GraphQLContext graphQLContext = initGraphQLContext(Context.of("key", "value"));
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry, graphQLContext);
Map<String, DataLoader<?, ?>> map = this.dataLoaderRegistry.getDataLoadersMap();
assertThat(map).hasSize(1).containsKey(Book.class.getName());
// Invoke DataLoader to check the context
((DataLoader<Long, Book>) map.get(Book.class.getName())).load(1L).get();
assertThat(valueRef.get()).isEqualTo("value");
}
@Test
@@ -69,11 +101,17 @@ public class DefaultBatchLoaderRegistryTests {
.withOptions(options -> options.setStatisticsCollector(() -> collector))
.registerBatchLoader((keys, environment) -> Flux.empty());
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry);
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry, GraphQLContext.newContext().build());
Map<String, DataLoader<?, ?>> map = dataLoaderRegistry.getDataLoadersMap();
assertThat(map).hasSize(1).containsKey(name);
assertThat(map.get(name).getStatistics()).isSameAs(collector.getStatistics());
}
private GraphQLContext initGraphQLContext(ContextView context) {
ExecutionInput executionInput = ExecutionInput.newExecutionInput().query("").build();
ReactorContextManager.setReactorContext(context, executionInput);
return executionInput.getGraphQLContext();
}
}