Support batched EntityMapping methods

See gh-922

spring-graphql/src/main/java/org/springframework/graphql/data/GraphQlArgumentBinder.java

@@ -142,13 +142,12 @@ public class GraphQlArgumentBinder {
 		Object rawValue = (name != null) ? environment.getArgument(name) : environment.getArguments();
 		boolean isOmitted = (name != null && !environment.getArguments().containsKey(name));
 
-		return bind(name, rawValue, isOmitted, targetType);
+		return bind(rawValue, isOmitted, targetType);
 	}
 
 	/**
 	 * Variant of {@link #bind(DataFetchingEnvironment, String, ResolvableType)}
 	 * with a pre-extracted raw value to bind from.
-	 * @param name the name of an argument, or {@code null} to use the full map
 	 * @param rawValue the raw argument value (Collection, Map, or scalar)
 	 * @param isOmitted {@code true} if the argument was omitted from the input
 	 * and {@code false} if it was provided, but possibly {@code null}
@@ -156,19 +155,13 @@ public class GraphQlArgumentBinder {
 	 * @since 1.3.0
 	 */
 	@Nullable
-	public Object bind(
-			@Nullable String name, @Nullable Object rawValue, boolean isOmitted, ResolvableType targetType)
-			throws BindException {
-
+	public Object bind(@Nullable Object rawValue, boolean isOmitted, ResolvableType targetType) throws BindException {
 		ArgumentsBindingResult bindingResult = new ArgumentsBindingResult(targetType);
-
-		Object value = bindRawValue(
-				"$", rawValue, isOmitted, targetType, targetType.resolve(Object.class), bindingResult);
-
+		Class<?> targetClass = targetType.resolve(Object.class);
+		Object value = bindRawValue("$", rawValue, isOmitted, targetType, targetClass, bindingResult);
 		if (bindingResult.hasErrors()) {
 			throw new BindException(bindingResult);
 		}
-
 		return value;
 	}
 

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntitiesDataFetcher.java

@@ -20,9 +20,11 @@ package org.springframework.graphql.data.federation;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletionException;
 
 import com.apollographql.federation.graphqljava._Entity;
@@ -38,6 +40,7 @@ import reactor.core.publisher.Mono;
 import org.springframework.graphql.data.method.annotation.support.HandlerDataFetcherExceptionResolver;
 import org.springframework.graphql.execution.ErrorType;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 
 /**
  * DataFetcher that handles the "_entities" query by invoking
@@ -65,6 +68,7 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 	public Mono<DataFetcherResult<List<Object>>> get(DataFetchingEnvironment environment) {
 		List<Map<String, Object>> representations = environment.getArgument(_Entity.argumentName);
 
+		Set<String> batched = new HashSet<>();
 		List<Mono<Object>> monoList = new ArrayList<>();
 		for (int index = 0; index < representations.size(); index++) {
 			Map<String, Object> map = representations.get(index);
@@ -79,15 +83,27 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 				monoList.add(resolveException(ex, environment, null, index));
 				continue;
 			}
-			monoList.add(invokeResolver(environment, handlerMethod, map, index));
+
+			if (!handlerMethod.isBatchHandlerMethod()) {
+				monoList.add(invokeEntityMethod(environment, handlerMethod, map, index));
+			}
+			else if (batched.contains(typename)) {
+				// zip needs a value, this will be replaced by batch results
+				monoList.add(Mono.just(Collections.emptyMap()));
+			}
+			else {
+				EntityBatchDelegate batchDelegate = new EntityBatchDelegate(environment, handlerMethod, typename);
+				monoList.add(batchDelegate.invokeEntityBatchMethod());
+				batched.add(typename);
+			}
 		}
 		return Mono.zip(monoList, Arrays::asList).map(EntitiesDataFetcher::toDataFetcherResult);
 	}
 
-	private Mono<Object> invokeResolver(
+	private Mono<Object> invokeEntityMethod(
 			DataFetchingEnvironment env, EntityHandlerMethod handlerMethod, Map<String, Object> map, int index) {
 
-		return handlerMethod.getEntity(env, map, index)
+		return handlerMethod.getEntity(env, map)
 				.switchIfEmpty(Mono.error(new RepresentationNotResolvedException(map, handlerMethod)))
 				.onErrorResume((ex) -> resolveException(ex, env, handlerMethod, index));
 	}
@@ -96,7 +112,7 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 			Throwable ex, DataFetchingEnvironment env, @Nullable EntityHandlerMethod handlerMethod, int index) {
 
 		Throwable theEx = (ex instanceof CompletionException) ? ex.getCause() : ex;
-		DataFetchingEnvironment theEnv = new EntityDataFetchingEnvironment(env, index);
+		DataFetchingEnvironment theEnv = new IndexedDataFetchingEnvironment(env, index);
 		Object handler = (handlerMethod != null) ? handlerMethod.getBean() : null;
 
 		return this.exceptionResolver.resolveException(theEx, theEnv, handler)
@@ -120,6 +136,9 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 		List<GraphQLError> errors = new ArrayList<>();
 		for (int i = 0; i < entities.size(); i++) {
 			Object entity = entities.get(i);
+			if (entity instanceof EntityBatchDelegate delegate) {
+				delegate.processResults(entities, errors);
+			}
 			if (entity instanceof ErrorContainer errorContainer) {
 				errors.addAll(errorContainer.errors());
 				entities.set(i, null);
@@ -129,11 +148,80 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 	}
 
 
-	private static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
+	private class EntityBatchDelegate {
+
+		private final DataFetchingEnvironment environment;
+
+		private final EntityHandlerMethod handlerMethod;
+
+		private final List<Map<String, Object>> representations = new ArrayList<>();
+
+		private final List<Integer> indexes = new ArrayList<>();
+
+		@Nullable
+		private List<?> resultList;
+
+		EntityBatchDelegate(DataFetchingEnvironment env, EntityHandlerMethod handlerMethod, String typeName) {
+			this.environment = env;
+			this.handlerMethod = handlerMethod;
+			List<Map<String, Object>> maps = env.getArgument(_Entity.argumentName);
+			for (int i = 0; i < maps.size(); i++) {
+				Map<String, Object> map = maps.get(i);
+				if (typeName.equals(map.get("__typename"))) {
+					this.representations.add(map);
+					this.indexes.add(i);
+				}
+			}
+		}
+
+		Mono<Object> invokeEntityBatchMethod() {
+			return this.handlerMethod.getEntities(this.environment, this.representations)
+					.mapNotNull((result) -> (((List<?>) result).isEmpty()) ? null : result)
+					.switchIfEmpty(Mono.defer(this::handleEmptyResult))
+					.onErrorResume(this::handleErrorResult)
+					.map((result) -> {
+						this.resultList = (List<?>) result;
+						return this;
+					});
+		}
+
+		Mono<Object> handleEmptyResult() {
+			List<Mono<Object>> exceptions = new ArrayList<>(this.indexes.size());
+			for (int i = 0; i < this.indexes.size(); i++) {
+				Map<String, Object> map = this.representations.get(i);
+				Exception ex = new RepresentationNotResolvedException(map, this.handlerMethod);
+				exceptions.add(resolveException(ex, this.environment, this.handlerMethod, this.indexes.get(i)));
+			}
+			return Mono.zip(exceptions, Arrays::asList);
+		}
+
+		Mono<List<Object>> handleErrorResult(Throwable ex) {
+			List<Mono<Object>> list = new ArrayList<>();
+			for (Integer index : this.indexes) {
+				list.add(resolveException(ex, this.environment, this.handlerMethod, index));
+			}
+			return Mono.zip(list, Arrays::asList);
+		}
+
+		void processResults(List<Object> entities, List<GraphQLError> errors) {
+			Assert.state(this.resultList != null, "Expected resultList");
+			for (int i = 0; i < this.resultList.size(); i++) {
+				Object entity = this.resultList.get(i);
+				if (entity instanceof ErrorContainer errorContainer) {
+					errors.addAll(errorContainer.errors());
+					entity = null;
+				}
+				entities.set(this.indexes.get(i), entity);
+			}
+		}
+	}
+
+
+	private static class IndexedDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
 
 		private final ExecutionStepInfo executionStepInfo;
 
-		EntityDataFetchingEnvironment(DataFetchingEnvironment env, int index) {
+		IndexedDataFetchingEnvironment(DataFetchingEnvironment env, int index) {
 			super(env);
 			this.executionStepInfo = ExecutionStepInfo.newExecutionStepInfo(env.getExecutionStepInfo())
 					.path(env.getExecutionStepInfo().getPath().segment(index))

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java

@@ -16,6 +16,8 @@
 
 package org.springframework.graphql.data.federation;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 import graphql.schema.DataFetchingEnvironment;
@@ -25,6 +27,7 @@ import org.springframework.core.ResolvableType;
 import org.springframework.graphql.data.GraphQlArgumentBinder;
 import org.springframework.graphql.data.method.annotation.Argument;
 import org.springframework.graphql.data.method.annotation.support.ArgumentMethodArgumentResolver;
+import org.springframework.lang.Nullable;
 import org.springframework.validation.BindException;
 
 /**
@@ -48,24 +51,52 @@ final class EntityArgumentMethodArgumentResolver extends ArgumentMethodArgumentR
 			DataFetchingEnvironment environment, String name, ResolvableType targetType) throws BindException {
 
 		if (environment instanceof EntityDataFetchingEnvironment entityEnv) {
-			Map<String, Object> entityMap = entityEnv.getRepresentation();
-			Object rawValue = entityMap.get(name);
-			boolean isOmitted = !entityMap.containsKey(name);
-			return getArgumentBinder().bind(name, rawValue, isOmitted, targetType);
+			return doBind(name, targetType, entityEnv.getRepresentation());
+		}
+		else if (environment instanceof EntityBatchDataFetchingEnvironment batchEnv) {
+			name = dePluralize(name);
+			targetType = targetType.getNested(2);
+			List<Object> values = new ArrayList<>();
+			for (Map<String, Object> representation : batchEnv.getRepresentations()) {
+				values.add(doBind(name, targetType, representation));
+			}
+			return values;
+		}
+		else {
+			throw new IllegalStateException("Expected decorated DataFetchingEnvironment");
 		}
-
-		throw new IllegalStateException("Expected decorated DataFetchingEnvironment");
 	}
 
+	@Nullable
+	private Object doBind(String name, ResolvableType targetType, Map<String, Object> entityMap) throws BindException {
+		Object rawValue = entityMap.get(name);
+		boolean isOmitted = !entityMap.containsKey(name);
+		return getArgumentBinder().bind(rawValue, isOmitted, targetType);
+	}
+
+	private static String dePluralize(String name) {
+		return (name.endsWith("List")) ? name.substring(0, name.length() - 4) : name;
+	}
+
+
 	/**
-	 * Wrap the environment in order to also expose the entity representation map.
+	 * Utility method for use from {@link EntityHandlerMethod} to make the entity
+	 * representation map available.
 	 */
 	static DataFetchingEnvironment wrap(DataFetchingEnvironment env, Map<String, Object> representation) {
 		return new EntityDataFetchingEnvironment(env, representation);
 	}
 
+	/**
+	 * Utility method for use from {@link EntityHandlerMethod} to make the list
+	 * of entity representation maps available.
+	 */
+	static DataFetchingEnvironment wrap(DataFetchingEnvironment env, List<Map<String, Object>> representations) {
+		return new EntityBatchDataFetchingEnvironment(env, representations);
+	}
 
-	private static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
+
+	static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
 
 		private final Map<String, Object> representation;
 
@@ -79,4 +110,19 @@ final class EntityArgumentMethodArgumentResolver extends ArgumentMethodArgumentR
 		}
 	}
 
