diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/BatchMapping.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/BatchMapping.java index b90de5c2..9a0e4714 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/BatchMapping.java +++ b/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 @@ public @interface BatchMapping { */ String typeName() default ""; + /** + * Set the maximum number of keys to include a single batch, before + * splitting into multiple batches of keys to load. + *

By default this is -1 in which case there is no limit. + * @since 1.1 + */ + int maxBatchSize() default -1; + } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java index e9413a7b..a86146f4 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java @@ -309,11 +309,11 @@ public class AnnotatedControllerConfigurer 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 @@ public class AnnotatedControllerConfigurer 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 @@ public class AnnotatedControllerConfigurer "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 @@ public class AnnotatedControllerConfigurer Class clazz = returnType.getParameterType(); Class nestedClass = (clazz.equals(Callable.class) ? returnType.nested().getNestedParameterType() : clazz); + BatchLoaderRegistry.RegistrationSpec 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 @@ public class AnnotatedControllerConfigurer 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 class AnnotatedControllerConfigurer return this.batchMapping; } + public int getMaxBatchSize() { + return this.maxBatchSize; + } + public HandlerMethod getHandlerMethod() { return this.handlerMethod; } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java index f05d1770..93e3f4da 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java +++ b/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.GraphQLContext; 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 @@ public class BatchMappingDetectionTests { "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 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 @@ public class BatchMappingDetectionTests { .hasMessageStartingWith("Expected either @BatchMapping or @SchemaMapping, not both"); } - private RuntimeWiring.Builder initRuntimeWiringBuilder(Class handlerType) { + @SuppressWarnings("unchecked") + private RuntimeWiring.Builder initRuntimeWiringBuilder(T handler) { AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); - context.registerBean(handlerType); + if (handler instanceof Class handlerType) { + context.registerBean(handlerType); + } + else { + context.registerBean((Class) handler.getClass(), () -> handler); + } context.registerBean(BatchLoaderRegistry.class, () -> this.batchLoaderRegistry); context.refresh(); @@ -136,6 +162,24 @@ public class BatchMappingDetectionTests { } + @Controller + @SuppressWarnings("unused") + private static class BatchSizeController { + + private final List batchSizes = new ArrayList<>(); + + public List getBatchSizes() { + return this.batchSizes; + } + + @BatchMapping(maxBatchSize = 5, typeName = "Book") + public List authors(List bookIds) { + this.batchSizes.add(bookIds.size()); + return Collections.emptyList(); + } + } + + @Controller private static class InvalidReturnTypeController {