From d0ce0fb1f145dd3920c5db8f6cdbb031f15e3400 Mon Sep 17 00:00:00 2001 From: James Bodkin Date: Mon, 8 Jun 2026 15:06:53 +0100 Subject: [PATCH 1/3] Add nested key entity support for GraphQL Federation Signed-off-by: James Bodkin --- .../EntityArgumentMethodArgumentResolver.java | 23 +++- .../org/springframework/graphql/Library.java | 4 + .../springframework/graphql/LibraryId.java | 8 ++ .../org/springframework/graphql/Location.java | 4 + .../springframework/graphql/LocationArea.java | 4 + .../graphql/LocationAreaId.java | 8 ++ .../EntityMappingInvocationTests.java | 116 ++++++++++++++++-- .../library/federation-schema.graphqls | 14 +++ 8 files changed, 165 insertions(+), 16 deletions(-) create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/Library.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/Location.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java create mode 100644 spring-graphql/src/test/resources/library/federation-schema.graphqls diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java index 8fc9280e..a76b8acc 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java +++ b/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 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 entityMap) { + if (entityMap.size() != 2) { + return false; + } + + for (Map.Entry 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) { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/Library.java b/spring-graphql/src/test/java/org/springframework/graphql/Library.java new file mode 100644 index 00000000..99fd217b --- /dev/null +++ b/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) { +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java b/spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java new file mode 100644 index 00000000..c3327294 --- /dev/null +++ b/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) { + } + +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/Location.java b/spring-graphql/src/test/java/org/springframework/graphql/Location.java new file mode 100644 index 00000000..f2c26d7a --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/Location.java @@ -0,0 +1,4 @@ +package org.springframework.graphql; + +public record Location(String id) { +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java b/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java new file mode 100644 index 00000000..b3bf959c --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java @@ -0,0 +1,4 @@ +package org.springframework.graphql; + +public record LocationArea(Location location) { +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java b/spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java new file mode 100644 index 00000000..0179a6d1 --- /dev/null +++ b/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) { + } + +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java index c025a9f2..e6b64642 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java +++ b/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 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 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 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,7 +244,7 @@ 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); @@ -196,13 +252,15 @@ void dataLoader() { @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 variables) { + private static ResponseHelper executeWith(Class controllerClass, Resource federationSchema, String document, + Map variables) { + ExecutionGraphQlRequest request = TestExecutionRequest.forDocumentAndVars(document, variables); - Mono responseMono = graphQlService(controllerClass).execute(request); + Mono 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.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 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 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")); + } + + } + } diff --git a/spring-graphql/src/test/resources/library/federation-schema.graphqls b/spring-graphql/src/test/resources/library/federation-schema.graphqls new file mode 100644 index 00000000..2c4bb8bd --- /dev/null +++ b/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! +} From c09b84be5f77d36af93b9cc052905ac14d76966f Mon Sep 17 00:00:00 2001 From: James Bodkin Date: Tue, 30 Jun 2026 11:42:38 +0100 Subject: [PATCH 2/3] Formatting Signed-off-by: James Bodkin --- .../org/springframework/graphql/Library.java | 16 +++++++++++++++ .../springframework/graphql/LibraryId.java | 20 +++++++++++++++++-- .../springframework/graphql/LocationArea.java | 16 +++++++++++++++ .../graphql/LocationAreaId.java | 20 +++++++++++++++++-- .../EntityMappingInvocationTests.java | 4 +--- 5 files changed, 69 insertions(+), 7 deletions(-) diff --git a/spring-graphql/src/test/java/org/springframework/graphql/Library.java b/spring-graphql/src/test/java/org/springframework/graphql/Library.java index 99fd217b..390dadec 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/Library.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/Library.java @@ -1,3 +1,19 @@ +/* + * 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) { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java b/spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java index c3327294..8a43a8f7 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/LibraryId.java @@ -1,8 +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) { - } + public record LocationId(String id) { + } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java b/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java index b3bf959c..6abfe236 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/LocationArea.java @@ -1,3 +1,19 @@ +/* + * 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) { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java b/spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java index 0179a6d1..94d0dfb4 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/LocationAreaId.java @@ -1,8 +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) { - } + public record LocationId(String id) { + } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java index e6b64642..0ae0a119 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java @@ -256,9 +256,7 @@ void unmappedEntity() { .withMessage("Unmapped entity types: 'Media', 'PrintedMedia', 'Book'"); } - private static ResponseHelper executeWith(Class controllerClass, Resource federationSchema, String document, - Map variables) { - + private static ResponseHelper executeWith(Class controllerClass, Resource federationSchema, String document, Map variables) { ExecutionGraphQlRequest request = TestExecutionRequest.forDocumentAndVars(document, variables); Mono responseMono = graphQlService(controllerClass, federationSchema).execute(request); return ResponseHelper.forResponse(responseMono); From 5d2d43ac35517200def345724994e9e03d61df1e Mon Sep 17 00:00:00 2001 From: James Bodkin Date: Wed, 1 Jul 2026 15:38:05 +0100 Subject: [PATCH 3/3] Refactor implementation for nested key entity support Signed-off-by: James Bodkin --- .../data/federation/EntitiesDataFetcher.java | 9 +- .../EntityArgumentMethodArgumentResolver.java | 53 +++++----- .../data/federation/EntityHandlerMethod.java | 8 +- .../data/federation/EntityKeyResolver.java | 98 +++++++++++++++++++ .../federation/FederationSchemaFactory.java | 6 +- .../org/springframework/graphql/Location.java | 16 +++ 6 files changed, 152 insertions(+), 38 deletions(-) create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityKeyResolver.java diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntitiesDataFetcher.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntitiesDataFetcher.java index 631da826..5e74b836 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntitiesDataFetcher.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntitiesDataFetcher.java @@ -57,14 +57,17 @@ final class EntitiesDataFetcher implements DataFetcher handlerMethods, Map 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 invokeEntityMethod( DataFetchingEnvironment environment, EntityHandlerMethod handlerMethod, Map 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 invokeEntitiesMethod( } } - return handlerMethod.getEntities(environment, typeRepresentations) + return handlerMethod.getEntities(environment, typeRepresentations, this.entityKeyResolver) .mapNotNull((result) -> (((List) result).isEmpty()) ? null : result) .switchIfEmpty(Mono.defer(() -> { List> exceptions = new ArrayList<>(originalIndexes.size()); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java index a76b8acc..500fd093 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java @@ -24,7 +24,6 @@ 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; @@ -52,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 representation : batchEnv.getRepresentations()) { - values.add(doBind(name, targetType, representation)); + Object key = batchEnv.getKey(representation); + values.add(doBind(name, targetType, key)); } return values; } @@ -68,26 +69,8 @@ else if (environment instanceof EntityBatchDataFetchingEnvironment batchEnv) { } } - private @Nullable Object doBind(String name, ResolvableType targetType, Map entityMap) throws BindException { - if (isScalarValue(entityMap)) { - Object rawValue = entityMap.get(name); - return getArgumentBinder().bind(rawValue, false, targetType); - } - return getArgumentBinder().bind(entityMap, false, targetType); - } - - private boolean isScalarValue(Map entityMap) { - if (entityMap.size() != 2) { - return false; - } - - for (Map.Entry entry : entityMap.entrySet()) { - if (!"__typename".equals(entry.getKey())) { - Object value = entry.getValue(); - return BeanUtils.isSimpleValueType(value.getClass()); - } - } - return false; + private @Nullable Object doBind(String name, ResolvableType targetType, Object key) throws BindException { + return getArgumentBinder().bind(key, false, targetType); } private static String dePluralize(String name) { @@ -99,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 representation) { - return new EntityDataFetchingEnvironment(env, representation); + static DataFetchingEnvironment wrap(DataFetchingEnvironment env, Map 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> representations) { - return new EntityBatchDataFetchingEnvironment(env, representations); + static DataFetchingEnvironment wrap(DataFetchingEnvironment env, List> representations, EntityKeyResolver resolver) { + return new EntityBatchDataFetchingEnvironment(env, representations, resolver); } @@ -118,15 +101,21 @@ static DataFetchingEnvironment wrap(DataFetchingEnvironment env, List representation; + private final EntityKeyResolver resolver; - EntityDataFetchingEnvironment(DataFetchingEnvironment env, Map representation) { + EntityDataFetchingEnvironment(DataFetchingEnvironment env, Map representation, EntityKeyResolver resolver) { super(env); this.representation = representation; + this.resolver = resolver; } Map getRepresentation() { return this.representation; } + + public Object getKey() { + return this.resolver.getKey(representation); + } } @@ -137,15 +126,21 @@ Map getRepresentation() { static class EntityBatchDataFetchingEnvironment extends DelegatingDataFetchingEnvironment { private final List> representations; + private final EntityKeyResolver resolver; - EntityBatchDataFetchingEnvironment(DataFetchingEnvironment env, List> representations) { + EntityBatchDataFetchingEnvironment(DataFetchingEnvironment env, List> representations, EntityKeyResolver resolver) { super(env); this.representations = representations; + this.resolver = resolver; } List> getRepresentations() { return this.representations; } + + public Object getKey(Map representation) { + return this.resolver.getKey(representation); + } } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java index 2541ff1b..8a49be6e 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java @@ -52,13 +52,13 @@ boolean isBatchHandlerMethod() { } - Mono getEntity(DataFetchingEnvironment env, Map representation) { - env = EntityArgumentMethodArgumentResolver.wrap(env, representation); + Mono getEntity(DataFetchingEnvironment env, Map representation, EntityKeyResolver resolver) { + env = EntityArgumentMethodArgumentResolver.wrap(env, representation, resolver); return doInvoke(env); } - Mono getEntities(DataFetchingEnvironment env, List> representations) { - env = EntityArgumentMethodArgumentResolver.wrap(env, representations); + Mono getEntities(DataFetchingEnvironment env, List> representations, EntityKeyResolver resolver) { + env = EntityArgumentMethodArgumentResolver.wrap(env, representations, resolver); return doInvoke(env); } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityKeyResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityKeyResolver.java new file mode 100644 index 00000000..c24f2cbf --- /dev/null +++ b/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. + *

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")}). + * + *

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 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. + * + *

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 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; + } + } + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java index bf629ca2..11dbeecb 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java @@ -188,8 +188,10 @@ public SchemaTransformer createSchemaTransformer(TypeDefinitionRegistry registry Map 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) diff --git a/spring-graphql/src/test/java/org/springframework/graphql/Location.java b/spring-graphql/src/test/java/org/springframework/graphql/Location.java index f2c26d7a..4a526b23 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/Location.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/Location.java @@ -1,3 +1,19 @@ +/* + * 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) {