Add nested key entity support for GraphQL Federation

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

@@ -57,14 +57,17 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 
 	private final HandlerDataFetcherExceptionResolver exceptionResolver;
 
+	private final EntityKeyResolver entityKeyResolver;
+
 
 	EntitiesDataFetcher(
 			Map<String, EntityHandlerMethod> handlerMethods, Map<String, String> objectToInterfaceMap,
-			HandlerDataFetcherExceptionResolver resolver) {
+			HandlerDataFetcherExceptionResolver resolver, EntityKeyResolver entityKeyResolver) {
 
 		this.handlerMethods = new LinkedHashMap<>(handlerMethods);
 		this.objectToInterfaceMap = objectToInterfaceMap;
 		this.exceptionResolver = resolver;
+		this.entityKeyResolver = entityKeyResolver;
 	}
 
 
@@ -120,7 +123,7 @@ private Mono<Object> invokeEntityMethod(
 			DataFetchingEnvironment environment, EntityHandlerMethod handlerMethod,
 			Map<String, Object> representation, int index) {
 
-		return handlerMethod.getEntity(environment, representation)
+		return handlerMethod.getEntity(environment, representation, this.entityKeyResolver)
 				.switchIfEmpty(Mono.error(new RepresentationNotResolvedException(representation, handlerMethod)))
 				.onErrorResume((ex) -> resolveException(ex, environment, handlerMethod, index));
 	}
@@ -141,7 +144,7 @@ private Mono<EntitiesResultContainer> invokeEntitiesMethod(
 			}
 		}
 
