From ddafc93c6c05684ef3b37ca3809ca66e090d94ca Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Mon, 13 Sep 2021 16:34:10 +0100 Subject: [PATCH] Add BatchLoaderRegistry The registry is an application API for registering batch loading functions along with DataLoaderOption's. See gh-63 --- .../boot/GraphQlServiceAutoConfiguration.java | 13 +- .../execution/BatchLoaderRegistry.java | 112 +++++++++++ .../execution/DataLoaderRegistrar.java | 35 ++++ .../execution/DefaultBatchLoaderRegistry.java | 181 ++++++++++++++++++ .../execution/ExecutionGraphQlService.java | 32 +++- .../org/springframework/graphql/Book.java | 14 ++ .../springframework/graphql/BookSource.java | 46 +++-- .../AnnotatedDataFetcherInvocationTests.java | 7 +- .../graphql/execution/BatchLoadingTests.java | 103 ++++++++++ .../DefaultBatchLoaderRegistryTests.java | 81 ++++++++ 10 files changed, 600 insertions(+), 24 deletions(-) create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/execution/BatchLoadingTests.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java diff --git a/graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/GraphQlServiceAutoConfiguration.java b/graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/GraphQlServiceAutoConfiguration.java index 6729dda5..ec436480 100644 --- a/graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/GraphQlServiceAutoConfiguration.java +++ b/graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/GraphQlServiceAutoConfiguration.java @@ -25,6 +25,8 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.graphql.GraphQlService; +import org.springframework.graphql.execution.BatchLoaderRegistry; +import org.springframework.graphql.execution.DefaultBatchLoaderRegistry; import org.springframework.graphql.execution.ExecutionGraphQlService; import org.springframework.graphql.execution.GraphQlSource; @@ -41,10 +43,19 @@ import org.springframework.graphql.execution.GraphQlSource; @AutoConfigureAfter(GraphQlAutoConfiguration.class) public class GraphQlServiceAutoConfiguration { + private final DefaultBatchLoaderRegistry batchLoaderRegistry = new DefaultBatchLoaderRegistry(); + + @Bean + public BatchLoaderRegistry batchLoaderRegistry() { + return this.batchLoaderRegistry; + } + @Bean @ConditionalOnMissingBean public GraphQlService graphQlService(GraphQlSource graphQlSource) { - return new ExecutionGraphQlService(graphQlSource); + ExecutionGraphQlService service = new ExecutionGraphQlService(graphQlSource); + service.addDataLoaderRegistrar(this.batchLoaderRegistry); + return service; } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java new file mode 100644 index 00000000..f33265c7 --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/BatchLoaderRegistry.java @@ -0,0 +1,112 @@ +/* + * 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.execution; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.BiFunction; +import java.util.function.Consumer; + +import org.dataloader.BatchLoaderEnvironment; +import org.dataloader.DataLoaderOptions; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +/** + * Registry of functions that batch load data values given a set of keys. + * + *

