Support batched EntityMapping methods

See gh-922

Author: rstoyanchev

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

@@ -142,33 +142,26 @@ public Object bind(
 		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}
 	 * @param targetType the type of Object to create
 	 * @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 @@
 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 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 @@ public Mono<DataFetcherResult<List<Object>>> get(DataFetchingEnvironment environ
 				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 @@ private Mono<Object> resolveException(
 			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 @@ private static DataFetcherResult<List<Object>> toDataFetcherResult(List<Object>
 		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 @@ private static DataFetcherResult<List<Object>> toDataFetcherResult(List<Object>
 	}
 
 
-	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.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 @@ protected Object doBind(
 			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");
+		}
+	}
+
+	@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);
+	}
 
-		throw new IllegalStateException("Expected decorated DataFetchingEnvironment");
+	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 @@ Map<String, Object> getRepresentation() {
 		}
 	}
 
+
+	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");
+		}
+	}
+
+}