Support for @BatchMapping methods

See gh-130
This commit is contained in:
Rossen Stoyanchev
2021-09-27 07:29:12 +01:00
parent fbe87359ac
commit 56ca50e6f3
6 changed files with 862 additions and 44 deletions

View File

@@ -0,0 +1,108 @@
/*
* Copyright 2002-2021 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.data.method;
import java.util.Collection;
import java.util.Map;
import org.dataloader.BatchLoaderEnvironment;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.core.CollectionFactory;
import org.springframework.core.MethodParameter;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
/**
* An extension of {@link HandlerMethod} for annotated handler methods adapted
* to a batch loader function with a list of values and {@link BatchLoaderEnvironment}
* as their input.
*
* @author Rossen Stoyanchev
* @since 1.0.0
*/
public class BatchLoadHandlerMethod extends InvocableHandlerMethodSupport {
public BatchLoadHandlerMethod(HandlerMethod handlerMethod) {
super(handlerMethod);
}
/**
* Invoke the underlying batch loading method, resolving its arguments from
* the given keys for batch loading and the {@link BatchLoaderEnvironment}.
*
* @param keys the batch loading keys
* @param environment the environment available to batch loaders
* @return a {@code Flux} of values or a {@code Mono} with map of key-value pairs.
*/
@Nullable
public <K> Object invoke(Collection<K> keys, BatchLoaderEnvironment environment) {
MethodParameter[] parameters = getMethodParameters();
Assert.notEmpty(parameters, "Batch loading methods should have at least " +
"one argument with the List of parent objects: " + getBridgedMethod().toGenericString());
Object[] args = new Object[parameters.length];
for (int i = 0; i < parameters.length; i++) {
args[i] = resolveArgument(parameters[i], keys, environment);
}
Object result;
try {
result = doInvoke(args);
}
catch (Exception ex) {
throw new IllegalStateException("...", ex);
}
if (result != null) {
if (result instanceof Collection) {
return Flux.fromIterable((Collection<?>) result);
}
else if (result instanceof Map) {
return Mono.just(result);
}
}
return result;
}
public <K> Object resolveArgument(
MethodParameter parameter, Collection<K> keys, BatchLoaderEnvironment environment) {
Class<?> parameterType = parameter.getParameterType();
if (Collection.class.isAssignableFrom(parameterType)) {
if (parameterType.isInstance(keys)) {
return keys;
}
Class<?> elementType = parameter.nested().getNestedParameterType();
Collection<K> collection = CollectionFactory.createCollection(parameterType, elementType, keys.size());
collection.addAll(keys);
return collection;
}
if (parameterType.isInstance(environment)) {
return environment;
}
throw new IllegalStateException(formatArgumentError(parameter, "Unexpected argument type."));
}
}

View File

@@ -61,7 +61,7 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
/**
* Invoke the method after resolving its argument values in the context of
* the given environment.
* the given {@link DataFetchingEnvironment}.
* <p>Argument values are commonly resolved through
* {@link HandlerMethodArgumentResolver HandlerMethodArgumentResolvers}.
* The {@code providedArgs} parameter however may supply argument values to

View File

@@ -0,0 +1,91 @@
/*
* Copyright 2002-2021 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.data.method.annotation;
import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import org.springframework.core.annotation.AliasFor;
/**
* Annotation for handler methods that batch load field values, given a list
* of source/parent values. For example:
*
* <pre class="code">
* &#064;BatchMapping
* public Flux&lt;Author&gt; author(List&lt;Book&gt; books) {
* // ...
* }
* </pre>
*
* <p>The annotated method is registered as a batch loading function and along
* with it, a {@link graphql.schema.DataFetcher} for the field is registered
* transparently that looks up the field through the registered
* {@code DataLoader}.
*
* <p>Effectively, a shortcut for:
*
* <pre class="code">
* &#064;Controller
* public class BookController {
*
* public BookController(BatchLoaderRegistry registry) {
* registry.forTypePair(Long.class, Author.class).registerBatchLoader((ids, env) -> ...);
* }
*
* &#064;SchemaMapping
* public Author author(Book book, DataLoader&lt;Long, Author&gt; dataLoader) {
* return dataLoader.load(book.getAuthorId());
* }
*
* }
* </pre>
*
* @author Rossen Stoyanchev
* @since 1.0.0
*/
@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface BatchMapping {
/**
* Customize the name of the GraphQL field to bind to.
* <p>By default, if not specified, this is initialized from the method name.
*/
@AliasFor("value")
String field() default "";
/**
* Effectively an alias for {@link #field()}.
*/
@AliasFor("field")
String value() default "";
/**
* Customizes the name of the parent/container type for the GraphQL field.
* <p>By default, if not specified, it is derived from the class name of the
* List of source/parent values injected into the handler method.
* <p>This value for this attribute can be initialized from a class-level
* {@link SchemaMapping @SchemaMapping}. When used on both levels, the one
* on the method level overrides the one at the class level.
*/
String typeName() default "";
}

