Pass localContext only if available

See gh-1066

Author: rstoyanchev

spring-graphql-docs/modules/ROOT/pages/controllers.adoc

@@ -743,7 +743,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.
 
 |===
 

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java

@@ -603,7 +603,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()));
 		}
 
 		@Override

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<ExecutionGraphQlResponse> responseMono = createGraphQlService(new BatchKeyContextsController()).execute(document);
+		Mono<ExecutionGraphQlResponse> responseMono = createGraphQlService(
+				BatchKeyContextsController.class, new BatchKeyContextsController()).execute(document);
 
 		List<Course> actualCourses = ResponseHelper.forResponse(responseMono).toList("courses", Course.class);
 		List<Course> courses = Course.allCourses();
@@ -228,7 +232,14 @@ public Callable<Map<Course, List<Person>>> students(List<Course> courses) {
 	}
 
 	@Controller
-	private static class BatchKeyContextsController extends CourseController {
+	private static class BatchKeyContextsController {
+
+		@QueryMapping
+		public DataFetcherResult<Collection<Course>> courses() {
+			return DataFetcherResult.<Collection<Course>>newResult().data(courseMap.values())
+					.localContext(GraphQLContext.newContext().build())
+					.build();
+		}
 
 		@BatchMapping
 		public List<Person> instructor(List<Course> courses, BatchLoaderEnvironment environment) {

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 <T> TestExecutionGraphQlService createGraphQlService(Class<T> 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();