From 3bee3e81cc801fc88dd8aae49e66d9c771433011 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Mon, 19 Dec 2022 20:16:13 +0000
Subject: [PATCH] Use ParameterNameDiscoverer with BatchMapping

Closes gh-562
---
 .../support/BatchLoaderHandlerMethod.java     |  7 ++++
 ...ntextValueMethodArgumentResolverTests.java | 33 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java
index 0be60b94..ae72c289 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java
@@ -29,7 +29,9 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.CollectionFactory;
+import org.springframework.core.DefaultParameterNameDiscoverer;
 import org.springframework.core.MethodParameter;
+import org.springframework.core.ParameterNameDiscoverer;
 import org.springframework.graphql.data.method.HandlerMethod;
 import org.springframework.graphql.data.method.InvocableHandlerMethodSupport;
 import org.springframework.graphql.data.method.annotation.ContextValue;
@@ -53,6 +55,9 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 			AnnotatedControllerConfigurer.class.getClassLoader());
 
 
+	private final ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer();
+
+
 	public BatchLoaderHandlerMethod(HandlerMethod handlerMethod, @Nullable Executor executor) {
 		super(handlerMethod, executor);
 	}
@@ -114,6 +119,8 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 	private  <K> Object resolveArgument(
 			MethodParameter parameter, Collection<K> keys, BatchLoaderEnvironment environment) {
 
+		parameter.initParameterNameDiscovery(this.parameterNameDiscoverer);
+
 		Class<?> parameterType = parameter.getParameterType();
 
 		if (Collection.class.isAssignableFrom(parameterType)) {
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java
index d4c0a434..c56efcd0 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java
@@ -17,12 +17,15 @@ package org.springframework.graphql.data.method.annotation.support;
 
 import java.lang.reflect.Method;
 import java.time.Duration;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Optional;
 import java.util.function.BiConsumer;
 
 import graphql.GraphQLContext;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.DataFetchingEnvironmentImpl;
+import org.dataloader.BatchLoaderEnvironment;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -30,7 +33,9 @@ import reactor.test.StepVerifier;
 import org.springframework.core.DefaultParameterNameDiscoverer;
 import org.springframework.core.MethodParameter;
 import org.springframework.core.annotation.SynthesizingMethodParameter;
+import org.springframework.graphql.Author;
 import org.springframework.graphql.Book;
+import org.springframework.graphql.BookSource;
 import org.springframework.graphql.data.method.HandlerMethod;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.data.method.annotation.ContextValue;
@@ -134,6 +139,28 @@ public class ContextValueMethodArgumentResolverTests {
 		StepVerifier.create((Mono<String>) handlerMethod.invoke(environment)).verifyComplete();
 	}
 
+	@Test // gh-562
+	void resolveFromParameterNameWithBatchMapping() throws Exception {
+
+		TestController controller = new TestController();
+
+		BatchLoaderHandlerMethod handlerMethod = new BatchLoaderHandlerMethod(
+				new HandlerMethod(controller,
+						TestController.class.getMethod("getAuthors", List.class, Long.class)), null);
+
+		GraphQLContext context = new GraphQLContext.Builder().build();
+		context.put("id", 123L);
+
+		BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment().context(context).build();
+		List<Book> keys = Arrays.asList(BookSource.getBook(1L), BookSource.getBook(2L), BookSource.getBook(3L));
+
+		StepVerifier.create(handlerMethod.invokeForIterable(keys, environment))
+				.expectNextCount(3)
+				.verifyComplete();
+
+		assertThat(controller.savedId).isEqualTo(context.get("id"));
+	}
+
 	@Nullable
 	private Object resolveValue(
 			@Nullable GraphQLContext localContext, @Nullable GraphQLContext graphQLContext, int index) {
@@ -166,11 +193,17 @@ public class ContextValueMethodArgumentResolverTests {
 
 	private static class TestController {
 
+		private Long savedId;
+
 		@Nullable
 		public String handleMono(@ContextValue Mono<String> stringMono) {
 			return stringMono.block(Duration.ofSeconds(1));
 		}
 
+		public List<Author> getAuthors(List<Book> books, @ContextValue Long id) {
+			this.savedId = id;
+			return books.stream().map(Book::getAuthor).toList();
+		}
 	}
 
 }