Support maxBatchSize via BatchMapping annotation

Closes gh-520

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/BatchMapping.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 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.
@@ -95,4 +95,12 @@
 	 */
 	String typeName() default "";
 
+	/**
+	 * Set the maximum number of keys to include a single batch, before
+	 * splitting into multiple batches of keys to load.
+	 * <p>By default this is -1 in which case there is no limit.
+	 * @since 1.1
+	 */
+	int maxBatchSize() default -1;
+
 }

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java

@@ -309,11 +309,11 @@ private MappingInfo getMappingInfo(Method method, Object handler, Class<?> handl
 		String typeName;
 		String field;
 		boolean batchMapping = false;
+		int batchSize = -1;
 		HandlerMethod handlerMethod = createHandlerMethod(method, handler, handlerType);
 
 		Annotation annotation = annotations.iterator().next();
-		if (annotation instanceof SchemaMapping) {
-			SchemaMapping mapping = (SchemaMapping) annotation;
+		if (annotation instanceof SchemaMapping mapping) {
 			typeName = mapping.typeName();
 			field = (StringUtils.hasText(mapping.field()) ? mapping.field() : method.getName());
 		}
@@ -322,6 +322,7 @@ private MappingInfo getMappingInfo(Method method, Object handler, Class<?> handl
 			typeName = mapping.typeName();
 			field = (StringUtils.hasText(mapping.field()) ? mapping.field() : method.getName());
 			batchMapping = true;
+			batchSize = mapping.maxBatchSize();
 		}
 
 		if (!StringUtils.hasText(typeName)) {
@@ -354,7 +355,7 @@ private MappingInfo getMappingInfo(Method method, Object handler, Class<?> handl
 				"No parentType specified, and a source/parent method argument was also not found: " +
 						handlerMethod.getShortLogMessage());
 
-		return new MappingInfo(typeName, field, batchMapping, handlerMethod);
+		return new MappingInfo(typeName, field, batchMapping, batchSize, handlerMethod);
 	}
 
 	private HandlerMethod createHandlerMethod(Method method, Object handler, Class<?> handlerType) {
@@ -394,11 +395,16 @@ private String registerBatchLoader(MappingInfo info) {
 		Class<?> clazz = returnType.getParameterType();
 		Class<?> nestedClass = (clazz.equals(Callable.class) ? returnType.nested().getNestedParameterType() : clazz);
 
+		BatchLoaderRegistry.RegistrationSpec<Object, Object> registration = registry.forName(dataLoaderKey);
+		if (info.getMaxBatchSize() > 0) {
+			registration.withOptions(options -> options.setMaxBatchSize(info.getMaxBatchSize()));
+		}
+
 		if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(nestedClass)) {
-			registry.forName(dataLoaderKey).registerBatchLoader(invocable::invokeForIterable);
+			registration.registerBatchLoader(invocable::invokeForIterable);
 		}
 		else if (clazz.equals(Mono.class) || nestedClass.equals(Map.class)) {
-			registry.forName(dataLoaderKey).registerMappedBatchLoader(invocable::invokeForMap);
+			registration.registerMappedBatchLoader(invocable::invokeForMap);
 		}
 		else {
 			throw new IllegalStateException("@BatchMapping method is expected to return " +
@@ -434,12 +440,18 @@ private static class MappingInfo {
 
 		private final boolean batchMapping;
 
+		private final int maxBatchSize;
+
 		private final HandlerMethod handlerMethod;
 
-		public MappingInfo(String typeName, String field, boolean batchMapping, HandlerMethod handlerMethod) {
+		public MappingInfo(
+				String typeName, String field, boolean batchMapping, int maxBatchSize,
+				HandlerMethod handlerMethod) {
+
 			this.coordinates = FieldCoordinates.coordinates(typeName, field);
-			this.handlerMethod = handlerMethod;
 			this.batchMapping = batchMapping;
+			this.maxBatchSize = maxBatchSize;
+			this.handlerMethod = handlerMethod;
 		}
 
 		public FieldCoordinates getCoordinates() {
@@ -451,6 +463,10 @@ public boolean isBatchMapping() {
 			return this.batchMapping;
 		}
 
+		public int getMaxBatchSize() {
+			return this.maxBatchSize;
+		}
+
 		public HandlerMethod getHandlerMethod() {
 			return this.handlerMethod;
 		}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java

@@ -15,6 +15,9 @@
  */
 package org.springframework.graphql.data.method.annotation.support;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Callable;
@@ -23,6 +26,7 @@
 import graphql.schema.DataFetcher;
 import graphql.schema.idl.RuntimeWiring;
 import org.dataloader.BatchLoaderEnvironment;
+import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderRegistry;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
@@ -71,6 +75,22 @@ void registerWithDefaultCoordinates() {
 				"Book.authorCallableMap", "Book.authorEnvironment");
 	}
 
+	@Test
+	void registerWithMaxBatchSize() {
+
+		BatchSizeController controller = new BatchSizeController();
+		initRuntimeWiringBuilder(controller).build();
+		DataLoaderRegistry registry = new DataLoaderRegistry();
+		this.batchLoaderRegistry.registerDataLoaders(registry, GraphQLContext.newContext().build());
+
+		DataLoader<Integer, String> dataLoader = registry.getDataLoader("Book.authors");
+		dataLoader.loadMany(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8));
+		dataLoader.dispatchAndJoin();
+
+		assertThat(controller.getBatchSizes()).containsExactly(5, 3);
+
+	}
+
 	@Test
 	void invalidReturnType() {
 		assertThatThrownBy(() -> initRuntimeWiringBuilder(InvalidReturnTypeController.class).build())
@@ -83,9 +103,15 @@ void schemaAndBatch() {
 				.hasMessageStartingWith("Expected either @BatchMapping or @SchemaMapping, not both");
 	}
 
-	private RuntimeWiring.Builder initRuntimeWiringBuilder(Class<?> handlerType) {
+	@SuppressWarnings("unchecked")
+	private <T> RuntimeWiring.Builder initRuntimeWiringBuilder(T handler) {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
-		context.registerBean(handlerType);
+		if (handler instanceof Class<?> handlerType) {
+			context.registerBean(handlerType);
+		}
+		else {
+			context.registerBean((Class<T>) handler.getClass(), () -> handler);
+		}
 		context.registerBean(BatchLoaderRegistry.class, () -> this.batchLoaderRegistry);
 		context.refresh();
 
@@ -136,6 +162,24 @@ public List<Author> authorEnvironment(BatchLoaderEnvironment environment, List<B
 	}
 
 
+	@Controller
+	@SuppressWarnings("unused")
+	private static class BatchSizeController {
+
+		private final List<Integer> batchSizes = new ArrayList<>();
+
+		public List<Integer> getBatchSizes() {
+			return this.batchSizes;
+		}
+
+		@BatchMapping(maxBatchSize = 5, typeName = "Book")
+		public List<String> authors(List<Integer> bookIds) {
+			this.batchSizes.add(bookIds.size());
+			return Collections.emptyList();
+		}
+	}
+
+
 	@Controller
 	private static class InvalidReturnTypeController {