From cf1c99f489fba9848ffbb44aca4369a38a4b2f90 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Thu, 13 Mar 2025 12:06:52 +0000 Subject: [PATCH] Support DataLoader in EntityMapping methods Closes gh-1095 --- .../modules/ROOT/pages/federation.adoc | 39 +++++++++++++++--- .../federation/FederationSchemaFactory.java | 2 + .../EntityMappingInvocationTests.java | 40 +++++++++++++++++++ 3 files changed, 76 insertions(+), 5 deletions(-) diff --git a/spring-graphql-docs/modules/ROOT/pages/federation.adoc b/spring-graphql-docs/modules/ROOT/pages/federation.adoc index 3bd8e5f8..40bfad3e 100644 --- a/spring-graphql-docs/modules/ROOT/pages/federation.adoc +++ b/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(@Argument int id, DataLoader dataLoader) { // <2> + return dataLoader.load(id); + } + + @BatchMapping + public Map author(List 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` +| 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. diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java index f3495838..c6182917 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java @@ -48,6 +48,7 @@ import org.springframework.graphql.data.method.annotation.support.Authentication 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 @@ public final class FederationSchemaFactory // Type based resolvers.addResolver(new DataFetchingEnvironmentMethodArgumentResolver()); + resolvers.addResolver(new DataLoaderMethodArgumentResolver()); if (springSecurityPresent) { ApplicationContext context = obtainApplicationContext(); resolvers.addResolver(new PrincipalMethodArgumentResolver()); diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java index 08b60768..e51e39fc 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java @@ -19,16 +19,19 @@ package org.springframework.graphql.data.federation; 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 @@ public class EntityMappingInvocationTests { assertError(helper, 2, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty"); } + @Test + void dataLoader() { + Map 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 class EntityMappingInvocationTests { } + @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(@Argument int id, DataLoader dataLoader) { + return dataLoader.load(id); + } + + @BatchMapping + public Flux author(List books) { + return Flux.fromIterable(books).map(book -> BookSource.getBook(book.getId()).getAuthor()); + } + } + + @SuppressWarnings("unused") @Controller private static class EmptyController {