+
+	static class EntityBatchDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
+
+		private final List<Map<String, Object>> representations;
+
+		EntityBatchDataFetchingEnvironment(DataFetchingEnvironment env, List<Map<String, Object>> representations) {
+			super(env);
+			this.representations = representations;
+		}
+
+		List<Map<String, Object>> getRepresentations() {
+			return this.representations;
+		}
+	}
+
 }

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentsMethodArgumentResolver.java

@@ -0,0 +1,74 @@
+/*
+ * 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.
+ * 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.federation;
+
+import java.util.List;
+import java.util.Map;
+
+import graphql.schema.DataFetchingEnvironment;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.core.ResolvableType;
+import org.springframework.graphql.data.GraphQlArgumentBinder;
+import org.springframework.graphql.data.federation.EntityArgumentMethodArgumentResolver.EntityBatchDataFetchingEnvironment;
+import org.springframework.graphql.data.federation.EntityArgumentMethodArgumentResolver.EntityDataFetchingEnvironment;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
+import org.springframework.util.Assert;
+
+/**
+ * Resolver for the representation map of an entity, or for all representations
+ * for the target schema type (batched handler methods).
+ *
+ * @author Rossen Stoyanchev
+ */
+final class EntityArgumentsMethodArgumentResolver implements HandlerMethodArgumentResolver {
+
+	private final GraphQlArgumentBinder argumentBinder;
+
+
+	EntityArgumentsMethodArgumentResolver(GraphQlArgumentBinder argumentBinder) {
+		Assert.notNull(argumentBinder, "GraphQlArgumentBinder is required");
+		this.argumentBinder = argumentBinder;
+	}
+
+
+	@Override
+	public boolean supportsParameter(MethodParameter param) {
+		if (param.getParameterType().equals(List.class)) {
+			param = param.nested(0);
+		}
+		if (param.getNestedParameterType().equals(Map.class)) {
+			return param.nested(0).getNestedParameterType().equals(String.class);
+		}
+		return false;
+	}
+
+	@Override
+	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment env) throws Exception {
+		ResolvableType targetType = ResolvableType.forMethodParameter(parameter);
+		if (env instanceof EntityDataFetchingEnvironment entityEnv) {
+			return this.argumentBinder.bind(entityEnv.getRepresentation(), false, targetType);
+		}
+		else if (env instanceof EntityBatchDataFetchingEnvironment batchEnv) {
+			return this.argumentBinder.bind(batchEnv.getRepresentations(), false, targetType);
+		}
+		else {
+			throw new IllegalStateException("Expected decorated DataFetchingEnvironment");
+		}
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java

@@ -16,6 +16,7 @@
 
 package org.springframework.graphql.data.federation;
 
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
@@ -23,7 +24,6 @@ import java.util.concurrent.Executor;
 import graphql.schema.DataFetchingEnvironment;
 import reactor.core.publisher.Mono;
 
-import org.springframework.graphql.data.method.HandlerMethod;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.data.method.annotation.support.DataFetcherHandlerMethodSupport;
 import org.springframework.lang.Nullable;
@@ -35,27 +35,53 @@ import org.springframework.lang.Nullable;
  */
 final class EntityHandlerMethod extends DataFetcherHandlerMethodSupport {
 
+	private final boolean batchHandlerMethod;
+
+
 	EntityHandlerMethod(
-			HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers,
+			FederationSchemaFactory.EntityMappingInfo info, HandlerMethodArgumentResolverComposite resolvers,
 			@Nullable Executor executor) {
 
-		super(handlerMethod, resolvers, executor);
+		super(info.handlerMethod(), resolvers, executor);
+		this.batchHandlerMethod = info.isBatchHandlerMethod();
 	}
 
 
-	Mono<Object> getEntity(
-			DataFetchingEnvironment environment, Map<String, Object> representation, int index) {
+	boolean isBatchHandlerMethod() {
+		return this.batchHandlerMethod;
+	}
 
+
+	Mono<Object> getEntity(DataFetchingEnvironment env, Map<String, Object> representation) {
 		Object[] args;
 		try {
-			environment = EntityArgumentMethodArgumentResolver.wrap(environment, representation);
-			args = getMethodArgumentValues(environment, representation);
+			env = EntityArgumentMethodArgumentResolver.wrap(env, representation);
+			args = getMethodArgumentValues(env);
 		}
 		catch (Throwable ex) {
 			return Mono.error(ex);
 		}
 
-		Object result = doInvoke(environment.getGraphQlContext(), args);
+		return doInvoke(env, args);
+	}
+
+	@SuppressWarnings("unchecked")
+	Mono<Object> getEntities(DataFetchingEnvironment env, List<Map<String, Object>> representations) {
+		Object[] args;
+		try {
+			env = EntityArgumentMethodArgumentResolver.wrap(env, representations);
+			args = getMethodArgumentValues(env);
+		}
+		catch (Throwable ex) {
+			return Mono.error(ex);
+		}
+
+		return doInvoke(env, args);
+	}
+
+	private Mono<Object> doInvoke(DataFetchingEnvironment env, Object[] args) {
+
+		Object result = doInvoke(env.getGraphQlContext(), args);
 
 		if (result instanceof Mono<?> mono) {
 			return mono.cast(Object.class);

spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java

@@ -18,7 +18,9 @@ package org.springframework.graphql.data.federation;
 
 import java.lang.reflect.Method;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletionStage;
 import java.util.function.BiFunction;
 
 import com.apollographql.federation.graphqljava.Federation;
@@ -28,10 +30,12 @@ import graphql.schema.GraphQLSchema;
 import graphql.schema.TypeResolver;
 import graphql.schema.idl.RuntimeWiring;
 import graphql.schema.idl.TypeDefinitionRegistry;
+import reactor.core.publisher.Mono;
 
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.expression.BeanFactoryResolver;
 import org.springframework.core.KotlinDetector;
+import org.springframework.core.MethodParameter;
 import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.graphql.data.GraphQlArgumentBinder;
 import org.springframework.graphql.data.method.HandlerMethod;
@@ -89,7 +93,7 @@ public final class FederationSchemaFactory
 
 		detectHandlerMethods().forEach((info) ->
 				this.handlerMethods.put(info.typeName(),
-						new EntityHandlerMethod(info.handlerMethod(), getArgumentResolvers(), getExecutor())));
+						new EntityHandlerMethod(info, getArgumentResolvers(), getExecutor())));
 
 		if (this.typeResolver == null) {
 			this.typeResolver = new ClassNameTypeResolver();
@@ -108,6 +112,7 @@ public final class FederationSchemaFactory
 		resolvers.addResolver(new ContextValueMethodArgumentResolver());
 		resolvers.addResolver(new LocalContextValueMethodArgumentResolver());
 		resolvers.addResolver(new EntityArgumentMethodArgumentResolver(argumentBinder));
+		resolvers.addResolver(new EntityArgumentsMethodArgumentResolver(argumentBinder));
 
 		// Type based
 		resolvers.addResolver(new DataFetchingEnvironmentMethodArgumentResolver());
@@ -172,6 +177,15 @@ public final class FederationSchemaFactory
 
 
 	public record EntityMappingInfo(String typeName, HandlerMethod handlerMethod) {
+
+		public boolean isBatchHandlerMethod() {
+			MethodParameter type = handlerMethod().getReturnType();
+			Class<?> paramType = type.getParameterType();
+			if (Mono.class.isAssignableFrom(paramType) || CompletionStage.class.isAssignableFrom(paramType)) {
+				type = type.nested();
+			}
+			return List.class.isAssignableFrom(type.getParameterType());
+		}
 	}
 
 }

spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java

@@ -16,6 +16,7 @@
 
 package org.springframework.graphql.data.federation;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -82,18 +83,10 @@ public class EntityMappingInvocationTests {
 						Map.of("__typename", "Book", "id", "3"),
 						Map.of("__typename", "Book", "id", "5")));
 
-		ExecutionGraphQlRequest request = TestExecutionRequest.forDocumentAndVars(document, variables);
-		Mono<ExecutionGraphQlResponse> responseMono = graphQlService().execute(request);
+		ResponseHelper helper = executeWith(BookController.class, variables);
 
-		ResponseHelper helper = ResponseHelper.forResponse(responseMono);
-
-		Author author = helper.toEntity("_entities[0].author", Author.class);
-		assertThat(author.getFirstName()).isEqualTo("Joseph");
-		assertThat(author.getLastName()).isEqualTo("Heller");
-
-		author = helper.toEntity("_entities[1].author", Author.class);
-		assertThat(author.getFirstName()).isEqualTo("George");
-		assertThat(author.getLastName()).isEqualTo("Orwell");
+		assertAuthor(0, "Joseph", "Heller", helper);
+		assertAuthor(1, "George", "Orwell", helper);
 	}
 
 	@Test
@@ -108,21 +101,77 @@ public class EntityMappingInvocationTests {
 						Map.of("__typename", "Book", "id", "3"),
 						Map.of("__typename", "Book", "id", "5")));
 
