From 5aaf5b1df4c067eb0d92bde18321584784d57a6c Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Wed, 3 Nov 2021 17:08:19 +0000 Subject: [PATCH] Context for batch loading functions 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 --- .../src/docs/asciidoc/index.adoc | 16 +++--- .../execution/BatchLoaderRegistry.java | 14 ++++- .../execution/DataLoaderRegistrar.java | 13 ++++- .../execution/DefaultBatchLoaderRegistry.java | 28 ++++++++-- .../execution/ExecutionGraphQlService.java | 20 +++---- .../execution/ReactorContextManager.java | 13 +++++ .../support/BatchMappingDetectionTests.java | 3 +- .../DataLoaderArgumentResolverTests.java | 3 +- .../DefaultBatchLoaderRegistryTests.java | 56 ++++++++++++++++--- 9 files changed, 128 insertions(+), 38 deletions(-) diff --git a/spring-graphql-docs/src/docs/asciidoc/index.adoc b/spring-graphql-docs/src/docs/asciidoc/index.adoc index 4e1a407f..6dab6f02 100644 --- a/spring-graphql-docs/src/docs/asciidoc/index.adoc +++ b/spring-graphql-docs/src/docs/asciidoc/index.adoc @@ -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 <> to access the -registered loader. Annotated controllers also support a -<> 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. + - <> for `@SchemaMapping` +methods to access the `DataLoader` for the field type. + - <> data controller methods that provide a +shortcut and avoid the need to use `DataLoader` directly. The Spring Boot starter declares a <>, so that applications can diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java index 1f273153..2c92f4aa 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java @@ -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}. + *

Note: 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 withOptions(Consumer 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}. + *

Note: 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 */ diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java index c9f896eb..bbeed39c 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java @@ -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); } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java index 24c97cab..9dee12bf 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java @@ -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 implements BatchLoaderWithContext { private final String name; @@ -168,11 +178,17 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { @Override public CompletionStage> load(List 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 implements MappedBatchLoaderWithContext { private final String name; @@ -200,8 +216,10 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { @Override public CompletionStage> load(Set keys, BatchLoaderEnvironment environment) { - return this.loader.apply(keys, environment).toFuture(); + ContextView contextView = ReactorContextManager.getReactorContext(environment); + return this.loader.apply(keys, environment).contextWrite(contextView).toFuture(); } + } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java index 5c3c9d61..64c8f70b 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java @@ -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 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; } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java index e7b57d3c..2815c985 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java @@ -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 diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java index ee5dce3e..8a32ab55 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java @@ -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"); } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataLoaderArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataLoaderArgumentResolverTests.java index 782e9f56..5a716537 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataLoaderArgumentResolverTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataLoaderArgumentResolverTests.java @@ -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(); } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java index 91a1f429..9b853ec2 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java @@ -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 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> map = this.dataLoaderRegistry.getDataLoadersMap(); assertThat(map).hasSize(1).containsKey(Book.class.getName()); + + // Invoke DataLoader to check the context + ((DataLoader) 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 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> map = this.dataLoaderRegistry.getDataLoadersMap(); assertThat(map).hasSize(1).containsKey(Book.class.getName()); + + // Invoke DataLoader to check the context + ((DataLoader) 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> 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(); + } + }