View File

@@ -15,12 +15,16 @@
*/
package org.springframework.graphql.data.method.annotation.support;
import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.lang.reflect.Type;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import graphql.schema.DataFetcher;
@@ -29,6 +33,9 @@ import graphql.schema.FieldCoordinates;
import graphql.schema.idl.RuntimeWiring;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dataloader.DataLoader;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.factory.InitializingBean;
@@ -37,14 +44,15 @@ import org.springframework.context.ApplicationContextAware;
import org.springframework.core.KotlinDetector;
import org.springframework.core.MethodIntrospector;
import org.springframework.core.MethodParameter;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.ResolvableType;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.graphql.data.method.BatchLoadHandlerMethod;
import org.springframework.graphql.data.method.DataFetcherHandlerMethod;
import org.springframework.graphql.data.method.HandlerMethod;
import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
import org.springframework.graphql.data.method.InvocableHandlerMethod;
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.RuntimeWiringConfigurer;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Controller;
@@ -65,7 +73,6 @@ public class AnnotatedDataFetcherConfigurer
private final static Log logger = LogFactory.getLog(AnnotatedDataFetcherConfigurer.class);
/**
* Bean name prefix for target beans behind scoped proxies. Used to exclude those
* targets from handler method detection, in favor of the corresponding proxies.
@@ -78,9 +85,6 @@ public class AnnotatedDataFetcherConfigurer
*/
private static final String SCOPED_TARGET_NAME_PREFIX = "scopedTarget.";
private static final ResolvableType MAP_RESOLVABLE_TYPE =
ResolvableType.forType(new ParameterizedTypeReference<Map<String, Object>>() {});
@Nullable
private ApplicationContext applicationContext;
@@ -119,12 +123,18 @@ public class AnnotatedDataFetcherConfigurer
@Override
public void configure(RuntimeWiring.Builder builder) {
Assert.state(this.argumentResolvers != null, "`argumentResolvers` is not initialized");
findHandlerMethods().forEach((info) -> {
FieldCoordinates coordinates = info.getCoordinates();
HandlerMethod handlerMethod = info.getHandlerMethod();
DataFetcher<?> dataFetcher = new SchemaMappingDataFetcher(coordinates, handlerMethod, this.argumentResolvers);
builder.type(coordinates.getTypeName(), typeBuilder ->
typeBuilder.dataFetcher(coordinates.getFieldName(), dataFetcher));
DataFetcher<?> dataFetcher;
if (!info.isBatchMapping()) {
dataFetcher = new SchemaMappingDataFetcher(info, this.argumentResolvers);
}
else {
String dataLoaderKey = registerBatchLoader(info);
dataFetcher = new BatchMappingDataFetcher(dataLoaderKey);
}
builder.type(info.getCoordinates().getTypeName(), typeBuilder ->
typeBuilder.dataFetcher(info.getCoordinates().getFieldName(), dataFetcher));
});
}
@@ -186,29 +196,59 @@ public class AnnotatedDataFetcherConfigurer
@Nullable
private MappingInfo getMappingInfo(Method method, Object handler, Class<?> handlerType) {
SchemaMapping annotation = AnnotatedElementUtils.findMergedAnnotation(method, SchemaMapping.class);
if (annotation == null) {
Set<Annotation> annotations = AnnotatedElementUtils.findAllMergedAnnotations(
method, new LinkedHashSet<>(Arrays.asList(BatchMapping.class, SchemaMapping.class)));
if (annotations.isEmpty()) {
return null;
}
String typeName = annotation.typeName();
String field = (StringUtils.hasText(annotation.field()) ? annotation.field() : method.getName());
if (annotations.size() != 1) {
throw new IllegalArgumentException(
"Expected either @BatchMapping or @SchemaMapping, not both: " + method.toGenericString());
}
String typeName;
String field;
boolean batchMapping = false;
HandlerMethod handlerMethod = createHandlerMethod(method, handler, handlerType);
Annotation annotation = annotations.iterator().next();
if (annotation instanceof SchemaMapping) {
SchemaMapping mapping = (SchemaMapping) annotation;
typeName = mapping.typeName();
field = (StringUtils.hasText(mapping.field()) ? mapping.field() : method.getName());
}
else {
BatchMapping mapping = (BatchMapping) annotation;
typeName = mapping.typeName();
field = (StringUtils.hasText(mapping.field()) ? mapping.field() : method.getName());
batchMapping = true;
}
if (!StringUtils.hasText(typeName)) {
SchemaMapping mapping = AnnotatedElementUtils.findMergedAnnotation(handlerType, SchemaMapping.class);
if (mapping != null) {
typeName = annotation.typeName();
typeName = mapping.typeName();
}
}
if (!StringUtils.hasText(typeName)) {
Assert.state(this.argumentResolvers != null, "`argumentResolvers` is not initialized");
for (MethodParameter parameter : handlerMethod.getMethodParameters()) {
HandlerMethodArgumentResolver resolver = this.argumentResolvers.getArgumentResolver(parameter);
if (resolver instanceof SourceMethodArgumentResolver) {
typeName = parameter.getParameterType().getSimpleName();
break;
if (!batchMapping) {
Assert.state(this.argumentResolvers != null, "`argumentResolvers` is not initialized");
HandlerMethodArgumentResolver resolver = this.argumentResolvers.getArgumentResolver(parameter);
if (resolver instanceof SourceMethodArgumentResolver) {
typeName = parameter.getParameterType().getSimpleName();
break;
}
}
else {
if (Collection.class.isAssignableFrom(parameter.getParameterType())) {
typeName = parameter.nested().getNestedParameterType().getSimpleName();
break;
}
}
}
}
@@ -217,7 +257,7 @@ public class AnnotatedDataFetcherConfigurer
"No parentType specified, and a source/parent method argument was also not found: " +
handlerMethod.getShortLogMessage());
return new MappingInfo(typeName, field, handlerMethod);
return new MappingInfo(typeName, field, batchMapping, handlerMethod);
}
private HandlerMethod createHandlerMethod(Method method, Object handler, Class<?> handlerType) {
@@ -227,40 +267,79 @@ public class AnnotatedDataFetcherConfigurer
new HandlerMethod(handler, invocableMethod));
}
private String formatMappings(Class<?> handlerType, Collection<MappingInfo> mappings) {
private String formatMappings(Class<?> handlerType, Collection<MappingInfo> infos) {
String formattedType = Arrays.stream(ClassUtils.getPackageName(handlerType).split("\\."))
.map(p -> p.substring(0, 1))
.collect(Collectors.joining(".", "", "." + handlerType.getSimpleName()));
return mappings.stream()
return infos.stream()
.map(mappingInfo -> {
Method method = mappingInfo.getHandlerMethod().getMethod();
String methodParameters = Arrays.stream(method.getParameterTypes())
.map(Class::getSimpleName)
String methodParameters = Arrays.stream(method.getGenericParameterTypes())
.map(Type::getTypeName)
.collect(Collectors.joining(",", "(", ")"));
return mappingInfo.getCoordinates() + " => " + method.getName() + methodParameters;
})
.collect(Collectors.joining("\n\t", "\n\t" + formattedType + ":" + "\n\t", ""));
}
@SuppressWarnings("unchecked")
private <P, F> String registerBatchLoader(MappingInfo info) {
if (!info.isBatchMapping()) {
throw new IllegalArgumentException("Not a @BatchMapping method: " + info);
}
String dataLoaderKey = info.getCoordinates().toString();
BatchLoadHandlerMethod invocable = new BatchLoadHandlerMethod(info.getHandlerMethod());
BatchLoaderRegistry registry = obtainApplicationContext().getBean(BatchLoaderRegistry.class);
Class<?> clazz = info.getHandlerMethod().getReturnType().getParameterType();
if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(clazz)) {
registry.<P,F>forName(dataLoaderKey).registerBatchLoader((values, env) ->
(Flux<F>) invocable.invoke(values, env));
}
else if (clazz.equals(Mono.class) || clazz.equals(Map.class)) {
registry.<P,F>forName(dataLoaderKey).registerMappedBatchLoader((values, env) ->
(Mono<Map<P, F>>) invocable.invoke(values, env));
}
else {
throw new IllegalStateException("@BatchMapping method is expected to return " +
"Flux<V>, List<V>, Mono<Map<K, V>>, or Map<K, V>: " + info.getHandlerMethod());
}
return dataLoaderKey;
}
private static class MappingInfo {
private final FieldCoordinates coordinates;
private final boolean batchMapping;
private final HandlerMethod handlerMethod;
public MappingInfo(String typeName, String field, HandlerMethod handlerMethod) {
public MappingInfo(String typeName, String field, boolean batchMapping, HandlerMethod handlerMethod) {
this.coordinates = FieldCoordinates.coordinates(typeName, field);
this.handlerMethod = handlerMethod;
this.batchMapping = batchMapping;
}
public FieldCoordinates getCoordinates() {
return this.coordinates;
}
public boolean isBatchMapping() {
return this.batchMapping;
}
public HandlerMethod getHandlerMethod() {
return this.handlerMethod;
}
@Override
public String toString() {
return this.coordinates + " -> " + this.handlerMethod.toString();
}
}
@@ -269,45 +348,53 @@ public class AnnotatedDataFetcherConfigurer
*/
static class SchemaMappingDataFetcher implements DataFetcher<Object> {
private final FieldCoordinates coordinates;
private final HandlerMethod handlerMethod;
private final MappingInfo info;
private final HandlerMethodArgumentResolverComposite argumentResolvers;
public SchemaMappingDataFetcher(FieldCoordinates coordinates, HandlerMethod handlerMethod,
HandlerMethodArgumentResolverComposite resolvers) {
this.coordinates = coordinates;
this.handlerMethod = handlerMethod;
public SchemaMappingDataFetcher(MappingInfo info, HandlerMethodArgumentResolverComposite resolvers) {
this.info = info;
this.argumentResolvers = resolvers;
}
/**
* Return the {@link FieldCoordinates} the HandlerMethod is mapped to.
*/
public FieldCoordinates getCoordinates() {
return this.coordinates;
return this.info.getCoordinates();
}
/**
* Return the {@link HandlerMethod} used to fetch data.
*/
public HandlerMethod getHandlerMethod() {
return this.handlerMethod;
return this.info.getHandlerMethod();
}
@Override
@SuppressWarnings("ConstantConditions")
public Object get(DataFetchingEnvironment environment) throws Exception {
return new DataFetcherHandlerMethod(getHandlerMethod(), this.argumentResolvers).invoke(environment);
}
}
InvocableHandlerMethod invocable =
new InvocableHandlerMethod(this.handlerMethod.createWithResolvedBean(), this.argumentResolvers);
return invocable.invoke(environment);
static class BatchMappingDataFetcher implements DataFetcher<Object> {
private final String dataLoaderKey;
public BatchMappingDataFetcher(String dataLoaderKey) {
this.dataLoaderKey = dataLoaderKey;
}
@Override
public Object get(DataFetchingEnvironment env) {
DataLoader<?, ?> dataLoader = env.getDataLoaderRegistry().getDataLoader(this.dataLoaderKey);
if (dataLoader == null) {
throw new IllegalStateException("No DataLoader for key '" + this.dataLoaderKey + "'");
}
return dataLoader.load(env.getSource());
}
}

View File

@@ -0,0 +1,146 @@
/*
* Copyright 2002-2021 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.data.method.annotation.support;
import java.util.List;
import java.util.Map;
import graphql.schema.DataFetcher;
import graphql.schema.idl.RuntimeWiring;
import org.dataloader.BatchLoaderEnvironment;
import org.dataloader.DataLoaderRegistry;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.graphql.Author;
import org.springframework.graphql.Book;
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.stereotype.Controller;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
* Unit tests for {@link AnnotatedDataFetcherConfigurer}, focusing on detection
* and mapping of handler methods to schema fields.
*
* @author Rossen Stoyanchev
*/
@SuppressWarnings({"rawtypes", "unused"})
public class BatchMappingDetectionTests {
private final DefaultBatchLoaderRegistry batchLoaderRegistry = new DefaultBatchLoaderRegistry();
@Test
void registerWithDefaultCoordinates() {
Map<String, Map<String, DataFetcher>> map =
initRuntimeWiringBuilder(BookController.class).build().getDataFetchers();
assertThat(map).containsOnlyKeys("Book");
assertThat(map.get("Book")).containsOnlyKeys(
"authorFlux", "authorList", "authorMonoMap", "authorMap", "authorEnvironment");
DataLoaderRegistry registry = new DataLoaderRegistry();
this.batchLoaderRegistry.registerDataLoaders(registry);
assertThat(registry.getDataLoadersMap()).containsOnlyKeys(
"Book.authorFlux", "Book.authorList", "Book.authorMonoMap", "Book.authorMap", "Book.authorEnvironment");
}
@Test
void invalidReturnType() {
assertThatThrownBy(() -> initRuntimeWiringBuilder(InvalidReturnTypeController.class).build())
.hasMessageStartingWith("@BatchMapping method is expected to return");
}
@Test
void schemaAndBatch() {
assertThatThrownBy(() -> initRuntimeWiringBuilder(SchemaAndBatchMappingController.class).build())
.hasMessageStartingWith("Expected either @BatchMapping or @SchemaMapping, not both");
}
private RuntimeWiring.Builder initRuntimeWiringBuilder(Class<?> handlerType) {
AnnotationConfigApplicationContext appContext = new AnnotationConfigApplicationContext();
appContext.registerBean(handlerType);
appContext.registerBean(BatchLoaderRegistry.class, () -> this.batchLoaderRegistry);
appContext.refresh();
AnnotatedDataFetcherConfigurer configurer = new AnnotatedDataFetcherConfigurer();
configurer.setApplicationContext(appContext);
configurer.afterPropertiesSet();
RuntimeWiring.Builder wiringBuilder = RuntimeWiring.newRuntimeWiring();
configurer.configure(wiringBuilder);
return wiringBuilder;
}
@Controller
@SuppressWarnings({"ConstantConditions", "unused"})
private static class BookController {
@BatchMapping
public Flux<Author> authorFlux(List<Book> books) {
return null;
}
@BatchMapping
public List<Author> authorList(List<Book> books) {
return null;
}
@BatchMapping
public Mono<Map<Book, Author>> authorMonoMap(List<Book> books) {
return null;
}
@BatchMapping
public Map<Book, Author> authorMap(List<Book> books) {
return null;
}
@BatchMapping
public List<Author> authorEnvironment(BatchLoaderEnvironment environment, List<Book> books) {
return null;
}
}
@Controller
private static class InvalidReturnTypeController {
@BatchMapping
public void authors(List<Book> books) {
}
}
@Controller
private static class SchemaAndBatchMappingController {
@BatchMapping
@SchemaMapping
public void authors(List<Book> books) {
}
}
}

View File

@@ -0,0 +1,386 @@
/*
* Copyright 2002-2021 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.data.method.annotation.support;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import graphql.ExecutionResult;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.graphql.GraphQlService;
import org.springframework.graphql.RequestInput;
import org.springframework.graphql.data.method.annotation.BatchMapping;
import org.springframework.graphql.data.method.annotation.QueryMapping;
import org.springframework.graphql.execution.DataLoaderRegistrar;
import org.springframework.graphql.execution.DefaultBatchLoaderRegistry;
import org.springframework.graphql.execution.ExecutionGraphQlService;
import org.springframework.graphql.execution.GraphQlSource;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Controller;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Test GraphQL requests handled through {@code @BatchMapping} methods.
*
* @author Rossen Stoyanchev
*/
@SuppressWarnings({"unchecked", "unused"})
public class BatchMappingInvocationTests {
private static final Map<Long, Course> courseMap = new HashMap<>();
private static final Map<Long, Person> personMap = new HashMap<>();
static {
Course.save(11L, "Ethical Hacking", 15L, Arrays.asList(22L, 26L, 31L));
Course.save(19L, "Docker and Kubernetes", 17L, Arrays.asList(31L, 39L, 44L, 45L));
Person.save(15L, "Josh", "Kelly");
Person.save(17L, "Albert", "Murray");
Person.save(22L, "Bonnie", "Gray");
Person.save(26L, "John", "Perry");
Person.save(31L, "Alaine", "Baily");
Person.save(39L, "Jeff", "Peterson");
Person.save(44L, "Jared", "Mccarthy");
Person.save(45L, "Benjamin", "Brown");
}
private static final String schema = "" +
"type Query {" +
" courses: [Course]" +
"}" +
"type Course {" +
" id: ID" +
" name: String" +
" instructor: Person" +
" students: [Person]" +
"}" +
"type Person {" +
" id: ID" +
" firstName: String" +
" lastName: String" +
"}";
private static Class<?>[] controllerClasses() {
return new Class[] {
BatchFluxController.class,
BatchListController.class,
BatchMonoMapController.class,
BatchMapController.class
};
}
@ParameterizedTest
@MethodSource("controllerClasses")
void oneToOne(Class<?> controllerClass) {
String query = "{ " +
" courses { " +
" name" +
" instructor {" +
" firstName" +
" lastName" +
" }" +
" }" +
"}";
ExecutionResult result = initGraphQlService(controllerClass, CourseConfig.class)
.execute(new RequestInput(query, null, null))
.block();
Map<String, Object> data = getData(result);
List<Map<String, Object>> actualCourses = (List<Map<String, Object>>) data.get("courses");
List<Course> courses = Course.allCourses();
assertThat(actualCourses).hasSize(courses.size());
for (int i = 0; i < courses.size(); i++) {
Map<String, Object> actualCourse = actualCourses.get(i);
Course course = courses.get(i);
assertThat(actualCourse.get("name")).isEqualTo(course.name());
Map<String, Object> actualInstructor = (Map<String, Object>) actualCourse.get("instructor");
assertThat(actualInstructor.get("firstName")).isEqualTo(course.instructor().firstName());
assertThat(actualInstructor.get("lastName")).isEqualTo(course.instructor().lastName());
}
}
@ParameterizedTest
@MethodSource("controllerClasses")
void oneToMany(Class<?> controllerClass) {
String query = "{ " +
" courses { " +
" name" +
" students {" +
" firstName" +
" lastName" +
" }" +
" }" +
"}";
ExecutionResult result = initGraphQlService(controllerClass, CourseConfig.class)
.execute(new RequestInput(query, null, null))
.block();
Map<String, Object> data = getData(result);
List<Map<String, Object>> actualCourses = (List<Map<String, Object>>) data.get("courses");
List<Course> courses = Course.allCourses();
assertThat(actualCourses).hasSize(courses.size());
for (int i = 0; i < courses.size(); i++) {
Map<String, Object> actualCourse = actualCourses.get(i);
Course course = courses.get(i);
assertThat(actualCourse.get("name")).isEqualTo(course.name());
List<Map<String, Object>> actualStudents = (List<Map<String, Object>>) actualCourse.get("students");
List<Person> students = course.students();
assertThat(actualStudents).hasSize(students.size());
for (int j = 0; j < actualStudents.size(); j++) {
assertThat(actualStudents.get(i).get("firstName")).isEqualTo(students.get(i).firstName());
assertThat(actualStudents.get(i).get("lastName")).isEqualTo(students.get(i).lastName());
}
}
}
private ExecutionGraphQlService initGraphQlService(Class<?>... configClasses) {
AnnotationConfigApplicationContext applicationContext = new AnnotationConfigApplicationContext();
applicationContext.register(configClasses);
applicationContext.refresh();
return applicationContext.getBean(ExecutionGraphQlService.class);
}
private <T> T getData(@Nullable ExecutionResult result) {
assertThat(result).isNotNull();
assertThat(result.getErrors()).isEmpty();
T data = result.getData();
assertThat(data).isNotNull();
return data;
}
private static class CourseController {
@QueryMapping
public Collection<Course> courses() {
return courseMap.values();
}
}
@Controller
private static class BatchFluxController extends CourseController {
@BatchMapping
public Flux<Person> instructor(List<Course> courses) {
return Flux.fromIterable(courses).map(Course::instructor);
}
@BatchMapping
public Flux<List<Person>> students(List<Course> courses) {
return Flux.fromIterable(courses).map(Course::students);
}
}
@Controller
private static class BatchListController extends CourseController {
@BatchMapping
public List<Person> instructor(List<Course> courses) {
return courses.stream().map(Course::instructor).collect(Collectors.toList());
}
@BatchMapping
public List<List<Person>> students(List<Course> courses) {
return courses.stream().map(Course::students).collect(Collectors.toList());
}
}
@Controller
private static class BatchMonoMapController extends CourseController {
@BatchMapping
public Mono<Map<Course, Person>> instructor(List<Course> courses) {
return Flux.fromIterable(Course.allCourses())
.collect(Collectors.toMap(Function.identity(), Course::instructor));
}
@BatchMapping
public Mono<Map<Course, List<Person>>> students(Set<Course> courses) {
return Flux.fromIterable(courses).collect(Collectors.toMap(Function.identity(), Course::students));
}
}
@Controller
private static class BatchMapController extends CourseController {
@BatchMapping
public Map<Course, Person> instructor(List<Course> courses) {
return Course.allCourses().stream().collect(
Collectors.toMap(Function.identity(), Course::instructor));
}
@BatchMapping
public Map<Course, List<Person>> students(List<Course> courses) {
return courses.stream().collect(Collectors.toMap(Function.identity(), Course::students));
}
}
private static class CourseConfig {
@Bean
public GraphQlSource graphQlSource(AnnotatedDataFetcherConfigurer configurer) {
return GraphQlSource.builder()
.schemaResources(new ByteArrayResource(schema.getBytes(StandardCharsets.UTF_8)))
.configureRuntimeWiring(configurer)
.build();
}
@Bean
public GraphQlService graphQlService(GraphQlSource source, DataLoaderRegistrar registrar) {
ExecutionGraphQlService service = new ExecutionGraphQlService(source);
service.addDataLoaderRegistrar(registrar);
return service;
}
@Bean
public AnnotatedDataFetcherConfigurer annotatedDataFetcherConfigurer() {
return new AnnotatedDataFetcherConfigurer();
}
@Bean
public DefaultBatchLoaderRegistry batchLoaderRegistry() {
return new DefaultBatchLoaderRegistry();
}
}
private static class Course {
private final Long id;
private final String name;
private final Long instructorId;
private final List<Long> studentIds;
public Course(Long id, String name, Long instructorId, List<Long> studentIds) {
this.id = id;
this.name = name;
this.instructorId = instructorId;
this.studentIds = studentIds;
}
public String name() {
return this.name;
}
public Long instructorId() {
return this.instructorId;
}
public List<Long> studentIds() {
return this.studentIds;
}
public List<Person> students() {
return this.studentIds.stream().map(personMap::get).collect(Collectors.toList());
}
public Person instructor() {
return personMap.get(this.instructorId);
}
public static void save(Long id, String name, Long instructorId, List<Long> studentIds) {
Course course = new Course(id, name, instructorId, studentIds);
courseMap.put(id, course);
}
public static List<Course> allCourses() {
return new ArrayList<>(courseMap.values());
}
// Course is a key in the DataLoader map
@Override
public boolean equals(Object other) {
if (this == other) {
return true;
}
if (other == null || getClass() != other.getClass()) {
return false;
}
return this.id.equals(((Course) other).id);
}
@Override
public int hashCode() {
return this.id.hashCode();
}
}
private static class Person {
private final Long id;
private final String firstName;
private final String lastName;
public Person(Long id, String firstName, String lastName) {
this.id = id;
this.firstName = firstName;
this.lastName = lastName;
}
public Long id() {
return this.id;
}
public String firstName() {
return this.firstName;
}
public String lastName() {
return this.lastName;
}
public static void save(Long id, String firstName, String lastName) {
Person person = new Person(id, firstName, lastName);
personMap.put(id, person);
}
}
}