Skip to content

Commit 1a5dacc

Browse files
rstoyanchevrstoyanchev
committed
SchemaMappingInspector @BatchMapping method support
Closes gh-673
1 parent c7c75da commit 1a5dacc

File tree

2 files changed

+87
-20
lines changed

2 files changed

+87
-20
lines changed

Diff for: ‎spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java

+44-19
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import java.util.Map;
2929
import java.util.Set;
3030
import java.util.concurrent.Callable;
31+
import java.util.concurrent.CompletableFuture;
3132
import java.util.concurrent.Executor;
3233
import java.util.function.BiConsumer;
3334
import java.util.function.Consumer;
@@ -349,8 +350,7 @@ public void configure(RuntimeWiring.Builder runtimeWiringBuilder) {
349350
info, this.argumentResolvers, this.validationHelper, this.exceptionResolver, this.executor);
350351
}
351352
else {
352-
String dataLoaderKey = registerBatchLoader(info);
353-
dataFetcher = new BatchMappingDataFetcher(dataLoaderKey);
353+
dataFetcher = registerBatchLoader(info);
354354
}
355355
runtimeWiringBuilder.type(info.getCoordinates().getTypeName(), typeBuilder ->
356356
typeBuilder.dataFetcher(info.getCoordinates().getFieldName(), dataFetcher));
@@ -502,38 +502,49 @@ private String formatMappings(Class<?> handlerType, Collection<MappingInfo> info
502502
.collect(Collectors.joining("\n\t", "\n\t" + formattedType + ":" + "\n\t", ""));
503503
}
504504

