Add nested key entity support for GraphQL Federation

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

@@ -24,6 +24,7 @@
 import graphql.schema.DelegatingDataFetchingEnvironment;
 import org.jspecify.annotations.Nullable;
 
+import org.springframework.beans.BeanUtils;
 import org.springframework.core.ResolvableType;
 import org.springframework.graphql.data.GraphQlArgumentBinder;
 import org.springframework.graphql.data.method.annotation.Argument;
@@ -68,9 +69,25 @@ 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);
+		if (isScalarValue(entityMap)) {
+			Object rawValue = entityMap.get(name);
+			return getArgumentBinder().bind(rawValue, false, targetType);
+		}
+		return getArgumentBinder().bind(entityMap, false, targetType);
+	}
+
+	private boolean isScalarValue(Map<String, Object> entityMap) {
+		if (entityMap.size() != 2) {
+			return false;
+		}
+
+		for (Map.Entry<String, Object> entry : entityMap.entrySet()) {
+			if (!"__typename".equals(entry.getKey())) {
+				Object value = entry.getValue();
+				return BeanUtils.isSimpleValueType(value.getClass());
+			}
+		}
+		return false;
 	}
 
 	private static String dePluralize(String name) {

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

@@ -0,0 +1,4 @@
+package org.springframework.graphql;
+
+public record Library(String id, Location location) {
+}

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

@@ -0,0 +1,8 @@
+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,4 @@
+package org.springframework.graphql;
+
+public record Location(String id) {
+}

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

@@ -0,0 +1,4 @@
+package org.springframework.graphql;
+
+public record LocationArea(Location location) {
+}

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

@@ -0,0 +1,8 @@
+package org.springframework.graphql;
+
+public record LocationAreaId(LocationId location) {
+
+    public record LocationId(String id) {
+    }
+
+}

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

@@ -42,6 +42,11 @@
 import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.ExecutionGraphQlResponse;
 import org.springframework.graphql.GraphQlSetup;
+import org.springframework.graphql.Library;
+import org.springframework.graphql.LibraryId;
+import org.springframework.graphql.Location;
+import org.springframework.graphql.LocationArea;
+import org.springframework.graphql.LocationAreaId;
 import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.TestExecutionGraphQlService;
 import org.springframework.graphql.TestExecutionRequest;
@@ -64,9 +69,9 @@
  */
 class EntityMappingInvocationTests {
 
-	private static final Resource federationSchema = new ClassPathResource("books/federation-schema.graphqls");
+	private static final Resource bookFederationSchema = new ClassPathResource("books/federation-schema.graphqls");
 
-	private static final String document = """
+	private static final String bookDocument = """
 			query Entities($representations: [_Any!]!) {
 				_entities(representations: $representations) {
 					...on Book {
@@ -81,6 +86,32 @@ query Entities($representations: [_Any!]!) {
 			}
 			""";
 
+	private static final Resource libraryFederationSchema = new ClassPathResource("library/federation-schema.graphqls");
+
+	private static final String locationAreaDocument = """
+			query Entities($representations: [_Any!]!) {
+				_entities(representations: $representations) {
+					...on LocationArea {
+						location {
+							id
+						}
+					}
+				}
+			}
+			""";
+
+	private static final String libraryDocument = """
+			query Entities($representations: [_Any!]!) {
+				_entities(representations: $representations) {
+					...on Library {
+						id
+						location {
+							id
+						}
+					}
+				}
+			}
+			""";
 
 	@Test
 	void fetchEntities() {
@@ -90,7 +121,7 @@ void fetchEntities() {
 						Map.of("__typename", "Book", "id", "5"),
 						Map.of("__typename", "PrintedMedia", "id", "42")));
 
-		ResponseHelper helper = executeWith(BookController.class, variables);
+		ResponseHelper helper = executeWith(BookController.class, bookFederationSchema, bookDocument, variables);
 
 		assertAuthor(0, "Joseph", "Heller", helper);
 		assertAuthor(1, "George", "Orwell", helper);
@@ -109,7 +140,7 @@ void fetchEntitiesWithExceptions() {
 						Map.of("__typename", "Book", "id", "3"),
 						Map.of("__typename", "Book", "id", "5")));
 
-		ResponseHelper helper = executeWith(BookController.class, variables);
+		ResponseHelper helper = executeWith(BookController.class, bookFederationSchema, bookDocument, variables);
 
 		assertError(helper, 0, "BAD_REQUEST", "Missing \"__typename\" argument");
 		assertError(helper, 1, "INTERNAL_ERROR", "No entity fetcher");
@@ -124,11 +155,36 @@ void fetchEntitiesWithExceptions() {
 	@Test // gh-1057
 	void fetchEntitiesWithEmptyList() {
 		Map<String, Object> vars = Map.of("representations", Collections.emptyList());
-		ResponseHelper helper = executeWith(BookController.class, vars);
+		ResponseHelper helper = executeWith(BookController.class, bookFederationSchema, bookDocument, vars);
 
 		assertThat(helper.toEntity("_entities.length()", Integer.class)).isEqualTo(0);
 	}
 
+	@Test
+	void fetchSingleNestedKeyEntity() {
+		Map<String, Object> variables = Map.of("representations", List.of(
+			Map.of("__typename", "LocationArea", "location", Map.of("id", "1"))
+		));
+
+		ResponseHelper helper = executeWith(LibraryController.class, libraryFederationSchema, locationAreaDocument, variables);
+
+		LocationArea locationArea = helper.toEntity("_entities[0]", LocationArea.class);
+		assertThat(locationArea.location().id()).isEqualTo("1");
+	}
+
+	@Test
+	void fetchMixedNestedKeyEntity() {
+		Map<String, Object> variables = Map.of("representations", List.of(
+				Map.of("__typename", "Library", "id", "1", "location", Map.of("id", "1"))
+		));
+
+		ResponseHelper helper = executeWith(LibraryController.class, libraryFederationSchema, libraryDocument, variables);
+
+		Library library = helper.toEntity("_entities[0]", Library.class);
+		assertThat(library.id()).isEqualTo("1");
+		assertThat(library.location().id()).isEqualTo("1");
+	}
+
 	@ValueSource(classes = {BookListController.class, BookFluxController.class})
 	@ParameterizedTest
 	void batching(Class<?> controllerClass) {
@@ -140,7 +196,7 @@ void batching(Class<?> controllerClass) {
 						Map.of("__typename", "Book", "id", "42"),
 						Map.of("__typename", "Book", "id", "53")));
 
-		ResponseHelper helper = executeWith(controllerClass, variables);
+		ResponseHelper helper = executeWith(controllerClass, bookFederationSchema, bookDocument, variables);
 
 		assertAuthor(0, "George", "Orwell", helper);
 		assertAuthor(1, "Virginia", "Woolf", helper);
@@ -158,7 +214,7 @@ void batchingWithError(Class<?> controllerClass) {
 						Map.of("__typename", "Book", "id", "4"),
 						Map.of("__typename", "Book", "id", "5")));
 
-		ResponseHelper helper = executeWith(controllerClass, variables);
+		ResponseHelper helper = executeWith(controllerClass, bookFederationSchema, bookDocument, variables);
 
 		assertError(helper, 0, "BAD_REQUEST", "handled");
 		assertError(helper, 1, "BAD_REQUEST", "handled");
@@ -174,7 +230,7 @@ void batchingWithoutResult(Class<?> controllerClass) {
 						Map.of("__typename", "Book", "id", "4"),
 						Map.of("__typename", "Book", "id", "5")));
 
-		ResponseHelper helper = executeWith(controllerClass, variables);
+		ResponseHelper helper = executeWith(controllerClass, bookFederationSchema, bookDocument, 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");
@@ -188,21 +244,23 @@ void dataLoader() {
 						Map.of("__typename", "Book", "id", "3"),
 						Map.of("__typename", "Book", "id", "5")));
 
-		ResponseHelper helper = executeWith(DataLoaderBookController.class, variables);
+		ResponseHelper helper = executeWith(DataLoaderBookController.class, bookFederationSchema, bookDocument, variables);
 
 		assertAuthor(0, "Joseph", "Heller", helper);
 		assertAuthor(1, "George", "Orwell", helper);
 	}
 
 	@Test
 	void unmappedEntity() {
-		assertThatIllegalStateException().isThrownBy(() -> executeWith(EmptyController.class, Map.of()))
+		assertThatIllegalStateException().isThrownBy(() -> executeWith(EmptyController.class, bookFederationSchema, bookDocument, Map.of()))
 				.withMessage("Unmapped entity types: 'Media', 'PrintedMedia', 'Book'");
 	}
 
-	private static ResponseHelper executeWith(Class<?> controllerClass, Map<String, Object> variables) {
+	private static ResponseHelper executeWith(Class<?> controllerClass, Resource federationSchema, String document,
+											  Map<String, Object> variables) {
+
 		ExecutionGraphQlRequest request = TestExecutionRequest.forDocumentAndVars(document, variables);
-		Mono<ExecutionGraphQlResponse> responseMono = graphQlService(controllerClass).execute(request);
+		Mono<ExecutionGraphQlResponse> responseMono = graphQlService(controllerClass, federationSchema).execute(request);
 		return ResponseHelper.forResponse(responseMono);
 	}
 
@@ -220,7 +278,7 @@ private static void assertError(ResponseHelper helper, int i, String errorType,
 		assertThat(helper.<Object>rawValue(path)).isNull();
 	}
 
-	private static TestExecutionGraphQlService graphQlService(Class<?> controllerClass) {
+	private static TestExecutionGraphQlService graphQlService(Class<?> controllerClass, Resource federationSchema) {
 		BatchLoaderRegistry registry = new DefaultBatchLoaderRegistry();
 
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
@@ -391,4 +449,36 @@ public GraphQLError handle(IllegalArgumentException ex, DataFetchingEnvironment
 		}
 	}
 
+	@SuppressWarnings("unused")
+	@Controller
+	public static class LibraryController {
+
+		@EntityMapping
+		public @Nullable Library library(@Argument LibraryId id, Map<String, Object> map) {
+
+			assertThat(id.id()).isNotNull();
+			assertThat(id.location().id()).isNotNull();
+
+			assertThat(map).hasSize(3)
+					.containsEntry("__typename", "Library")
+					.containsEntry("id", "1")
+					.containsEntry("location", Map.of("id", "1"));
+
+			return new Library("1", new Location("1"));
+		}
+
+		@EntityMapping
+		public @Nullable LocationArea locationArea(@Argument LocationAreaId id, Map<String, Object> map) {
+
+			assertThat(id.location().id()).isNotNull();
+
+			assertThat(map).hasSize(2)
+					.containsEntry("__typename", "LocationArea")
+					.containsEntry("location", Map.of("id", "1"));
+
+			return new LocationArea(new Location("1"));
+		}
+
+	}
+
 }

spring-graphql/src/test/resources/library/federation-schema.graphqls

@@ -0,0 +1,14 @@
+extend schema @link(url: "https://specs.apollo.dev/federation/v2.9", import: ["@key", "@extends", "@external"] )
+
+type Location {
+    id: ID!
+}
+
+type Library @key(fields: "id location { id }") {
+    id: ID!
+    location: Location!
+}
+
+type LocationArea @key(fields: "location { id }") {
+    location: Location!
+}