Add BatchLoaderRegistry

The registry is an application API for registering batch loading
functions along with DataLoaderOption's.

See gh-63

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;
 	}
 
 }

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.
+ *
+ * <p>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 <a href="https://www.graphql-java.com/documentation/v16/batching/">Using DataLoader</a>
+ * @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 <K> the key type
+	 * @param <V> the value type
+	 * @return a spec to complete the registration
+	 */
+	<K, V> RegistrationSpec<K, V> forTypePair(Class<K> keyType, Class<V> valueType);
+
+
+	/**
+	 * Spec to complete the registration of a batch loading function.
+	 *
+	 * @param <K> the type of the key that identifies the value
+	 * @param <V> the type of the data value
+	 */
+	interface RegistrationSpec<K, V> {
+
+		/**
+		 * Customize the name under which the {@link org.dataloader.DataLoader}
+		 * is registered and can be accessed in the data layer.
+		 * <p>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<K, V> 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<K, V> withOptions(Consumer<DataLoaderOptions> 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<K, V> withOptions(DataLoaderOptions options);
+
+		/**
+		 * Register the give batch loading function.
+		 * <p>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<List<K>, BatchLoaderEnvironment, Flux<V>> 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<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader);
+	}
+
+}

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);
+
+}

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<ReactorBatchLoader<?,?>> loaders = new ArrayList<>();
+
+	private final List<ReactorMappedBatchLoader<?,?>> mappedLoaders = new ArrayList<>();
+
+
+	@Override
+	public <K, V> RegistrationSpec<K, V> forTypePair(Class<K> keyType, Class<V> 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<K, V> implements RegistrationSpec<K, V> {
+
+		private String name;
+
+		private DataLoaderOptions options = DataLoaderOptions.newOptions();
+
+		public DefaultRegistrationSpec(String name) {
+			this.name = name;
+		}
+
+		@Override
+		public RegistrationSpec<K, V> withName(String name) {
+			this.name = name;
+			return this;
+		}
+
+		@Override
+		public RegistrationSpec<K, V> withOptions(Consumer<DataLoaderOptions> optionsConsumer) {
+			optionsConsumer.accept(this.options);
+			return this;
+		}
+
+		@Override
+		public RegistrationSpec<K, V> withOptions(DataLoaderOptions options) {
+			this.options = options;
+			return this;
+		}
+
+		@Override
+		public void registerBatchLoader(BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader) {
+			DefaultBatchLoaderRegistry.this.loaders.add(
+					new ReactorBatchLoader<>(this.name, loader, this.options));
+		}
+
+		@Override
+		public void registerMappedBatchLoader(BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader) {
+			DefaultBatchLoaderRegistry.this.mappedLoaders.add(
+					new ReactorMappedBatchLoader<>(this.name, loader, this.options));
+		}
+	}
+
+
+	private static class ReactorBatchLoader<K, V> implements BatchLoaderWithContext<K, V> {
+
+		private final String name;
+
+		private final BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader;
+
+		private final DataLoaderOptions options;
+
+		private ReactorBatchLoader(String name,
+				BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> 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<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
+			return this.loader.apply(keys, environment).collectList().toFuture();
+		}
+	}
+
+
+	private static class ReactorMappedBatchLoader<K, V> implements MappedBatchLoaderWithContext<K, V> {
+
+		private final String name;
+
+		private final BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader;
+
+		private final DataLoaderOptions options;
+
+		private ReactorMappedBatchLoader(String name,
+				BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> 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<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
+			return this.loader.apply(keys, environment).toFuture();
+		}
+	}
+
+}

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<DataLoaderRegistrar> 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<ExecutionResult> 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;
+	}
+
 }

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;
 	}

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<Long, Book> booksMap = new HashMap<>();
 
+	private static final Map<Long, Book> booksWithoutAuthorsMap;
+
 	private static final Map<Long, Author> 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<Long, Book> booksMap() {
-		return booksMap;
-	}
-
 	public static List<Book> books() {
 		return new ArrayList<>(booksMap.values());
 	}
 
+	public static List<Book> 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<Book> findBooksByAuthor(String author) {
 		return Flux.fromIterable(books())

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<String, Object> 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

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<String, Object> 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<Long, Author> 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<String, Object> data = result.getData();
+		assertThat(data).isNotNull();
+
+		List<Map<String, Object>> bookList = getValue(data, "booksByCriteria");
+		assertThat(bookList).hasSize(2);
+		Map<String, Object> authorMap = (Map<String, Object>) 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> T getValue(Map<String, Object> data, String key) {
+		return (T) data.get(key);
+	}
+
+}

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<String, DataLoader<?, ?>> 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<String, DataLoader<?, ?>> 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<String, DataLoader<?, ?>> map = dataLoaderRegistry.getDataLoadersMap();
+		assertThat(map).hasSize(1).containsKey(name);
+		assertThat(map.get(name).getStatistics()).isSameAs(collector.getStatistics());
+	}
+
+}