Propagate context to Dataloader in @BatchMapping

Prior to this commit, `@BatchMapping` controller methods would be
automatically registered as date fetchers delegating to data loader
calls. Those calls would not include the current local context or main
context. As a result, injecting the `BatchLoaderEnvironment` in the
controller method signature would not contain the `getKeyContext()`.

This commit ensures that dataloader calls not only use the current
source, but also the current local context/main context so that it will
be present in the key contexts map.

Fixes gh-1071
This commit is contained in:
Brian Clozel
2024-10-17 12:21:52 +02:00
parent b8b93a0df5
commit 6849d3bf38
2 changed files with 45 additions and 3 deletions

View File

@@ -755,7 +755,7 @@ public class AnnotatedControllerConfigurer implements ApplicationContextAware, I
public Object get(DataFetchingEnvironment env) {
DataLoader<?, ?> dataLoader = env.getDataLoaderRegistry().getDataLoader(this.dataLoaderKey);
Assert.state(dataLoader != null, "No DataLoader for key '" + this.dataLoaderKey + "'");
return dataLoader.load(env.getSource());
return dataLoader.load(env.getSource(), (env.getLocalContext() != null) ? env.getLocalContext() : env.getGraphQlContext());
}
}

View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2023 the original author or authors.
* Copyright 2002-2024 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.
@@ -24,6 +24,9 @@ import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import graphql.GraphQLContext;
import org.dataloader.BatchLoaderEnvironment;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
@@ -126,12 +129,33 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
}
}
@Test
void shouldBindKeyContextsToEnvironment() {
String document = "{ " +
" courses { " +
" id" +
" name" +
" students {" +
" id" +
" firstName" +
" lastName" +
" }" +
" }" +
"}";
Mono<ExecutionGraphQlResponse> responseMono = createGraphQlService(new BatchKeyContextsController()).execute(document);
List<Course> actualCourses = ResponseHelper.forResponse(responseMono).toList("courses", Course.class);
List<Course> courses = Course.allCourses();
assertThat(actualCourses).hasSize(courses.size());
}
@Controller
private static class BatchMonoMapController extends CourseController {
@BatchMapping
public Mono<Map<Course, Person>> instructor(List<Course> courses) {
public Mono<Map<Course, Person>> instructor(List<Course> courses, BatchLoaderEnvironment environment) {
return Flux.fromIterable(courses).collect(Collectors.toMap(Function.identity(), Course::instructor));
}
@@ -203,5 +227,23 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
}
@Controller
private static class BatchKeyContextsController extends CourseController {
@BatchMapping
public List<Person> instructor(List<Course> courses, BatchLoaderEnvironment environment) {
assertThat(environment.getKeyContexts().keySet()).containsAll(Course.allCourses());
assertThat(environment.getKeyContexts().values()).allSatisfy(value -> assertThat(value).isInstanceOf(GraphQLContext.class));
return courses.stream().map(Course::instructor).collect(Collectors.toList());
}
@BatchMapping
public List<List<Person>> students(List<Course> courses, BatchLoaderEnvironment environment) {
assertThat(environment.getKeyContexts().keySet()).containsAll(Course.allCourses());
assertThat(environment.getKeyContexts().values()).allSatisfy(value -> assertThat(value).isInstanceOf(GraphQLContext.class));
return courses.stream().map(Course::students).collect(Collectors.toList());
}
}
}