At request time, each function is registered as a + * {@link org.dataloader.DataLoader} in the {@link org.dataloader.DataLoaderRegistry} + * and can be accessed in the data layer to load related entities while avoiding + * the N+1 select problem. + * + * @author Rossen Stoyanchev + * @since 1.0.0 + * @see Using DataLoader + * @see org.dataloader.BatchLoader + * @see org.dataloader.MappedBatchLoader + * @see org.dataloader.DataLoader + */ +public interface BatchLoaderRegistry { + + /** + * Start the registration of a new function for batch loading data values by + * specifying the key and value types. + * @param keyType the type of the key that identifies the value + * @param valueType the type of the data value + * @param the key type + * @param the value type + * @return a spec to complete the registration + */ + RegistrationSpec forTypePair(Class keyType, Class valueType); + + + /** + * Spec to complete the registration of a batch loading function. + * + * @param the type of the key that identifies the value + * @param the type of the data value + */ + interface RegistrationSpec { + + /** + * Customize the name under which the {@link org.dataloader.DataLoader} + * is registered and can be accessed in the data layer. + *

By default, this is the full class name of the value type. + * @param name the name to use + * @return a spec to complete the registration + */ + RegistrationSpec withName(String name); + + /** + * Customize the {@link DataLoaderOptions} to use to create the + * {@link org.dataloader.DataLoader} via {@link org.dataloader.DataLoaderFactory}. + * @param optionsConsumer callback to customize the options, invoked + * immediately and given access to the options instance + * @return a spec to complete the registration + */ + RegistrationSpec withOptions(Consumer optionsConsumer); + + /** + * Replace the {@link DataLoaderOptions} to use to create the + * {@link org.dataloader.DataLoader} via {@link org.dataloader.DataLoaderFactory}. + * @param options the options to use + * @return a spec to complete the registration + */ + RegistrationSpec withOptions(DataLoaderOptions options); + + /** + * Register the give batch loading function. + *

The values returned from the function must match the order and + * the number of keys, with {@code null} for missing values. + * Please, see {@link org.dataloader.BatchLoader}. + * @param loader the loader function + * @see org.dataloader.BatchLoader + */ + void registerBatchLoader(BiFunction, BatchLoaderEnvironment, Flux> loader); + + /** + * A variant of {@link #registerBatchLoader(BiFunction)} that returns a + * Map of key-value pairs, which is useful is there aren't values for all keys. + * Please see {@link org.dataloader.MappedBatchLoader}. + * @param loader the loader function + * @see org.dataloader.MappedBatchLoader + */ + void registerMappedBatchLoader(BiFunction, BatchLoaderEnvironment, Mono>> loader); + } + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java new file mode 100644 index 00000000..7f0cb604 --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataLoaderRegistrar.java @@ -0,0 +1,35 @@ +/* + * 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.execution; + +import org.dataloader.DataLoaderRegistry; + +/** + * Contract for access to the {@link DataLoaderRegistry} at request time. + * + * @author Rossen Stoyanchev + * @since 1.0.0 + */ +public interface DataLoaderRegistrar { + + /** + * Callback that provides access to the {@link DataLoaderRegistry} from the + * the {@link graphql.ExecutionInput}. + * @param registry the registry to make registrations against + */ + void registerDataLoaders(DataLoaderRegistry registry); + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java new file mode 100644 index 00000000..05633e32 --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java @@ -0,0 +1,181 @@ +/* + * 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.execution; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletionStage; +import java.util.function.BiFunction; +import java.util.function.Consumer; + +import org.dataloader.BatchLoaderEnvironment; +import org.dataloader.BatchLoaderWithContext; +import org.dataloader.DataLoader; +import org.dataloader.DataLoaderFactory; +import org.dataloader.DataLoaderOptions; +import org.dataloader.DataLoaderRegistry; +import org.dataloader.MappedBatchLoaderWithContext; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +/** + * A default implementation of {@link BatchLoaderRegistry} that accepts + * registrations, and also an implementation of {@link DataLoaderRegistrar} to + * apply those registrations to a {@link DataLoaderRegistry}. + * + * @author Rossen Stoyanchev + * @since 1.0.0 + */ +public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry, DataLoaderRegistrar { + + private final List> loaders = new ArrayList<>(); + + private final List> mappedLoaders = new ArrayList<>(); + + + @Override + public RegistrationSpec forTypePair(Class keyType, Class valueType) { + return new DefaultRegistrationSpec<>(valueType.getName()); + } + + @Override + public void registerDataLoaders(DataLoaderRegistry registry) { + for (ReactorBatchLoader loader : this.loaders) { + DataLoader dataLoader = DataLoaderFactory.newDataLoader(loader, loader.getOptions()); + registerDataLoader(loader.getName(), dataLoader, registry); + } + for (ReactorMappedBatchLoader loader : this.mappedLoaders) { + DataLoader dataLoader = DataLoaderFactory.newMappedDataLoader(loader, loader.getOptions()); + registerDataLoader(loader.getName(), dataLoader, registry); + } + } + + private void registerDataLoader(String name, DataLoader dataLoader, DataLoaderRegistry registry) { + if (registry.getDataLoader(name) != null) { + throw new IllegalStateException("More than one DataLoader named '" + name + "'"); + } + registry.register(name, dataLoader); + } + + + private class DefaultRegistrationSpec implements RegistrationSpec { + + private String name; + + private DataLoaderOptions options = DataLoaderOptions.newOptions(); + + public DefaultRegistrationSpec(String name) { + this.name = name; + } + + @Override + public RegistrationSpec withName(String name) { + this.name = name; + return this; + } + + @Override + public RegistrationSpec withOptions(Consumer optionsConsumer) { + optionsConsumer.accept(this.options); + return this; + } + + @Override + public RegistrationSpec withOptions(DataLoaderOptions options) { + this.options = options; + return this; + } + + @Override + public void registerBatchLoader(BiFunction, BatchLoaderEnvironment, Flux> loader) { + DefaultBatchLoaderRegistry.this.loaders.add( + new ReactorBatchLoader<>(this.name, loader, this.options)); + } + + @Override + public void registerMappedBatchLoader(BiFunction, BatchLoaderEnvironment, Mono>> loader) { + DefaultBatchLoaderRegistry.this.mappedLoaders.add( + new ReactorMappedBatchLoader<>(this.name, loader, this.options)); + } + } + + + private static class ReactorBatchLoader implements BatchLoaderWithContext { + + private final String name; + + private final BiFunction, BatchLoaderEnvironment, Flux> loader; + + private final DataLoaderOptions options; + + private ReactorBatchLoader(String name, + BiFunction, BatchLoaderEnvironment, Flux> loader, + DataLoaderOptions options) { + + this.name = name; + this.loader = loader; + this.options = options; + } + + public String getName() { + return this.name; + } + + public DataLoaderOptions getOptions() { + return this.options; + } + + @Override + public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { + return this.loader.apply(keys, environment).collectList().toFuture(); + } + } + + + private static class ReactorMappedBatchLoader implements MappedBatchLoaderWithContext { + + private final String name; + + private final BiFunction, BatchLoaderEnvironment, Mono>> loader; + + private final DataLoaderOptions options; + + private ReactorMappedBatchLoader(String name, + BiFunction, BatchLoaderEnvironment, Mono>> loader, + DataLoaderOptions options) { + + this.name = name; + this.loader = loader; + this.options = options; + } + + public String getName() { + return this.name; + } + + public DataLoaderOptions getOptions() { + return this.options; + } + + @Override + public CompletionStage> load(Set keys, BatchLoaderEnvironment environment) { + return this.loader.apply(keys, environment).toFuture(); + } + } + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java index 3a35c72b..5c3c9d61 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java @@ -16,9 +16,13 @@ package org.springframework.graphql.execution; +import java.util.ArrayList; +import java.util.List; + import graphql.ExecutionInput; import graphql.ExecutionResult; import graphql.GraphQL; +import org.dataloader.DataLoaderRegistry; import reactor.core.publisher.Mono; import org.springframework.graphql.GraphQlService; @@ -35,14 +39,27 @@ public class ExecutionGraphQlService implements GraphQlService { private final GraphQlSource graphQlSource; + private final List dataLoaderRegistrars = new ArrayList<>(); + + public ExecutionGraphQlService(GraphQlSource graphQlSource) { this.graphQlSource = graphQlSource; } + + /** + * Add a registrar to get access to and configure the + * {@link DataLoaderRegistry} for each request. + * @param registrar the registrar to add + */ + public void addDataLoaderRegistrar(DataLoaderRegistrar registrar) { + this.dataLoaderRegistrars.add(registrar); + } + + @Override public final Mono execute(RequestInput requestInput) { - ExecutionInput executionInput = requestInput.toExecutionInput(); - + ExecutionInput executionInput = initExecutionInput(requestInput); GraphQL graphQl = this.graphQlSource.graphQl(); return Mono.deferContextual((contextView) -> { @@ -51,4 +68,15 @@ public class ExecutionGraphQlService implements GraphQlService { }); } + private ExecutionInput initExecutionInput(RequestInput requestInput) { + ExecutionInput input = requestInput.toExecutionInput(); + if (!this.dataLoaderRegistrars.isEmpty()) { + DataLoaderRegistry previousRegistry = input.getDataLoaderRegistry(); + DataLoaderRegistry newRegistry = DataLoaderRegistry.newRegistry().registerAll(previousRegistry).build(); + this.dataLoaderRegistrars.forEach(registrar -> registrar.registerDataLoaders(newRegistry)); + input = input.transform(builder -> builder.dataLoaderRegistry(newRegistry)); + } + return input; + } + } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/Book.java b/spring-graphql/src/test/java/org/springframework/graphql/Book.java index dbe24efb..b3f2a2ab 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/Book.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/Book.java @@ -22,14 +22,24 @@ public class Book { String name; + Long authorId; + Author author; public Book() { } + public Book(Long id, String name, Long authorId) { + this.id = id; + this.name = name; + this.authorId = authorId; + this.author = null; + } + public Book(Long id, String name, Author author) { this.id = id; this.name = name; + this.authorId = author.getId(); this.author = author; } @@ -49,6 +59,10 @@ public class Book { this.name = name; } + public Long getAuthorId() { + return this.authorId; + } + public Author getAuthor() { return this.author; } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/BookSource.java b/spring-graphql/src/test/java/org/springframework/graphql/BookSource.java index 271d4513..43e33209 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/BookSource.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/BookSource.java @@ -19,6 +19,8 @@ import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Function; +import java.util.stream.Collectors; import reactor.core.publisher.Flux; @@ -26,38 +28,48 @@ public class BookSource { private static final Map booksMap = new HashMap<>(); + private static final Map booksWithoutAuthorsMap; + private static final Map authorsMap = new HashMap<>(); static { - authorsMap.put(1L, new Author(1L, "George", "Orwell")); - authorsMap.put(2L, new Author(2L, "F. Scott", "Fitzgerald")); - authorsMap.put(3L, new Author(3L, "Joseph", "Heller")); - authorsMap.put(4L, new Author(4L, "Virginia", "Woolf")); - authorsMap.put(5L, new Author(5L, "Douglas", "Adams")); - authorsMap.put(6L, new Author(6L, "Vince", "Gilligan")); + authorsMap.put(101L, new Author(101L, "George", "Orwell")); + authorsMap.put(102L, new Author(102L, "F. Scott", "Fitzgerald")); + authorsMap.put(103L, new Author(103L, "Joseph", "Heller")); + authorsMap.put(104L, new Author(104L, "Virginia", "Woolf")); + authorsMap.put(105L, new Author(105L, "Douglas", "Adams")); + authorsMap.put(106L, new Author(106L, "Vince", "Gilligan")); - booksMap.put(1L, new Book(1L, "Nineteen Eighty-Four", authorsMap.get(1L))); - booksMap.put(2L, new Book(2L, "The Great Gatsby", authorsMap.get(2L))); - booksMap.put(3L, new Book(3L, "Catch-22", authorsMap.get(3L))); - booksMap.put(4L, new Book(4L, "To The Lighthouse", authorsMap.get(4L))); - booksMap.put(5L, new Book(5L, "Animal Farm", authorsMap.get(1L))); - booksMap.put(42L, new Book(42L, "Hitchhiker's Guide to the Galaxy", authorsMap.get(5L))); - booksMap.put(53L, new Book(53L, "Breaking Bad", authorsMap.get(6L))); + booksMap.put(1L, new Book(1L, "Nineteen Eighty-Four", authorsMap.get(101L))); + booksMap.put(2L, new Book(2L, "The Great Gatsby", authorsMap.get(102L))); + booksMap.put(3L, new Book(3L, "Catch-22", authorsMap.get(103L))); + booksMap.put(4L, new Book(4L, "To The Lighthouse", authorsMap.get(104L))); + booksMap.put(5L, new Book(5L, "Animal Farm", authorsMap.get(101L))); + booksMap.put(42L, new Book(42L, "Hitchhiker's Guide to the Galaxy", authorsMap.get(105L))); + booksMap.put(53L, new Book(53L, "Breaking Bad", authorsMap.get(106L))); + + booksWithoutAuthorsMap = booksMap.values().stream() + .map(book -> new Book(book.getId(), book.getName(), book.getAuthorId())) + .collect(Collectors.toMap(Book::getId, Function.identity())); } - public static Map booksMap() { - return booksMap; - } - public static List books() { return new ArrayList<>(booksMap.values()); } + public static List booksWithoutAuthors() { + return new ArrayList<>(booksWithoutAuthorsMap.values()); + } + public static Book getBook(Long id) { return booksMap.get(id); } + public static Book getBookWithoutAuthor(Long id) { + return booksWithoutAuthorsMap.get(id); + } + @SuppressWarnings("ConstantConditions") public static List findBooksByAuthor(String author) { return Flux.fromIterable(books()) diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedDataFetcherInvocationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedDataFetcherInvocationTests.java index d725f175..ce42bbf1 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedDataFetcherInvocationTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedDataFetcherInvocationTests.java @@ -39,7 +39,6 @@ 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.data.method.annotation.support.AnnotatedDataFetcherConfigurer; import org.springframework.graphql.execution.GraphQlSource; import org.springframework.http.codec.ServerCodecConfigurer; import org.springframework.stereotype.Controller; @@ -104,7 +103,7 @@ public class AnnotatedDataFetcherInvocationTests { @Test void queryWithArgumentViaDataFetchingEnvironment() { String query = "{ " + - " authorById(id:\"1\") { " + + " authorById(id:\"101\") { " + " id" + " firstName" + " lastName" + @@ -119,7 +118,7 @@ public class AnnotatedDataFetcherInvocationTests { assertThat(data).isNotNull(); Map author = getValue(data, "authorById"); - assertThat(author.get("id")).isEqualTo("1"); + assertThat(author.get("id")).isEqualTo("101"); assertThat(author.get("firstName")).isEqualTo("George"); assertThat(author.get("lastName")).isEqualTo("Orwell"); @@ -210,7 +209,7 @@ public class AnnotatedDataFetcherInvocationTests { @QueryMapping public Book bookById(@Argument Long id) { - return new Book(id, BookSource.getBook(id).getName(), null); + return BookSource.getBookWithoutAuthor(id); } @QueryMapping diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/BatchLoadingTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/BatchLoadingTests.java new file mode 100644 index 00000000..63439080 --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/BatchLoadingTests.java @@ -0,0 +1,103 @@ +/* + * 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.execution; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import graphql.ExecutionResult; +import org.dataloader.DataLoader; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; + +import org.springframework.core.io.ClassPathResource; +import org.springframework.graphql.Author; +import org.springframework.graphql.Book; +import org.springframework.graphql.BookSource; +import org.springframework.graphql.RequestInput; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for requests with batch loading, performed through an + * {@link ExecutionGraphQlService} configured with a {@link BatchLoaderRegistry}. + * + * @author Rossen Stoyanchev + */ +public class BatchLoadingTests { + + @Test + void batchLoader() { + ExecutionGraphQlService service = initExecutionGraphQlService(wiring -> { + wiring.type("Query", builder -> builder.dataFetcher("booksByCriteria", env -> { + Map criteria = env.getArgument("criteria"); + String authorName = (String) criteria.get("author"); + return BookSource.findBooksByAuthor(authorName).stream() + .map(book -> new Book(book.getId(), book.getName(), book.getAuthorId())) + .collect(Collectors.toList()); + })); + wiring.type("Book", builder -> builder.dataFetcher("author", env -> { + Book book = env.getSource(); + DataLoader dataLoader = env.getDataLoader(Author.class.getName()); + return dataLoader.load(book.getAuthorId()); + })); + }); + + DefaultBatchLoaderRegistry registry = new DefaultBatchLoaderRegistry(); + registry.forTypePair(Long.class, Author.class) + .registerBatchLoader((ids, env) -> Flux.fromIterable(ids).map(BookSource::getAuthor)); + + service.addDataLoaderRegistrar(registry); + + String query = "{ " + + " booksByCriteria(criteria: {author:\"Orwell\"}) { " + + " author {" + + " firstName, " + + " lastName " + + " }" + + " }" + + "}"; + + RequestInput input = new RequestInput(query, null, null); + ExecutionResult result = service.execute(input).block(); + + assertThat(result.getErrors()).isEmpty(); + Map data = result.getData(); + assertThat(data).isNotNull(); + + List> bookList = getValue(data, "booksByCriteria"); + assertThat(bookList).hasSize(2); + Map authorMap = (Map) bookList.get(0).get("author"); + assertThat(authorMap).isNotNull(); + assertThat(authorMap).containsEntry("firstName", "George"); + assertThat(authorMap).containsEntry("lastName", "Orwell"); + } + + private ExecutionGraphQlService initExecutionGraphQlService(RuntimeWiringConfigurer configurer) { + GraphQlSource graphQlSource = GraphQlSource.builder() + .schemaResources(new ClassPathResource("books/schema.graphqls")) + .configureRuntimeWiring(configurer) + .build(); + return new ExecutionGraphQlService(graphQlSource); + } + + @SuppressWarnings("unchecked") + private T getValue(Map data, String key) { + return (T) data.get(key); + } + +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java new file mode 100644 index 00000000..c4962d00 --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java @@ -0,0 +1,81 @@ +/* + * 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.execution; + +import java.util.Map; + +import org.dataloader.DataLoader; +import org.dataloader.DataLoaderRegistry; +import org.dataloader.stats.NoOpStatisticsCollector; +import org.dataloader.stats.StatisticsCollector; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import org.springframework.graphql.Book; + +import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat; + +/** + * Unit tests for {@link DefaultBatchLoaderRegistry}. + */ +public class DefaultBatchLoaderRegistryTests { + + private final DefaultBatchLoaderRegistry batchLoaderRegistry = new DefaultBatchLoaderRegistry(); + + private final DataLoaderRegistry dataLoaderRegistry = DataLoaderRegistry.newRegistry().build(); + + + @Test + void batchLoader() { + this.batchLoaderRegistry.forTypePair(String.class, Book.class).registerBatchLoader((keys, environment) -> Flux.empty()); + this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry); + + Map> map = this.dataLoaderRegistry.getDataLoadersMap(); + assertThat(map).hasSize(1).containsKey(Book.class.getName()); + } + + @Test + void mappedBatchLoader() { + this.batchLoaderRegistry + .forTypePair(String.class, Book.class) + .registerMappedBatchLoader((keys, environment) -> Mono.empty()); + + this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry); + + Map> map = this.dataLoaderRegistry.getDataLoadersMap(); + assertThat(map).hasSize(1).containsKey(Book.class.getName()); + } + + @Test + void batchLoaderWithCustomNameAndOptions() { + String name = "myLoader"; + StatisticsCollector collector = new NoOpStatisticsCollector(); + + this.batchLoaderRegistry + .forTypePair(String.class, Book.class) + .withName(name) + .withOptions(options -> options.setStatisticsCollector(() -> collector)) + .registerBatchLoader((keys, environment) -> Flux.empty()); + + this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry); + + Map> map = dataLoaderRegistry.getDataLoadersMap(); + assertThat(map).hasSize(1).containsKey(name); + assertThat(map.get(name).getStatistics()).isSameAs(collector.getStatistics()); + } + +}