+		ResponseHelper helper = executeWith(BookController.class, variables);
+
+		assertError(helper, 0, "BAD_REQUEST", "Missing \"__typename\" argument");
+		assertError(helper, 1, "INTERNAL_ERROR", "No entity fetcher");
+		assertError(helper, 2, "BAD_REQUEST", "handled");
+		assertError(helper, 3, "INTERNAL_ERROR", "not handled");
+		assertError(helper, 4, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
+
+		assertAuthor(5, "Joseph", "Heller", helper);
+		assertAuthor(6, "George", "Orwell", helper);
+	}
+
+	@Test
+	void batching() {
+		Map<String, Object> variables =
+				Map.of("representations", List.of(
+						Map.of("__typename", "Book", "id", "1"),
+						Map.of("__typename", "Book", "id", "4"),
+						Map.of("__typename", "Book", "id", "5"),
+						Map.of("__typename", "Book", "id", "42"),
+						Map.of("__typename", "Book", "id", "53")));
+
+		ResponseHelper helper = executeWith(BookBatchController.class, variables);
+
+		assertAuthor(0, "George", "Orwell", helper);
+		assertAuthor(1, "Virginia", "Woolf", helper);
+		assertAuthor(2, "George", "Orwell", helper);
+		assertAuthor(3, "Douglas", "Adams", helper);
+		assertAuthor(4, "Vince", "Gilligan", helper);
+	}
+
+	@Test
+	void batchingWithError() {
+		Map<String, Object> variables =
+				Map.of("representations", List.of(
+						Map.of("__typename", "Book", "id", "-97"),
+						Map.of("__typename", "Book", "id", "4"),
+						Map.of("__typename", "Book", "id", "5")));
+
+		ResponseHelper helper = executeWith(BookBatchController.class, variables);
+
+		assertError(helper, 0, "BAD_REQUEST", "handled");
+		assertError(helper, 1, "BAD_REQUEST", "handled");
+		assertError(helper, 2, "BAD_REQUEST", "handled");
+	}
+
+	@Test
+	void batchingWithoutResult() {
+		Map<String, Object> variables =
+				Map.of("representations", List.of(
+						Map.of("__typename", "Book", "id", "-99"),
+						Map.of("__typename", "Book", "id", "4"),
+						Map.of("__typename", "Book", "id", "5")));
+
+		ResponseHelper helper = executeWith(BookBatchController.class, variables);
+
+		assertError(helper, 0, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
+		assertError(helper, 1, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
+		assertError(helper, 2, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
+	}
+
+	private static ResponseHelper executeWith(Class<?> controllerClass, Map<String, Object> variables) {
 		ExecutionGraphQlRequest request = TestExecutionRequest.forDocumentAndVars(document, variables);
-		Mono<ExecutionGraphQlResponse> responseMono = graphQlService().execute(request);
+		Mono<ExecutionGraphQlResponse> responseMono = graphQlService(controllerClass).execute(request);
+		return ResponseHelper.forResponse(responseMono);
+	}
 
-		ResponseHelper helper = ResponseHelper.forResponse(responseMono);
-
-		int i = 0;
-
-		assertError(helper, i++, "BAD_REQUEST", "Missing \"__typename\" argument");
-		assertError(helper, i++, "INTERNAL_ERROR", "No entity fetcher");
-		assertError(helper, i++, "BAD_REQUEST", "handled");
-		assertError(helper, i++, "INTERNAL_ERROR", "not handled");
-		assertError(helper, i++, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
-
-		assertThat(helper.toEntity("_entities[" + i++ + "].author", Author.class).getLastName()).isEqualTo("Heller");
-		assertThat(helper.toEntity("_entities[" + i++ + "].author", Author.class).getLastName()).isEqualTo("Orwell");
+	private static void assertAuthor(int index, String firstName, String lastName, ResponseHelper helper) {
+		Author author = helper.toEntity("_entities[" + index + "].author", Author.class);
+		assertThat(author.getFirstName()).isEqualTo(firstName);
+		assertThat(author.getLastName()).isEqualTo(lastName);
 	}
 
 	private static void assertError(ResponseHelper helper, int i, String errorType, String msg) {
@@ -133,11 +182,11 @@ public class EntityMappingInvocationTests {
 		assertThat(helper.<Object>rawValue(path)).isNull();
 	}
 
-	private TestExecutionGraphQlService graphQlService() {
+	private static TestExecutionGraphQlService graphQlService(Class<?> controllerClass) {
 		BatchLoaderRegistry registry = new DefaultBatchLoaderRegistry();
 
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
-		context.register(AuthorController.class);
+		context.register(controllerClass);
 		context.registerBean(BatchLoaderRegistry.class, () -> registry);
 		context.refresh();
 
@@ -159,7 +208,7 @@ public class EntityMappingInvocationTests {
 
 	@SuppressWarnings("unused")
 	@Controller
-	private static class AuthorController {
+	private static class BookController {
 
 		@Nullable
 		@EntityMapping
@@ -191,4 +240,43 @@ public class EntityMappingInvocationTests {
 		}
 	}
 
+	@SuppressWarnings("unused")
+	@Controller
+	private static class BookBatchController {
+
+		@EntityMapping
+		public List<Book> book(@Argument List<Integer> idList, List<Map<String, Object>> representations) {
+
+			if (idList.get(0) == -97) {
+				throw new IllegalArgumentException("handled");
+			}
+
+			if (idList.get(0) == -99) {
+				return Collections.emptyList();
+			}
+
+			assertThat(representations).hasSize(5).containsExactly(
+					Map.of("__typename", "Book", "id", "1"),
+					Map.of("__typename", "Book", "id", "4"),
+					Map.of("__typename", "Book", "id", "5"),
+					Map.of("__typename", "Book", "id", "42"),
+					Map.of("__typename", "Book", "id", "53"));
+
+			return idList.stream().map(id -> new Book((long) id, null, (Long) null)).toList();
+		}
+
+		@BatchMapping
+		public Flux<Author> author(List<Book> books) {
+			return Flux.fromIterable(books).map(book -> BookSource.getBook(book.getId()).getAuthor());
+		}
+
+		@GraphQlExceptionHandler
+		public GraphQLError handle(IllegalArgumentException ex, DataFetchingEnvironment env) {
+			return GraphqlErrorBuilder.newError(env)
+					.errorType(ErrorType.BAD_REQUEST)
+					.message(ex.getMessage())
+					.build();
+		}
+	}
+
 }