From 19a935fc18fbd38bedcc87c51fbb8fdce5e9c79f Mon Sep 17 00:00:00 2001 From: Brian Clozel Date: Mon, 7 Apr 2025 15:29:57 +0200 Subject: [PATCH] Add batch loading info to SelfDescribingDataFetcher Prior to this commit, the `SelfDescribingDataFetcher` would augment the `DataFetcher` contract and provide more information about the data fetcher itself. This commit adds a new `isBatchLoading` method to indicate whether the current data fetcher is using a `DataLoader` for fetching elements. In Spring for GraphQL, this can typically happen if the method is annotated with `@BatchMapping` or if the `@SchemaMapping` method as a `DataLoader` parameter. This change is required for instrumentation purposes: such data fetchers should not be instrumented as data fetching operations, but instead delegate to the `DataLoaderRegistry` being itself instrumented. Closes gh-1176 --- .../AnnotatedControllerConfigurer.java | 23 +++++++++ .../ContextDataFetcherDecorator.java | 51 ++++++++++++++++++- .../execution/SelfDescribingDataFetcher.java | 11 +++- .../support/BatchMappingDetectionTests.java | 12 ++++- .../support/SchemaMappingDetectionTests.java | 31 ++++++++++- 5 files changed, 123 insertions(+), 5 deletions(-) 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 ea06482a..9bed5816 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 @@ -446,6 +446,8 @@ public class AnnotatedControllerConfigurer private final boolean subscription; + private final boolean usesDataLoader; + SchemaMappingDataFetcher( DataFetcherMappingInfo info, HandlerMethodArgumentResolverComposite argumentResolvers, @Nullable ValidationHelper helper, HandlerDataFetcherExceptionResolver exceptionResolver, @@ -462,6 +464,17 @@ public class AnnotatedControllerConfigurer this.executor = executor; this.invokeAsync = invokeAsync; this.subscription = this.mappingInfo.getCoordinates().getTypeName().equalsIgnoreCase("Subscription"); + this.usesDataLoader = hasDataLoaderParameter(); + } + + private boolean hasDataLoaderParameter() { + Method handlerMethod = this.mappingInfo.getHandlerMethod().getMethod(); + for (Class parameterType : handlerMethod.getParameterTypes()) { + if (DataLoader.class.equals(parameterType)) { + return true; + } + } + return false; } @Override @@ -551,6 +564,11 @@ public class AnnotatedControllerConfigurer .switchIfEmpty(Mono.error(ex)); } + @Override + public boolean isBatchLoading() { + return this.usesDataLoader; + } + @Override public String toString() { return getDescription(); @@ -595,6 +613,11 @@ public class AnnotatedControllerConfigurer dataLoader.load(env.getSource())); } + @Override + public boolean isBatchLoading() { + return true; + } + @Override public String toString() { return getDescription(); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java index f680738f..800ab5f1 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java @@ -17,6 +17,7 @@ package org.springframework.graphql.execution; import java.util.List; +import java.util.Map; import graphql.ExecutionInput; import graphql.GraphQLContext; @@ -40,6 +41,7 @@ import org.reactivestreams.Publisher; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import org.springframework.core.ResolvableType; import org.springframework.lang.Nullable; import org.springframework.util.Assert; @@ -51,11 +53,13 @@ import org.springframework.util.Assert; *
  • Re-establish Reactor Context passed via {@link ExecutionInput}. *
  • Re-establish ThreadLocal context passed via {@link ExecutionInput}. *
  • Resolve exceptions from a GraphQL subscription {@link Publisher}. + *
  • Propagate the cancellation signal to {@code DataFetcher} from the transport layer. * * * @author Rossen Stoyanchev + * @author Brian Clozel */ -final class ContextDataFetcherDecorator implements DataFetcher { +class ContextDataFetcherDecorator implements DataFetcher { private final DataFetcher delegate; @@ -146,6 +150,17 @@ final class ContextDataFetcherDecorator implements DataFetcher { return new ContextTypeVisitor(resolvers); } + private static ContextDataFetcherDecorator decorate( + DataFetcher delegate, boolean handlesSubscription, + SubscriptionExceptionResolver subscriptionExceptionResolver) { + if (delegate instanceof SelfDescribingDataFetcher selfDescribingDataFetcher) { + return new SelfDescribingDecorator(selfDescribingDataFetcher, handlesSubscription, subscriptionExceptionResolver); + } + else { + return new ContextDataFetcherDecorator(delegate, handlesSubscription, subscriptionExceptionResolver); + } + } + /** * Type visitor to apply {@link ContextDataFetcherDecorator}. @@ -171,7 +186,7 @@ final class ContextDataFetcherDecorator implements DataFetcher { if (applyDecorator(dataFetcher)) { boolean handlesSubscription = visitorHelper.isSubscriptionType(parent); - dataFetcher = new ContextDataFetcherDecorator(dataFetcher, handlesSubscription, this.exceptionResolver); + dataFetcher = ContextDataFetcherDecorator.decorate(dataFetcher, handlesSubscription, this.exceptionResolver); codeRegistry.dataFetcher(fieldCoordinates, dataFetcher); } @@ -192,4 +207,36 @@ final class ContextDataFetcherDecorator implements DataFetcher { } } + private static final class SelfDescribingDecorator extends ContextDataFetcherDecorator implements SelfDescribingDataFetcher { + + private final SelfDescribingDataFetcher selfDescribingDataFetcher; + + private SelfDescribingDecorator( + SelfDescribingDataFetcher delegate, boolean subscription, + SubscriptionExceptionResolver subscriptionExceptionResolver) { + super(delegate, subscription, subscriptionExceptionResolver); + this.selfDescribingDataFetcher = delegate; + } + + @Override + public boolean isBatchLoading() { + return this.selfDescribingDataFetcher.isBatchLoading(); + } + + @Override + public Map getArguments() { + return this.selfDescribingDataFetcher.getArguments(); + } + + @Override + public ResolvableType getReturnType() { + return this.selfDescribingDataFetcher.getReturnType(); + } + + @Override + public String getDescription() { + return this.selfDescribingDataFetcher.getDescription(); + } + } + } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/SelfDescribingDataFetcher.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/SelfDescribingDataFetcher.java index e37764ee..e43456f1 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/SelfDescribingDataFetcher.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/SelfDescribingDataFetcher.java @@ -1,5 +1,5 @@ /* - * Copyright 2020-2024 the original author or authors. + * Copyright 2020-2025 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. @@ -59,4 +59,13 @@ public interface SelfDescribingDataFetcher extends DataFetcher { return Collections.emptyMap(); } + /** + * Return whether this {@code DataFetcher} is using batch loading. + * @return {@code true} if the data fetcher is batch loading elements + * @since 1.4.0 + */ + default boolean isBatchLoading() { + return false; + } + } 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 abb925d7..fc15eeb0 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 @@ -1,5 +1,5 @@ /* - * Copyright 2002-2022 the original author or authors. + * Copyright 2002-2025 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. @@ -41,6 +41,7 @@ import org.springframework.graphql.data.method.annotation.BatchMapping; import org.springframework.graphql.data.method.annotation.SchemaMapping; import org.springframework.graphql.execution.BatchLoaderRegistry; import org.springframework.graphql.execution.DefaultBatchLoaderRegistry; +import org.springframework.graphql.execution.SelfDescribingDataFetcher; import org.springframework.stereotype.Controller; import static org.assertj.core.api.Assertions.assertThat; @@ -76,6 +77,15 @@ public class BatchMappingDetectionTests { "Book.authorCallableMap", "Book.authorEnvironment"); } + @Test + void dataFetchersMarkedAsBatchLoading() { + Map> dataFetcherMap = + initRuntimeWiringBuilder(BookController.class).build().getDataFetchers(); + assertThat(dataFetcherMap.get("Book").values()).allMatch(dataFetcher -> + dataFetcher instanceof SelfDescribingDataFetcher selfDescribingDataFetcher + && selfDescribingDataFetcher.isBatchLoading()); + } + @Test void registerWithMaxBatchSize() { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingDetectionTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingDetectionTests.java index a7f68427..7483668d 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingDetectionTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingDetectionTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2021 the original author or authors. + * Copyright 2002-2025 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. @@ -17,10 +17,12 @@ package org.springframework.graphql.data.method.annotation.support; import java.util.Map; +import java.util.concurrent.CompletableFuture; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; import graphql.schema.idl.RuntimeWiring; +import org.dataloader.DataLoader; import org.junit.jupiter.api.Test; import reactor.core.publisher.Flux; @@ -32,6 +34,7 @@ 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.execution.SelfDescribingDataFetcher; import org.springframework.stereotype.Controller; import org.springframework.util.StringUtils; @@ -81,6 +84,22 @@ public class SchemaMappingDetectionTests { assertMapping(map, "Book.authorCustomized", "authorWithNonMatchingMethodName"); } + @Test + void batchLoadingDataFetchers() { + Map> map = + initRuntimeWiringBuilder(BookController.class).build().getDataFetchers(); + Map queries = map.get("Book"); + assertThat(queries.values()).allMatch(dataFetcher -> + dataFetcher instanceof SelfDescribingDataFetcher selfDescribingDataFetcher + && !selfDescribingDataFetcher.isBatchLoading()); + + map = initRuntimeWiringBuilder(BatchLoadingController.class).build().getDataFetchers(); + queries = map.get("Book"); + assertThat(queries.values()).allMatch(dataFetcher -> + dataFetcher instanceof SelfDescribingDataFetcher selfDescribingDataFetcher + && selfDescribingDataFetcher.isBatchLoading()); + } + private RuntimeWiring.Builder initRuntimeWiringBuilder(Class handlerType) { AnnotationConfigApplicationContext appContext = new AnnotationConfigApplicationContext(); appContext.registerBean(handlerType); @@ -152,6 +171,16 @@ public class SchemaMappingDetectionTests { public Author authorWithNonMatchingMethodName(Book book) { return null; } + + } + + @Controller + private static class BatchLoadingController { + + @SchemaMapping + public CompletableFuture authorBatch(Book book, DataLoader loader) { + return null; + } } }