Add DataLoaderMethodArgumentResolver

This commit adds support for DataLoader<K,V> arguments on annotated
handler methods, where the key for the lookup is the class name of the
value type or the parameter name.

See gh-63

graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/GraphQlWebFluxAutoConfiguration.java

@@ -76,9 +76,9 @@ public class GraphQlWebFluxAutoConfiguration {
 
 	@Bean
 	public AnnotatedDataFetcherConfigurer annotatedDataFetcherConfigurer(ServerCodecConfigurer configurer) {
-		AnnotatedDataFetcherConfigurer registrar = new AnnotatedDataFetcherConfigurer();
-		registrar.setServerCodecConfigurer(configurer);
-		return registrar;
+		AnnotatedDataFetcherConfigurer dataFetcherConfigurer = new AnnotatedDataFetcherConfigurer();
+		dataFetcherConfigurer.setServerCodecConfigurer(configurer);
+		return dataFetcherConfigurer;
 	}
 
 	@Bean

graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/GraphQlWebMvcAutoConfiguration.java

@@ -84,9 +84,9 @@ public class GraphQlWebMvcAutoConfiguration {
 
 	@Bean
 	public AnnotatedDataFetcherConfigurer annotatedDataFetcherConfigurer(HttpMessageConverters converters) {
-		AnnotatedDataFetcherConfigurer registrar = new AnnotatedDataFetcherConfigurer();
-		registrar.setJsonMessageConverter(getJsonConverter(converters));
-		return registrar;
+		AnnotatedDataFetcherConfigurer dataFetcherConfigurer = new AnnotatedDataFetcherConfigurer();
+		dataFetcherConfigurer.setJsonMessageConverter(getJsonConverter(converters));
+		return dataFetcherConfigurer;
 	}
 
 	@SuppressWarnings("unchecked")

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

@@ -161,6 +161,9 @@ public class AnnotatedDataFetcherConfigurer
 		this.argumentResolvers.addResolver(initInputArgumentMethodArgumentResolver());
 		this.argumentResolvers.addResolver(new ArgumentMapMethodArgumentResolver());
 		this.argumentResolvers.addResolver(new DataFetchingEnvironmentMethodArgumentResolver());
+		this.argumentResolvers.addResolver(new DataLoaderMethodArgumentResolver());
+
+		// This works as a fallback, after all other resolvers
 		this.argumentResolvers.addResolver(new SourceMethodArgumentResolver());
 	}
 

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

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2002-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.data.method.annotation.support;
+
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
+
+import graphql.schema.DataFetchingEnvironment;
+import org.dataloader.DataLoader;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.core.ResolvableType;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * Resolver that retrieves arguments of type {@link DataLoader} from the
+ * {@link DataFetchingEnvironment}.
+ *
+ * <p>The {@code DataLoader} is looked up by deriving the key using one of the following:
+ * <ol>
+ * <li>The full name of the value type from the DataLoader generic types.</li>
+ * <li>The method parameter name.</li>
+ * </ol>
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public class DataLoaderMethodArgumentResolver implements HandlerMethodArgumentResolver {
+	
+	@Override
+	public boolean supportsParameter(MethodParameter parameter) {
+		return parameter.getParameterType().equals(DataLoader.class);
+	}
+
+	@Override
+	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment environment) {
+		DataLoader<Object, Object> dataLoader = null;
+
+		Class<?> valueType = getValueType(parameter);
+		if (valueType != null) {
+			dataLoader = environment.getDataLoader(valueType.getName());
+		}
+
+		String parameterName = null;
+		if (dataLoader == null) {
+			parameterName = parameter.getParameterName();
+			if (parameterName != null) {
+				dataLoader = environment.getDataLoader(parameterName);
+			}
+		}
+
+		if (dataLoader == null) {
+			String message = getErrorMessage(parameter, environment, valueType, parameterName);
+			throw new IllegalArgumentException(message);
+		}
+
+		return dataLoader;
+	}
+
+	@Nullable
+	private Class<?> getValueType(MethodParameter param) {
+		Assert.isAssignable(DataLoader.class, param.getParameterType());
+		Type genericType = param.getGenericParameterType();
+		if (genericType instanceof ParameterizedType) {
+			ParameterizedType parameterizedType = (ParameterizedType) genericType;
+			if (parameterizedType.getActualTypeArguments().length == 2) {
+				Type valueType = parameterizedType.getActualTypeArguments()[1];
+				return (valueType instanceof Class ?
+						(Class<?>) valueType : ResolvableType.forType(valueType).resolve());
+			}
+		}
+		return null;
+	}
+
+	private String getErrorMessage(
+			MethodParameter parameter, DataFetchingEnvironment environment,
+			@Nullable Class<?> valueType, @Nullable String parameterName) {
+
+		String message = "Cannot resolve DataLoader for parameter" +
+				(parameterName != null ? " '" + parameterName + "'" : "[" + parameter.getParameterIndex() + "]" ) +
+				" in method " + parameter.getMethod().toGenericString() + ". ";
+
+		if (valueType == null) {
+			message += "If the batch loader was registered without a name, " +
+					"then declaring the DataLoader argument with generic types should help " +
+					"to look up the DataLoader based on the value type name.";
+		}
+		else if (parameterName == null) {
+			message += "If the batch loader was registered with a name, " +
+					"then compiling with \"-parameters\" should help " +
+					"to look up the DataLoader based on the parameter name.";
+		}
+		else {
+			message += "Neither the name of the declared value type '" + valueType + "' " +
+					"nor the method parameter name '" + parameterName + "' match to any DataLoader. " +
+					"The DataLoaderRegistry contains: " + environment.getDataLoaderRegistry().getKeys();
+		}
+
+		return message;
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedDataFetcherDetectionTests.java

@@ -42,7 +42,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  *
  * @author Rossen Stoyanchev
  */
-public class AnnotatedDataFetcherConfigurerTests {
+public class AnnotatedDataFetcherDetectionTests {
 
 	@Test
 	void registerWithDefaultCoordinates() {

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedDataFetcherInvocationTests.java

@@ -17,28 +17,36 @@ package org.springframework.graphql.data.method.annotation.support;
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 
-import graphql.ExecutionInput;
 import graphql.ExecutionResult;
-import graphql.GraphQL;
 import graphql.GraphQLContext;
 import graphql.schema.DataFetchingEnvironment;
+import org.dataloader.DataLoader;
 import org.junit.jupiter.api.Test;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.test.StepVerifier;
 
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.graphql.Author;
 import org.springframework.graphql.Book;
 import org.springframework.graphql.BookCriteria;
 import org.springframework.graphql.BookSource;
+import org.springframework.graphql.GraphQlService;
+import org.springframework.graphql.RequestInput;
 import org.springframework.graphql.data.method.annotation.Argument;
 import org.springframework.graphql.data.method.annotation.MutationMapping;
 import org.springframework.graphql.data.method.annotation.QueryMapping;
 import org.springframework.graphql.data.method.annotation.SchemaMapping;
 import org.springframework.graphql.data.method.annotation.SubscriptionMapping;
+import org.springframework.graphql.execution.BatchLoaderRegistry;
+import org.springframework.graphql.execution.DefaultBatchLoaderRegistry;
+import org.springframework.graphql.execution.ExecutionGraphQlService;
 import org.springframework.graphql.execution.GraphQlSource;
 import org.springframework.http.codec.ServerCodecConfigurer;
 import org.springframework.stereotype.Controller;
@@ -64,7 +72,9 @@ public class AnnotatedDataFetcherInvocationTests {
 				"  }" +
 				"}";
 
-		ExecutionResult result = initGraphQl(BookController.class).execute(query);
+		ExecutionResult result = initGraphQlService(BookController.class)
+				.execute(new RequestInput(query, null, null))
+				.block();
 
 		assertThat(result.getErrors()).isEmpty();
 		Map<String, Object> data = result.getData();
@@ -88,7 +98,9 @@ public class AnnotatedDataFetcherInvocationTests {
 				"  }" +
 				"}";
 
-		ExecutionResult result = initGraphQl(BookController.class).execute(query);
+		ExecutionResult result = initGraphQlService(BookController.class)
+				.execute(new RequestInput(query, null, null))
+				.block();
 
 		assertThat(result.getErrors()).isEmpty();
 		Map<String, Object> data = result.getData();
@@ -110,8 +122,16 @@ public class AnnotatedDataFetcherInvocationTests {
 				"  }" +
 				"}";
 
-		ExecutionInput input = ExecutionInput.newExecutionInput().query(query).build();
-		ExecutionResult result = initGraphQl(BookController.class).execute(input);
+		AtomicReference<GraphQLContext> contextRef = new AtomicReference<>();
+		RequestInput requestInput = new RequestInput(query, null, null);
+		requestInput.configureExecutionInput((executionInput, builder) -> {
+			contextRef.set(executionInput.getGraphQLContext());
+			return executionInput;
+		});
+
+		ExecutionResult result = initGraphQlService(BookController.class)
+				.execute(requestInput)
+				.block();
 
 		assertThat(result.getErrors()).isEmpty();
 		Map<String, Object> data = result.getData();
@@ -122,7 +142,7 @@ public class AnnotatedDataFetcherInvocationTests {
 		assertThat(author.get("firstName")).isEqualTo("George");
 		assertThat(author.get("lastName")).isEqualTo("Orwell");
 
-		assertThat(input.getGraphQLContext().<String>get("key")).isEqualTo("value");
+		assertThat(contextRef.get().<String>get("key")).isEqualTo("value");
 	}
 
 	@Test
@@ -135,7 +155,9 @@ public class AnnotatedDataFetcherInvocationTests {
 				"  }" +
 				"}";
 
-		ExecutionResult result = initGraphQl(BookController.class).execute(operation);
+		ExecutionResult result = initGraphQlService(BookController.class)
+				.execute(new RequestInput(operation, null, null))
+				.block();
 
 		assertThat(result.getErrors()).isEmpty();
 		Map<String, Object> data = result.getData();
@@ -156,7 +178,9 @@ public class AnnotatedDataFetcherInvocationTests {
 				"  }" +
 				"}";
 
-		ExecutionResult result = initGraphQl(BookController.class).execute(operation);
+		ExecutionResult result = initGraphQlService(BookController.class)
+				.execute(new RequestInput(operation, null, null))
+				.block();
 
 		assertThat(result.getErrors()).isEmpty();
 		Publisher<ExecutionResult> publisher = result.getData();
@@ -180,22 +204,12 @@ public class AnnotatedDataFetcherInvocationTests {
 	}
 
 
-	private GraphQL initGraphQl(Class<?> beanClass) {
+	private ExecutionGraphQlService initGraphQlService(Class<?> beanClass) {
 		AnnotationConfigApplicationContext applicationContext = new AnnotationConfigApplicationContext();
-		applicationContext.registerBean(beanClass);
+		applicationContext.register(TestConfig.class);
 		applicationContext.refresh();
 
-		AnnotatedDataFetcherConfigurer configurer = new AnnotatedDataFetcherConfigurer();
-		configurer.setApplicationContext(applicationContext);
-		configurer.setServerCodecConfigurer(ServerCodecConfigurer.create());
-		configurer.afterPropertiesSet();
-
-		GraphQlSource graphQlSource = GraphQlSource.builder()
-				.schemaResources(new ClassPathResource("books/schema.graphqls"))
-				.configureRuntimeWiring(configurer::configure)
-				.build();
-
-		return graphQlSource.graphQl();
+		return applicationContext.getBean(ExecutionGraphQlService.class);
 	}
 
 	@SuppressWarnings("unchecked")
@@ -204,9 +218,52 @@ public class AnnotatedDataFetcherInvocationTests {
 	}
 
 
+	@Configuration
+	static class TestConfig {
+
+		@Bean
+		public BookController bookController() {
+			return new BookController(batchLoaderRegistry());
+		}
+
+		@Bean
+		public GraphQlService graphQlService(GraphQlSource graphQlSource) {
+			ExecutionGraphQlService service = new ExecutionGraphQlService(graphQlSource);
+			service.addDataLoaderRegistrar(batchLoaderRegistry());
+			return service;
+		}
+
+		@Bean
+		public GraphQlSource graphQlSource() {
+			return GraphQlSource.builder()
+					.schemaResources(new ClassPathResource("books/schema.graphqls"))
+					.configureRuntimeWiring(annotatedDataFetcherConfigurer())
+					.build();
+		}
+
+		@Bean
+		public AnnotatedDataFetcherConfigurer annotatedDataFetcherConfigurer() {
+			AnnotatedDataFetcherConfigurer registrar = new AnnotatedDataFetcherConfigurer();
+			registrar.setServerCodecConfigurer(ServerCodecConfigurer.create());
+			return registrar;
+		}
+
+		@Bean
+		public DefaultBatchLoaderRegistry batchLoaderRegistry() {
+			return new DefaultBatchLoaderRegistry();
+		}
+
+	}
+
+
 	@Controller
 	private static class BookController {
 
+		public BookController(BatchLoaderRegistry batchLoaderRegistry) {
+			batchLoaderRegistry.forTypePair(Long.class, Author.class)
+					.registerBatchLoader((ids, env) -> Flux.fromIterable(ids).map(BookSource::getAuthor));
+		}
+
 		@QueryMapping
 		public Book bookById(@Argument Long id) {
 			return BookSource.getBookWithoutAuthor(id);
@@ -218,8 +275,8 @@ public class AnnotatedDataFetcherInvocationTests {
 		}
 
 		@SchemaMapping
-		public Author author(Book book) {
-			return BookSource.getBook(book.getId()).getAuthor();
+		public CompletableFuture<Author> author(Book book, DataLoader<Long, Author> dataLoader) {
+			return dataLoader.load(book.getAuthorId());
 		}
 
 		@QueryMapping

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataLoaderArgumentResolverTests.java

@@ -0,0 +1,144 @@
+/*
+ * Copyright 2002-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.data.method.annotation.support;
+
+import java.lang.reflect.Method;
+import java.util.function.Consumer;
+
+import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DataFetchingEnvironmentImpl;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+
+import org.springframework.core.DefaultParameterNameDiscoverer;
+import org.springframework.core.MethodParameter;
+import org.springframework.graphql.Author;
+import org.springframework.graphql.Book;
+import org.springframework.graphql.execution.BatchLoaderRegistry;
+import org.springframework.graphql.execution.DefaultBatchLoaderRegistry;
+import org.springframework.util.ClassUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Unit tests for {@link DataLoaderMethodArgumentResolver}.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class DataLoaderArgumentResolverTests {
+
+	private static final Method method = ClassUtils.getMethod(
+			DataLoaderArgumentResolverTests.class, "handle", (Class<?>[]) null);
+
+	private final DataLoaderMethodArgumentResolver resolver = new DataLoaderMethodArgumentResolver();
+
+
+	@Test
+	void supportsParameter() {
+		assertThat(this.resolver.supportsParameter(initParameter(0))).isTrue();
+		assertThat(this.resolver.supportsParameter(initParameter(3))).isFalse();
+	}
+
+	@Test
+	void resolveArgument() {
+		DataFetchingEnvironment environment = initEnvironment(registry ->
+				registry.forTypePair(Long.class, Author.class).registerBatchLoader((ids, env) -> Flux.empty()));
+
+		Object argument = this.resolver.resolveArgument(initParameter(0), environment);
+		assertThat(argument).isNotNull();
+	}
+
+	@Test
+	void resolveArgumentViaParameterName() {
+		DataFetchingEnvironment environment = initEnvironment(registry ->
+				registry.forTypePair(Long.class, Author.class)
+						.withName("namedDataLoader")
+						.registerBatchLoader((ids, env) -> Flux.empty()));
+
+		MethodParameter parameter = initParameter(1);
+		parameter.initParameterNameDiscovery(new DefaultParameterNameDiscoverer());
+
+		Object argument = this.resolver.resolveArgument(parameter, environment);
+		assertThat(argument).isNotNull();
+	}
+
+	@Test
+	void resolveArgumentFailureWithoutGenericType() {
+		DataFetchingEnvironment environment = initEnvironment(registry ->
+				registry.forTypePair(Long.class, Author.class).registerBatchLoader((ids, env) -> Flux.empty()));
+
+		assertThatThrownBy(() -> this.resolver.resolveArgument(initParameter(2), environment))
+				.hasMessageContaining("declaring the DataLoader argument with generic types should help");
+	}
+
+	@Test
+	void resolveArgumentFailureWithoutParameterName() {
+		DataFetchingEnvironment environment = initEnvironment(registry ->
+				registry.forTypePair(Long.class, Author.class)
+						.withName("namedDataLoader")
+						.registerBatchLoader((ids, env) -> Flux.empty()));
+
+		MethodParameter parameter = initParameter(1);
+		// Skip ParameterNameDiscovery
+
+		assertThatThrownBy(() -> this.resolver.resolveArgument(parameter, environment))
+				.hasMessageContaining("compiling with \"-parameters\" should help");
+	}
+
+	@Test
+	void resolveArgumentFailureNoMatch() {
+		DataFetchingEnvironment environment = initEnvironment(registry ->
+				registry.forTypePair(Long.class, Book.class)
+						.withName("bookDataLoader")
+						.registerBatchLoader((ids, env) -> Flux.empty()));
+
+		MethodParameter parameter = initParameter(0);
+		parameter.initParameterNameDiscovery(new DefaultParameterNameDiscoverer());
+
+		assertThatThrownBy(() -> this.resolver.resolveArgument(parameter, environment))
+				.hasMessageContaining(
+						"Neither the name of the declared value type 'class org.springframework.graphql.Author' " +
+								"nor the method parameter name 'authorDataLoader' match to any DataLoader. " +
+								"The DataLoaderRegistry contains: [bookDataLoader]");
+	}
+
+	private DataFetchingEnvironment initEnvironment(Consumer<BatchLoaderRegistry> registryConsumer) {
+		DefaultBatchLoaderRegistry batchLoaderRegistry = new DefaultBatchLoaderRegistry();
+		registryConsumer.accept(batchLoaderRegistry);
+
+		DataLoaderRegistry registry = DataLoaderRegistry.newRegistry().build();
+		batchLoaderRegistry.registerDataLoaders(registry);
+
+		return DataFetchingEnvironmentImpl.newDataFetchingEnvironment().dataLoaderRegistry(registry).build();
+	}
+
+	private MethodParameter initParameter(int index) {
+		return new MethodParameter(method, index);
+	}
+
+
+	@SuppressWarnings({"unused", "rawtypes"})
+	public void handle(
+			DataLoader<Long, Author> authorDataLoader,
+			DataLoader<Long, Author> namedDataLoader,
+			DataLoader rawDataLoader,
+			DataFetchingEnvironment environment) {
+	}
+
+}