505-
private String registerBatchLoader(MappingInfo info) {
505+
private DataFetcher<Object> registerBatchLoader(MappingInfo info) {
506506
if (!info.isBatchMapping()) {
507507
throw new IllegalArgumentException("Not a @BatchMapping method: " + info);
508508
}
509509

510510
String dataLoaderKey = info.getCoordinates().toString();
511511
BatchLoaderRegistry registry = obtainApplicationContext().getBean(BatchLoaderRegistry.class);
512+
BatchLoaderRegistry.RegistrationSpec<Object, Object> registration = registry.forName(dataLoaderKey);
513+
if (info.getMaxBatchSize() > 0) {
514+
registration.withOptions(options -> options.setMaxBatchSize(info.getMaxBatchSize()));
515+
}
512516

513517
HandlerMethod handlerMethod = info.getHandlerMethod();
514518
BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(handlerMethod, this.executor);
515519

516520
MethodParameter returnType = handlerMethod.getReturnType();
517521
Class<?> clazz = returnType.getParameterType();
518-
Class<?> nestedClass = (clazz.equals(Callable.class) ? returnType.nested().getNestedParameterType() : clazz);
519522

520-
BatchLoaderRegistry.RegistrationSpec<Object, Object> registration = registry.forName(dataLoaderKey);
521-
if (info.getMaxBatchSize() > 0) {
522-
registration.withOptions(options -> options.setMaxBatchSize(info.getMaxBatchSize()));
523+
if (clazz.equals(Callable.class)) {
524+
returnType = returnType.nested();
525+
clazz = returnType.getNestedParameterType();
523526
}
524527

525-
if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(nestedClass)) {
528+
if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(clazz)) {
526529
registration.registerBatchLoader(invocable::invokeForIterable);
530+
ResolvableType valueType = ResolvableType.forMethodParameter(returnType.nested());
531+
return new BatchMappingDataFetcher(info, valueType, dataLoaderKey);
527532
}
528-
else if (clazz.equals(Mono.class) || nestedClass.equals(Map.class)) {
529-
registration.registerMappedBatchLoader(invocable::invokeForMap);
533+
534+
if (clazz.equals(Mono.class)) {
535+
returnType = returnType.nested();
536+
clazz = returnType.getNestedParameterType();
530537
}
531-
else {
532-
throw new IllegalStateException("@BatchMapping method is expected to return " +
533-
"Flux<V>, List<V>, Mono<Map<K, V>>, or Map<K, V>: " + handlerMethod);
538+
539+
if (Map.class.isAssignableFrom(clazz)) {
540+
registration.registerMappedBatchLoader(invocable::invokeForMap);
541+
ResolvableType valueType = ResolvableType.forMethodParameter(returnType.nested(1));
542+
return new BatchMappingDataFetcher(info, valueType, dataLoaderKey);
534543
}
535544

536-
return dataLoaderKey;
545+
throw new IllegalStateException(
546+
"@BatchMapping method is expected to return " +
547+
"Mono<Map<K, V>>, Map<K, V>, Flux<V>, or Collection<V>: " + handlerMethod);
537548
}
538549

539550
/**
@@ -715,20 +726,34 @@ public String toString() {
715726
}
716727

717728

718-
static class BatchMappingDataFetcher implements DataFetcher<Object> {
729+
static class BatchMappingDataFetcher implements DataFetcher<Object>, SelfDescribingDataFetcher<Object> {
730+
731+
private final MappingInfo info;
732+
733+
private final ResolvableType returnType;
719734

720735
private final String dataLoaderKey;
721736

722-
BatchMappingDataFetcher(String dataLoaderKey) {
737+
BatchMappingDataFetcher(MappingInfo info, ResolvableType valueType, String dataLoaderKey) {
738+
this.info = info;
739+
this.returnType = ResolvableType.forClassWithGenerics(CompletableFuture.class, valueType);
723740
this.dataLoaderKey = dataLoaderKey;
724741
}
725742

743+
@Override
744+
public String getDescription() {
745+
return "@BatchMapping " + this.info.getHandlerMethod().getShortLogMessage();
746+
}
747+
748+
@Override
749+
public ResolvableType getReturnType() {
750+
return this.returnType;
751+
}
752+
726753
@Override
727754
public Object get(DataFetchingEnvironment env) {
728755
DataLoader<?, ?> dataLoader = env.getDataLoaderRegistry().getDataLoader(this.dataLoaderKey);
729-
if (dataLoader == null) {
730-
throw new IllegalStateException("No DataLoader for key '" + this.dataLoaderKey + "'");
731-
}
756+
Assert.state(dataLoader != null, "No DataLoader for key '" + this.dataLoaderKey + "'");
732757
return dataLoader.load(env.getSource());
733758
}
734759
}

Diff for: ‎spring-graphql/src/test/java/org/springframework/graphql/execution/SchemaMappingInspectorTests.java

+43-1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.Arrays;
2020
import java.util.Collections;
2121
import java.util.List;
22+
import java.util.Map;
2223
import java.util.Set;
2324
import java.util.concurrent.CompletableFuture;
2425

@@ -30,13 +31,15 @@
3031
import org.junit.jupiter.api.Nested;
3132
import org.junit.jupiter.api.Test;
3233
import reactor.core.publisher.Flux;
34+
import reactor.core.publisher.Mono;
3335

3436
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
3537
import org.springframework.data.domain.OffsetScrollPosition;
3638
import org.springframework.data.domain.Window;
3739
import org.springframework.graphql.Author;
3840
import org.springframework.graphql.Book;
3941
import org.springframework.graphql.data.method.annotation.Argument;
42+
import org.springframework.graphql.data.method.annotation.BatchMapping;
4043
import org.springframework.graphql.data.method.annotation.MutationMapping;
4144
import org.springframework.graphql.data.method.annotation.QueryMapping;
4245
import org.springframework.graphql.data.method.annotation.SchemaMapping;
@@ -309,6 +312,29 @@ void reportIsEmptyWhenFieldHasDataFetcherMapping() {
309312
assertThatReport(report).hasUnmappedFieldCount(0).hasSkippedTypeCount(0);
310313
}
311314

315+
@Test
316+
void reportIsEmptyWhenFieldHasBatchMapping() {
317+
String schema = """
318+
type Query {
319+
books: [Book]
320+
}
321+
322+
type Book {
323+
id: ID
324+
name: String
325+
author: Author
326+
}
327+
328+
type Author {
329+
id: ID
330+
firstName: String
331+
lastName: String
332+
}
333+
""";
334+
SchemaMappingInspector.Report report = inspectSchema(schema, BatchMappingBookController.class);
335+
assertThatReport(report).hasUnmappedFieldCount(0).hasSkippedTypeCount(0);
336+
}
337+
312338
@Test
313339
void reportHasUnmappedField() {
314340
String schema = """
@@ -527,6 +553,7 @@ private RuntimeWiring createRuntimeWiring(Class<?>... controllerTypes) {
527553
for (Class<?> controllerType : controllerTypes) {
528554
context.registerBean(controllerType);
529555
}
556+
context.registerBean(BatchLoaderRegistry.class, () -> new DefaultBatchLoaderRegistry());
530557
context.refresh();
531558

532559
AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
@@ -604,6 +631,21 @@ public Flux<List<Book>> bookSearch(@Argument String author) {
604631
}
605632

606633

634+
@Controller
635+
private static class BatchMappingBookController {
636+
637+
@QueryMapping
638+
public List<Book> books() {
639+
return Collections.emptyList();
640+
}
641+
642+
@BatchMapping
643+
public Mono<Map<Book, Author>> author(List<Book> books) {
644+
return Mono.empty();
645+
}
646+
}
647+
648+
607649
@Controller
608650
static class TeamController {
609651
@QueryMapping
@@ -684,7 +726,7 @@ public SchemaInspectionReportAssert hasUnmappedDataFetcherCount(int expected) {
684726
public SchemaInspectionReportAssert hasSkippedTypeCount(int expected) {
685727
isNotNull();
686728
if (this.actual.skippedTypes().size() != expected) {
687-
failWithMessage("Expected %s skipped types, found %d.", expected, this.actual.skippedTypes());
729+
failWithMessage("Expected %s skipped types, found %s.", expected, this.actual.skippedTypes());
688730
}
689731
return this;
690732
}

0 commit comments

Comments
 (0)