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

Author: bclozel

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

@@ -446,6 +446,8 @@ static class SchemaMappingDataFetcher implements SelfDescribingDataFetcher<Objec
 
 		private final boolean subscription;
 
+		private final boolean usesDataLoader;
+
 		SchemaMappingDataFetcher(
 				DataFetcherMappingInfo info, HandlerMethodArgumentResolverComposite argumentResolvers,
 				@Nullable ValidationHelper helper, HandlerDataFetcherExceptionResolver exceptionResolver,
@@ -462,6 +464,17 @@ static class SchemaMappingDataFetcher implements SelfDescribingDataFetcher<Objec
 			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 @@ private <T> Publisher<T> handleSubscriptionError(
 					.switchIfEmpty(Mono.error(ex));
 		}
 
+		@Override
+		public boolean isBatchLoading() {
+			return this.usesDataLoader;
+		}
+
 		@Override
 		public String toString() {
 			return getDescription();
@@ -595,6 +613,11 @@ public Object get(DataFetchingEnvironment env) {
 					dataLoader.load(env.getSource()));
 		}
 
+		@Override
+		public boolean isBatchLoading() {
+			return true;
+		}
+
 		@Override
 		public String toString() {
 			return getDescription();

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 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 @@
  * <li>Re-establish Reactor Context passed via {@link ExecutionInput}.
  * <li>Re-establish ThreadLocal context passed via {@link ExecutionInput}.
  * <li>Resolve exceptions from a GraphQL subscription {@link Publisher}.
+ * <li>Propagate the cancellation signal to {@code DataFetcher} from the transport layer.
  * </ul>
  *
  * @author Rossen Stoyanchev
+ * @author Brian Clozel
  */
-final class ContextDataFetcherDecorator implements DataFetcher<Object> {
+class ContextDataFetcherDecorator implements DataFetcher<Object> {
 
 	private final DataFetcher<?> delegate;
 
@@ -146,6 +150,17 @@ static GraphQLTypeVisitor createVisitor(List<SubscriptionExceptionResolver> reso
 		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 @@ public TraversalControl visitGraphQLFieldDefinition(
 
 			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 @@ private boolean applyDecorator(DataFetcher<?> dataFetcher) {
 		}
 	}
 
+	private static final class SelfDescribingDecorator extends ContextDataFetcherDecorator implements SelfDescribingDataFetcher<Object> {
+
+		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<String, ResolvableType> getArguments() {
+			return this.selfDescribingDataFetcher.getArguments();
+		}
+
+		@Override
+		public ResolvableType getReturnType() {
+			return this.selfDescribingDataFetcher.getReturnType();
+		}
+
+		@Override
+		public String getDescription() {
+			return this.selfDescribingDataFetcher.getDescription();
+		}
+	}
+
 }

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 @@ default Map<String, ResolvableType> getArguments() {
 		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;
+	}
+
 }

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.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 @@ void registerWithDefaultCoordinates() {
 				"Book.authorCallableMap", "Book.authorEnvironment");
 	}
 
+	@Test
+	void dataFetchersMarkedAsBatchLoading() {
+		Map<String, Map<String, DataFetcher>> dataFetcherMap =
+				initRuntimeWiringBuilder(BookController.class).build().getDataFetchers();
+		assertThat(dataFetcherMap.get("Book").values()).allMatch(dataFetcher ->
+				dataFetcher instanceof SelfDescribingDataFetcher<?> selfDescribingDataFetcher
+						&& selfDescribingDataFetcher.isBatchLoading());
+	}
+
 	@Test
 	void registerWithMaxBatchSize() {
 

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.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 @@ void registerWithExplicitCoordinates() {
 		assertMapping(map, "Book.authorCustomized", "authorWithNonMatchingMethodName");
 	}
 
+	@Test
+	void batchLoadingDataFetchers() {
+		Map<String, Map<String, DataFetcher>> map =
+				initRuntimeWiringBuilder(BookController.class).build().getDataFetchers();
+		Map<String, DataFetcher> 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 Flux<Book> bookSearchWithNonMatchingMethodName(@Argument String author) {
 		public Author authorWithNonMatchingMethodName(Book book) {
 			return null;
 		}
+
+	}
+
+	@Controller
+	private static class BatchLoadingController {
+
+		@SchemaMapping
+		public CompletableFuture<Author> authorBatch(Book book, DataLoader<Long, Author> loader) {
+			return null;
+		}
 	}
 
 }