diff --git a/spring-graphql-docs/modules/ROOT/pages/controllers.adoc b/spring-graphql-docs/modules/ROOT/pages/controllers.adoc index 92df154b..04165e52 100644 --- a/spring-graphql-docs/modules/ROOT/pages/controllers.adoc +++ b/spring-graphql-docs/modules/ROOT/pages/controllers.adoc @@ -646,7 +646,15 @@ Batch mapping methods support the following arguments: | `BatchLoaderEnvironment` | The environment that is available in GraphQL Java to a -`org.dataloader.BatchLoaderWithContext`. + `org.dataloader.BatchLoaderWithContext`. + + The `context` property of `BatchLoaderEnvironment` returns the same + `GraphQLContext` instance that is also available to `@SchemaMapping` + methods through the `DataFetchingEnvironment`. + + The `keyContexts` property of `BatchLoaderEnvironment` returns the + localContext obtained from the `DataFetchingEnvironment` when the + `DataLoader` was called for each key. |=== diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java index a49a40c6..10793002 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java @@ -755,7 +755,9 @@ public ResolvableType getReturnType() { public Object get(DataFetchingEnvironment env) { DataLoader dataLoader = env.getDataLoaderRegistry().getDataLoader(this.dataLoaderKey); Assert.state(dataLoader != null, "No DataLoader for key '" + this.dataLoaderKey + "'"); - return dataLoader.load(env.getSource(), (env.getLocalContext() != null) ? env.getLocalContext() : env.getGraphQlContext()); + return ((env.getLocalContext() != null) ? + dataLoader.load(env.getSource(), env.getLocalContext()) : + dataLoader.load(env.getSource())); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingInvocationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingInvocationTests.java index 1f99da0b..bf66fef9 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingInvocationTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingInvocationTests.java @@ -16,6 +16,7 @@ package org.springframework.graphql.data.method.annotation.support; +import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Set; @@ -25,6 +26,7 @@ import java.util.stream.Stream; import graphql.GraphQLContext; +import graphql.execution.DataFetcherResult; import org.dataloader.BatchLoaderEnvironment; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -36,6 +38,7 @@ import org.springframework.graphql.ExecutionGraphQlResponse; import org.springframework.graphql.ResponseHelper; import org.springframework.graphql.data.method.annotation.BatchMapping; +import org.springframework.graphql.data.method.annotation.QueryMapping; import org.springframework.stereotype.Controller; import static org.assertj.core.api.Assertions.assertThat; @@ -143,7 +146,8 @@ void shouldBindKeyContextsToEnvironment() { " }" + "}"; - Mono responseMono = createGraphQlService(new BatchKeyContextsController()).execute(document); + Mono responseMono = createGraphQlService( + BatchKeyContextsController.class, new BatchKeyContextsController()).execute(document); List actualCourses = ResponseHelper.forResponse(responseMono).toList("courses", Course.class); List courses = Course.allCourses(); @@ -228,7 +232,14 @@ public Callable>> students(List courses) { } @Controller - private static class BatchKeyContextsController extends CourseController { + private static class BatchKeyContextsController { + + @QueryMapping + public DataFetcherResult> courses() { + return DataFetcherResult.>newResult().data(courseMap.values()) + .localContext(GraphQLContext.newContext().build()) + .build(); + } @BatchMapping public List instructor(List courses, BatchLoaderEnvironment environment) { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingTestSupport.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingTestSupport.java index 38e02c2b..28062a5f 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingTestSupport.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingTestSupport.java @@ -82,10 +82,14 @@ public class BatchMappingTestSupport { protected TestExecutionGraphQlService createGraphQlService(CourseController controller) { + return createGraphQlService(CourseController.class, controller); + } + + protected TestExecutionGraphQlService createGraphQlService(Class controllerClass, T controller) { BatchLoaderRegistry registry = new DefaultBatchLoaderRegistry(); AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); - context.registerBean(CourseController.class, () -> controller); + context.registerBean(controllerClass, () -> controller); context.registerBean(BatchLoaderRegistry.class, () -> registry); context.refresh();