Spring/spring-graphql
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add batch loading info to SelfDescribingDataFetcher

Browse Source 
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
This commit is contained in:
Brian Clozel
2025-04-07 15:29:57 +02:00
parent 462e73ede6
commit 19a935fc18
5 changed files with 123 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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();

51
spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java
View File

@@ -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;
* <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 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
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<Object> {
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<Object> {
}
}
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();
}
}
}

11
spring-graphql/src/main/java/org/springframework/graphql/execution/SelfDescribingDataFetcher.java
View File

@@ -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<T> extends DataFetcher<T> {
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;
}
}

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

@@ -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<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() {

31
spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingDetectionTests.java
View File

@@ -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<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 class SchemaMappingDetectionTests {
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;
}
}
}