From 6849d3bf38e3b9a5f008d4017e56ca66f82d2db0 Mon Sep 17 00:00:00 2001
From: Brian Clozel <brian.clozel@broadcom.com>
Date: Thu, 17 Oct 2024 12:21:52 +0200
Subject: [PATCH] Propagate context to Dataloader in `@BatchMapping`

Prior to this commit, `@BatchMapping` controller methods would be
automatically registered as date fetchers delegating to data loader
calls. Those calls would not include the current local context or main
context. As a result, injecting the `BatchLoaderEnvironment` in the
controller method signature would not contain the `getKeyContext()`.

This commit ensures that dataloader calls not only use the current
source, but also the current local context/main context so that it will
be present in the key contexts map.

Fixes gh-1071
---
 .../AnnotatedControllerConfigurer.java        |  2 +-
 .../support/BatchMappingInvocationTests.java  | 46 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

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 58ead114..a49a40c6 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,7 @@ public class AnnotatedControllerConfigurer implements ApplicationContextAware, I
 		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());
+			return dataLoader.load(env.getSource(), (env.getLocalContext() != null) ? env.getLocalContext() : env.getGraphQlContext());
 		}
 	}
 
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 45282685..1f99da0b 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,6 +24,9 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import graphql.GraphQLContext;
+import org.dataloader.BatchLoaderEnvironment;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -126,12 +129,33 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 		}
 	}
 
+	@Test
+	void shouldBindKeyContextsToEnvironment() {
+		String document = "{ " +
+				"  courses { " +
+				"    id" +
+				"    name" +
+				"    students {" +
+				"      id" +
+				"      firstName" +
+				"      lastName" +
+				"    }" +
+				"  }" +
+				"}";
+
+		Mono<ExecutionGraphQlResponse> responseMono = createGraphQlService(new BatchKeyContextsController()).execute(document);
+
+		List<Course> actualCourses = ResponseHelper.forResponse(responseMono).toList("courses", Course.class);
+		List<Course> courses = Course.allCourses();
+		assertThat(actualCourses).hasSize(courses.size());
+	}
+
 
 	@Controller
 	private static class BatchMonoMapController extends CourseController {
 
 		@BatchMapping
-		public Mono<Map<Course, Person>> instructor(List<Course> courses) {
+		public Mono<Map<Course, Person>> instructor(List<Course> courses, BatchLoaderEnvironment environment) {
 			return Flux.fromIterable(courses).collect(Collectors.toMap(Function.identity(), Course::instructor));
 		}
 
@@ -203,5 +227,23 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 
 	}
 
+	@Controller
+	private static class BatchKeyContextsController extends CourseController {
+
+		@BatchMapping
+		public List<Person> instructor(List<Course> courses, BatchLoaderEnvironment environment) {
+			assertThat(environment.getKeyContexts().keySet()).containsAll(Course.allCourses());
+			assertThat(environment.getKeyContexts().values()).allSatisfy(value -> assertThat(value).isInstanceOf(GraphQLContext.class));
+			return courses.stream().map(Course::instructor).collect(Collectors.toList());
+		}
+
+		@BatchMapping
+		public List<List<Person>> students(List<Course> courses, BatchLoaderEnvironment environment) {
+			assertThat(environment.getKeyContexts().keySet()).containsAll(Course.allCourses());
+			assertThat(environment.getKeyContexts().values()).allSatisfy(value -> assertThat(value).isInstanceOf(GraphQLContext.class));
+			return courses.stream().map(Course::students).collect(Collectors.toList());
+		}
+	}
+
 
 }