Support DataLoader in EntityMapping methods

Closes gh-1095

Author: rstoyanchev

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

@@ -80,11 +80,8 @@ xref:federation.adoc#federation.entity-mapping.signature[Method Signature] for s
 method argument and return value types.
 <2> `@SchemaMapping` methods can be used for the rest of the graph.
 
-An `@EntityMapping` method can batch load federated entities of a given type. To do that,
-declare the `@Argument` method parameter as a list, and return the corresponding entity
-instances as a list in the same order.
-
-For example:
+You can load federated entities of the same type together by accepting a `List` of id's,
+and returning a `List` or `Flux` of entities:
 
 [source,java,indent=0,subs="verbatim,quotes"]
 ----
@@ -108,6 +105,35 @@ look up the correct value in the "representation" input map. You can also set th
 argument name through the annotation.
 <2> `@BatchMapping` methods can be used for the rest of the graph.
 
+You can load federated entities with a `DataLoader`:
+
+[source,java,indent=0,subs="verbatim,quotes"]
+----
+	@Controller
+	private static class BookController {
+
+		@Autowired
+		public DataLoaderBookController(BatchLoaderRegistry registry) { // <1>
+			registry.forTypePair(Integer.class, Book.class).registerBatchLoader((bookIds, environment) -> {
+				// load entities...
+			});
+		}
+
+		@EntityMapping
+		public Future<Book> book(@Argument int id, DataLoader<Integer, Book> dataLoader) { // <2>
+			return dataLoader.load(id);
+		}
+
+		@BatchMapping
+		public Map<Book, Author> author(List<Book> books) { // <3>
+			// ...
+		}
+}
+----
+
+<1> Register a batch loader for the federated entity type.
+<2> Declare a `DataLoader` argument to the `@EntityMapping` method.
+<3> `@BatchMapping` methods can be used for the rest of the graph.
 
 
 [[federation.entity-mapping.signature]]
@@ -153,6 +179,9 @@ Entity mapping methods support the following arguments:
 | `DataFetchingEnvironment`
 | For direct access to the underlying `DataFetchingEnvironment`.
 
+| `DataLoader<I, E>`
+| To load federated entities with a `DataLoader` where `I` is the id type, and `E` is the entity type.
+
 |===
 
 `@EntityMapping` methods can return `Mono`, `CompletableFuture`, `Callable`, or the actual entity.

spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java

@@ -48,6 +48,7 @@
 import org.springframework.graphql.data.method.annotation.support.ContextValueMethodArgumentResolver;
 import org.springframework.graphql.data.method.annotation.support.ContinuationHandlerMethodArgumentResolver;
 import org.springframework.graphql.data.method.annotation.support.DataFetchingEnvironmentMethodArgumentResolver;
+import org.springframework.graphql.data.method.annotation.support.DataLoaderMethodArgumentResolver;
 import org.springframework.graphql.data.method.annotation.support.LocalContextValueMethodArgumentResolver;
 import org.springframework.graphql.data.method.annotation.support.PrincipalMethodArgumentResolver;
 import org.springframework.graphql.execution.ClassNameTypeResolver;
@@ -125,6 +126,7 @@ protected HandlerMethodArgumentResolverComposite initArgumentResolvers() {
 
 		// Type based
 		resolvers.addResolver(new DataFetchingEnvironmentMethodArgumentResolver());
+		resolvers.addResolver(new DataLoaderMethodArgumentResolver());
 		if (springSecurityPresent) {
 			ApplicationContext context = obtainApplicationContext();
 			resolvers.addResolver(new PrincipalMethodArgumentResolver());

spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java

@@ -19,16 +19,19 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Future;
 
 import graphql.GraphQLError;
 import graphql.GraphqlErrorBuilder;
 import graphql.schema.DataFetchingEnvironment;
+import org.dataloader.DataLoader;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.core.io.Resource;
@@ -176,6 +179,19 @@ void batchingWithoutResult(Class<?> controllerClass) {
 		assertError(helper, 2, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
 	}
 
+	@Test
+	void dataLoader() {
+		Map<String, Object> variables =
+				Map.of("representations", List.of(
+						Map.of("__typename", "Book", "id", "3"),
+						Map.of("__typename", "Book", "id", "5")));
+
+		ResponseHelper helper = executeWith(DataLoaderBookController.class, variables);
+
+		assertAuthor(0, "Joseph", "Heller", helper);
+		assertAuthor(1, "George", "Orwell", helper);
+	}
+
 	@Test
 	void unmappedEntity() {
 		assertThatIllegalStateException().isThrownBy(() -> executeWith(EmptyController.class, Map.of()))
@@ -306,6 +322,30 @@ public GraphQLError handle(IllegalArgumentException ex, DataFetchingEnvironment
 	}
 
 
+	@SuppressWarnings("unused")
+	@Controller
+	private static class DataLoaderBookController {
+
+		@Autowired
+		public DataLoaderBookController(BatchLoaderRegistry batchLoaderRegistry) {
+			batchLoaderRegistry.forTypePair(Integer.class, Book.class)
+					.registerBatchLoader((ids, env) ->
+							Flux.fromIterable(ids).map(id -> new Book((long) id, null, (Long) null)));
+		}
+
+		@Nullable
+		@EntityMapping
+		public Future<Book> book(@Argument int id, DataLoader<Integer, Book> dataLoader) {
+			return dataLoader.load(id);
+		}
+
+		@BatchMapping
+		public Flux<Author> author(List<Book> books) {
+			return Flux.fromIterable(books).map(book -> BookSource.getBook(book.getId()).getAuthor());
+		}
+	}
+
+
 	@SuppressWarnings("unused")
 	@Controller
 	private static class EmptyController {