-		return handlerMethod.getEntities(environment, typeRepresentations)
+		return handlerMethod.getEntities(environment, typeRepresentations, this.entityKeyResolver)
 				.mapNotNull((result) -> (((List<?>) result).isEmpty()) ? null : result)
 				.switchIfEmpty(Mono.defer(() -> {
 					List<Mono<?>> exceptions = new ArrayList<>(originalIndexes.size());

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

@@ -51,14 +51,16 @@ final class EntityArgumentMethodArgumentResolver extends ArgumentMethodArgumentR
 			DataFetchingEnvironment environment, String name, ResolvableType targetType) throws BindException {
 
 		if (environment instanceof EntityDataFetchingEnvironment entityEnv) {
-			return doBind(name, targetType, entityEnv.getRepresentation());
+			Object key = entityEnv.getKey();
+			return doBind(name, targetType, key);
 		}
 		else if (environment instanceof EntityBatchDataFetchingEnvironment batchEnv) {
 			name = dePluralize(name);
 			targetType = targetType.getNested(2);
 			List<@Nullable Object> values = new ArrayList<>();
 			for (Map<String, Object> representation : batchEnv.getRepresentations()) {
-				values.add(doBind(name, targetType, representation));
+				Object key = batchEnv.getKey(representation);
+				values.add(doBind(name, targetType, key));
 			}
 			return values;
 		}
@@ -67,10 +69,8 @@ else if (environment instanceof EntityBatchDataFetchingEnvironment batchEnv) {
 		}
 	}
 
-	private @Nullable 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 @Nullable Object doBind(String name, ResolvableType targetType, Object key) throws BindException {
+		return getArgumentBinder().bind(key, false, targetType);
 	}
 
 	private static String dePluralize(String name) {
@@ -82,16 +82,16 @@ private static String dePluralize(String name) {
 	 * 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);
+	static DataFetchingEnvironment wrap(DataFetchingEnvironment env, Map<String, Object> representation, EntityKeyResolver resolver) {
+		return new EntityDataFetchingEnvironment(env, representation, resolver);
 	}
 
 	/**
 	 * 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);
+	static DataFetchingEnvironment wrap(DataFetchingEnvironment env, List<Map<String, Object>> representations, EntityKeyResolver resolver) {
+		return new EntityBatchDataFetchingEnvironment(env, representations, resolver);
 	}
 
 
@@ -101,15 +101,21 @@ static DataFetchingEnvironment wrap(DataFetchingEnvironment env, List<Map<String
 	static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
 
 		private final Map<String, Object> representation;
+		private final EntityKeyResolver resolver;
 
-		EntityDataFetchingEnvironment(DataFetchingEnvironment env, Map<String, Object> representation) {
+		EntityDataFetchingEnvironment(DataFetchingEnvironment env, Map<String, Object> representation, EntityKeyResolver resolver) {
 			super(env);
 			this.representation = representation;
+			this.resolver = resolver;
 		}
 
 		Map<String, Object> getRepresentation() {
 			return this.representation;
 		}
+
+		public Object getKey() {
+			return this.resolver.getKey(representation);
+		}
 	}
 
 
@@ -120,15 +126,21 @@ Map<String, Object> getRepresentation() {
 	static class EntityBatchDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
 
 		private final List<Map<String, Object>> representations;
+		private final EntityKeyResolver resolver;
 
-		EntityBatchDataFetchingEnvironment(DataFetchingEnvironment env, List<Map<String, Object>> representations) {
+		EntityBatchDataFetchingEnvironment(DataFetchingEnvironment env, List<Map<String, Object>> representations, EntityKeyResolver resolver) {
 			super(env);
 			this.representations = representations;
+			this.resolver = resolver;
 		}
 
 		List<Map<String, Object>> getRepresentations() {
 			return this.representations;
 		}
+
+		public Object getKey(Map<String, Object> representation) {
+			return this.resolver.getKey(representation);
+		}
 	}
 
 }

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

@@ -52,13 +52,13 @@ boolean isBatchHandlerMethod() {
 	}
 
 
-	Mono<Object> getEntity(DataFetchingEnvironment env, Map<String, Object> representation) {
-		env = EntityArgumentMethodArgumentResolver.wrap(env, representation);
+	Mono<Object> getEntity(DataFetchingEnvironment env, Map<String, Object> representation, EntityKeyResolver resolver) {
+		env = EntityArgumentMethodArgumentResolver.wrap(env, representation, resolver);
 		return doInvoke(env);
 	}
 
-	Mono<Object> getEntities(DataFetchingEnvironment env, List<Map<String, Object>> representations) {
-		env = EntityArgumentMethodArgumentResolver.wrap(env, representations);
+	Mono<Object> getEntities(DataFetchingEnvironment env, List<Map<String, Object>> representations, EntityKeyResolver resolver) {
+		env = EntityArgumentMethodArgumentResolver.wrap(env, representations, resolver);
 		return doInvoke(env);
 	}
 

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

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2002-present 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.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+
+import graphql.language.Argument;
+import graphql.language.Directive;
+import graphql.language.Document;
+import graphql.language.Field;
+import graphql.language.OperationDefinition;
+import graphql.language.SelectionSet;
+import graphql.language.StringValue;
+import graphql.language.TypeDefinition;
+import graphql.parser.Parser;
+import graphql.schema.idl.TypeDefinitionRegistry;
+
+/**
+ * Resolves a simple entity key from a federated entity representation map.
+ * <p>This resolver inspects the GraphQL type definitions for {@code @key(fields: "...")} directives and records keys
+ * only when the directive declares exactly one top-level scalar field (for example {@code @key(fields: "id")}).
+ *
+ * <p>For matching types, {@link #getKey(Map)} returns the scalar key value from the representation.
+ * For all other types (composite keys, nested keys), it returns the full representation map.
+ *
+ * @author James Bodkin
+ * @since 2.0.5
+ */
+public class EntityKeyResolver {
+
+	private final Map<String, String> keyField = new LinkedHashMap<>();
+
+	EntityKeyResolver(TypeDefinitionRegistry registry) {
+		for (TypeDefinition<?> type : registry.types().values()) {
+			for (Directive directive : type.getDirectives("key")) {
+				processDirective(type, directive);
+			}
+		}
+	}
+
+	private void processDirective(TypeDefinition<?> type, Directive directive) {
+		Argument argument = directive.getArgument("fields");
+		if (argument != null) {
+			Object value = argument.getValue();
+			if (value instanceof StringValue sv) {
+				Parser parser = new Parser();
+				Document document = parser.parseDocument("{" + sv.getValue() + "}");
+				OperationDefinition operationDefinition = document.getDefinitionsOfType(OperationDefinition.class).get(0);
+				SelectionSet selectionSet = operationDefinition.getSelectionSet();
+				if (selectionSet.getSelections().size() > 1) {
+					return;
+				}
+
+				Field field = selectionSet.getSelectionsOfType(Field.class).get(0);
+				if (field.getSelectionSet() == null || field.getSelectionSet().getSelections().isEmpty()) {
+					this.keyField.put(type.getName(), field.getName());
+				}
+			}
+		}
+	}
+
+	/**
+	 * Resolve the entity key for a federated entity representation map.
+	 *
+	 * <p>If the representation type has a registered simple key field, this method returns that field value.
+	 * Otherwise, it returns the full representation as-is.
+	 * @param representation the federated entity representation map, expected to contain {@code __typename}
+	 * @return scalar key value for simple key entities, or the full representation map for complex key entities
+	 */
+	@SuppressWarnings("NullAway")
+	public Object getKey(Map<String, Object> representation) {
+		String typeName = (String) Objects.requireNonNull(representation.get("__typename"));
+		if (this.keyField.containsKey(typeName)) {
+			String fieldName = this.keyField.get(typeName);
+			return representation.get(fieldName);
+		}
+		else {
+			return representation;
+		}
+	}
+
+}

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

@@ -188,8 +188,10 @@ public SchemaTransformer createSchemaTransformer(TypeDefinitionRegistry registry
 		Map<String, String> objectToInterfaceTypeMap = detectInterfaceImplementationTypes(registry);
 		checkEntityMappings(registry, objectToInterfaceTypeMap);
 
-		EntitiesDataFetcher entitiesDataFetcher =
-				new EntitiesDataFetcher(this.handlerMethods, objectToInterfaceTypeMap, getExceptionResolver());
+		EntityKeyResolver entityKeyResolver = new EntityKeyResolver(registry);
+
+		EntitiesDataFetcher entitiesDataFetcher = new EntitiesDataFetcher(this.handlerMethods, objectToInterfaceTypeMap,
+				getExceptionResolver(), entityKeyResolver);
 
 		return Federation.transform(registry, wiring)
 				.fetchEntities(entitiesDataFetcher)

spring-graphql/src/test/java/org/springframework/graphql/Library.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2020-present 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;
+
+public record Library(String id, Location location) {
+}

spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2020-present 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;
+
+public record LibraryId(String id, LocationId location) {
+
+	public record LocationId(String id) {
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/Location.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2020-present 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;
+
+public record Location(String id) {
+}

spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2020-present 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;
+
+public record LocationArea(Location location) {
+}

spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2020-present 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;
+
+public record LocationAreaId(LocationId location) {
+
+	public record LocationId(String id) {
+	}
+
+}