From 3ea4b38ace5e32b1122d9de29f4cd9cabc952efe Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Wed, 12 Jun 2024 09:41:09 +0100 Subject: [PATCH] Support Flux for batched entities query Closes gh-991 --- .../federation/FederationSchemaFactory.java | 4 + .../EntityMappingInvocationTests.java | 74 +++++++++++++++---- 2 files changed, 64 insertions(+), 14 deletions(-) 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 f89a672f..51fb1789 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 @@ -31,6 +31,7 @@ import graphql.schema.GraphQLSchema; import graphql.schema.TypeResolver; import graphql.schema.idl.RuntimeWiring; import graphql.schema.idl.TypeDefinitionRegistry; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import org.springframework.context.ApplicationContext; @@ -189,6 +190,9 @@ public final class FederationSchemaFactory public boolean isBatchHandlerMethod() { MethodParameter type = handlerMethod().getReturnType(); Class paramType = type.getParameterType(); + if (Flux.class.isAssignableFrom(paramType)) { + return true; + } if (Mono.class.isAssignableFrom(paramType) || CompletionStage.class.isAssignableFrom(paramType)) { type = type.nested(); } 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 47572dc5..a03ad867 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 @@ -24,6 +24,8 @@ import graphql.GraphQLError; import graphql.GraphqlErrorBuilder; import graphql.schema.DataFetchingEnvironment; 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; @@ -113,8 +115,9 @@ public class EntityMappingInvocationTests { assertAuthor(6, "George", "Orwell", helper); } - @Test - void batching() { + @ValueSource(classes = {BookListController.class, BookFluxController.class}) + @ParameterizedTest + void batching(Class controllerClass) { Map variables = Map.of("representations", List.of( Map.of("__typename", "Book", "id", "1"), @@ -123,7 +126,7 @@ public class EntityMappingInvocationTests { Map.of("__typename", "Book", "id", "42"), Map.of("__typename", "Book", "id", "53"))); - ResponseHelper helper = executeWith(BookBatchController.class, variables); + ResponseHelper helper = executeWith(controllerClass, variables); assertAuthor(0, "George", "Orwell", helper); assertAuthor(1, "Virginia", "Woolf", helper); @@ -132,30 +135,32 @@ public class EntityMappingInvocationTests { assertAuthor(4, "Vince", "Gilligan", helper); } - @Test - void batchingWithError() { + @ValueSource(classes = {BookListController.class, BookFluxController.class}) + @ParameterizedTest + void batchingWithError(Class controllerClass) { Map variables = Map.of("representations", List.of( Map.of("__typename", "Book", "id", "-97"), Map.of("__typename", "Book", "id", "4"), Map.of("__typename", "Book", "id", "5"))); - ResponseHelper helper = executeWith(BookBatchController.class, variables); + ResponseHelper helper = executeWith(controllerClass, variables); assertError(helper, 0, "BAD_REQUEST", "handled"); assertError(helper, 1, "BAD_REQUEST", "handled"); assertError(helper, 2, "BAD_REQUEST", "handled"); } - @Test - void batchingWithoutResult() { + @ValueSource(classes = {BookListController.class, BookFluxController.class}) + @ParameterizedTest + void batchingWithoutResult(Class controllerClass) { Map variables = Map.of("representations", List.of( Map.of("__typename", "Book", "id", "-99"), Map.of("__typename", "Book", "id", "4"), Map.of("__typename", "Book", "id", "5"))); - ResponseHelper helper = executeWith(BookBatchController.class, variables); + ResponseHelper helper = executeWith(controllerClass, variables); assertError(helper, 0, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty"); assertError(helper, 1, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty"); @@ -242,10 +247,53 @@ public class EntityMappingInvocationTests { @SuppressWarnings("unused") @Controller - private static class BookBatchController { + private static class BookListController { + + private final BookBatchService batchService = new BookBatchService(); @EntityMapping public List book(@Argument List idList, List> representations) { + return this.batchService.book(idList, representations); + } + + @BatchMapping + public List author(List books) { + return this.batchService.author(books); + } + + @GraphQlExceptionHandler + public GraphQLError handle(IllegalArgumentException ex, DataFetchingEnvironment env) { + return this.batchService.handle(ex, env); + } + } + + + @SuppressWarnings("unused") + @Controller + private static class BookFluxController { + + private final BookBatchService batchService = new BookBatchService(); + + @EntityMapping + public Flux book(@Argument List idList, List> representations) { + return Flux.fromIterable(this.batchService.book(idList, representations)); + } + + @BatchMapping + public Flux author(List books) { + return Flux.fromIterable(this.batchService.author(books)); + } + + @GraphQlExceptionHandler + public GraphQLError handle(IllegalArgumentException ex, DataFetchingEnvironment env) { + return this.batchService.handle(ex, env); + } + } + + + private static class BookBatchService { + + public List book(List idList, List> representations) { if (idList.get(0) == -97) { throw new IllegalArgumentException("handled"); @@ -265,12 +313,10 @@ public class EntityMappingInvocationTests { return idList.stream().map(id -> new Book((long) id, null, (Long) null)).toList(); } - @BatchMapping - public Flux author(List books) { - return Flux.fromIterable(books).map(book -> BookSource.getBook(book.getId()).getAuthor()); + public List author(List books) { + return books.stream().map(book -> BookSource.getBook(book.getId()).getAuthor()).toList(); } - @GraphQlExceptionHandler public GraphQLError handle(IllegalArgumentException ex, DataFetchingEnvironment env) { return GraphqlErrorBuilder.newError(env) .errorType(ErrorType.BAD_REQUEST)