Support interfaces in EntityMapping's

Closes gh-1356

Author: rstoyanchev

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

@@ -52,13 +52,18 @@ final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<Li
 
 	private final Map<String, EntityHandlerMethod> handlerMethods;
 
+	/** schema object type to interface types it implements. */
+	private final Map<String, String> objectToInterfaceMap;
+
 	private final HandlerDataFetcherExceptionResolver exceptionResolver;
 
 
 	EntitiesDataFetcher(
-			Map<String, EntityHandlerMethod> handlerMethods, HandlerDataFetcherExceptionResolver resolver) {
+			Map<String, EntityHandlerMethod> handlerMethods, Map<String, String> objectToInterfaceMap,
+			HandlerDataFetcherExceptionResolver resolver) {
 
 		this.handlerMethods = new LinkedHashMap<>(handlerMethods);
+		this.objectToInterfaceMap = objectToInterfaceMap;
 		this.exceptionResolver = resolver;
 	}
 
@@ -84,6 +89,10 @@ public Mono<DataFetcherResult<List<Object>>> get(DataFetchingEnvironment env) {
 				continue;
 			}
 			EntityHandlerMethod handlerMethod = this.handlerMethods.get(type);
+			if (handlerMethod == null) {
+				String interfaceType = this.objectToInterfaceMap.get(type);
+				handlerMethod = this.handlerMethods.get(interfaceType);
+			}
 			if (handlerMethod == null) {
 				Exception ex = new RepresentationException(map, "No entity fetcher");
 				monoList.add(resolveException(ex, env, null, index));

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

@@ -29,7 +29,10 @@
 import graphql.language.Argument;
 import graphql.language.BooleanValue;
 import graphql.language.Directive;
+import graphql.language.ObjectTypeDefinition;
+import graphql.language.Type;
 import graphql.language.TypeDefinition;
+import graphql.language.TypeName;
 import graphql.schema.DataFetcher;
 import graphql.schema.GraphQLSchema;
 import graphql.schema.TypeResolver;
@@ -180,18 +183,49 @@ public GraphQLSchema createGraphQLSchema(TypeDefinitionRegistry registry, Runtim
 	 * @param wiring the existing runtime wiring
 	 */
 	public SchemaTransformer createSchemaTransformer(TypeDefinitionRegistry registry, RuntimeWiring wiring) {
-		checkEntityMappings(registry);
-		Assert.state(this.typeResolver != null, "afterPropertiesSet not called");
+		Assert.state(this.typeResolver != null, "Not initialized: was afterPropertiesSet called?");
+
+		Map<String, String> objectToInterfaceTypeMap = detectInterfaceImplementationTypes(registry);
+		checkEntityMappings(registry, objectToInterfaceTypeMap);
+
+		EntitiesDataFetcher entitiesDataFetcher =
+				new EntitiesDataFetcher(this.handlerMethods, objectToInterfaceTypeMap, getExceptionResolver());
+
 		return Federation.transform(registry, wiring)
-				.fetchEntities(new EntitiesDataFetcher(this.handlerMethods, getExceptionResolver()))
+				.fetchEntities(entitiesDataFetcher)
 				.resolveEntityType(this.typeResolver);
 	}
 
-	private void checkEntityMappings(TypeDefinitionRegistry registry) {
+	/**
+	 * For all schema interface types referenced in entity mappings, return a
+	 * lookup from schema object type to the interface it implements.
+	 */
+	private Map<String, String> detectInterfaceImplementationTypes(TypeDefinitionRegistry registry) {
+		Map<String, String> map = new LinkedHashMap<>();
+		for (TypeDefinition<?> typeDef : registry.types().values()) {
+			if (typeDef instanceof ObjectTypeDefinition objectDef) {
+				for (Type<?> type : objectDef.getImplements()) {
+					String interfaceName = ((TypeName) type).getName();
+					if (this.handlerMethods.containsKey(interfaceName)) {
+						String objectTypeName = objectDef.getName();
+						String existingInterface = map.put(objectTypeName, interfaceName);
+						Assert.state(existingInterface == null, () ->
+								"Object type '" + objectTypeName + "' implements two EntityMapping interfaces: '" +
+										interfaceName + "' and '" + existingInterface + "'.");
+					}
+				}
+			}
+		}
+		return map;
+	}
+
+	private void checkEntityMappings(TypeDefinitionRegistry registry, Map<String, String> objectToInterfaceTypeMap) {
 		List<String> unmappedEntities = new ArrayList<>();
 		for (TypeDefinition<?> type : registry.types().values()) {
 			type.getDirectives().forEach((directive) -> {
-				if (isResolvableKeyDirective(directive) && !this.handlerMethods.containsKey(type.getName())) {
+				if (isResolvableKeyDirective(directive) &&
+						!this.handlerMethods.containsKey(type.getName()) &&
+						!objectToInterfaceTypeMap.containsKey(type.getName())) {
 					unmappedEntities.add(type.getName());
 				}
 			});

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

@@ -195,7 +195,7 @@ void dataLoader() {
 	@Test
 	void unmappedEntity() {
 		assertThatIllegalStateException().isThrownBy(() -> executeWith(EmptyController.class, Map.of()))
-				.withMessage("Unmapped entity types: 'Book'");
+				.withMessage("Unmapped entity types: 'Media', 'Book'");
 	}
 
 	private static ResponseHelper executeWith(Class<?> controllerClass, Map<String, Object> variables) {
@@ -247,7 +247,7 @@ private static TestExecutionGraphQlService graphQlService(Class<?> controllerCla
 	private static class BookController {
 
 		@Nullable
-		@EntityMapping
+		@EntityMapping("Media")
 		public Book book(@Argument int id, Map<String, Object> map) {
 
 			assertThat(map).hasSize(2)
@@ -282,7 +282,7 @@ private static class BookListController {
 
 		private final BookBatchService batchService = new BookBatchService();
 
-		@EntityMapping
+		@EntityMapping("Media")
 		public List<Book> book(@Argument List<Integer> idList, List<Map<String, Object>> representations) {
 			return this.batchService.book(idList, representations);
 		}
@@ -305,7 +305,7 @@ private static class BookFluxController {
 
 		private final BookBatchService batchService = new BookBatchService();
 
-		@EntityMapping
+		@EntityMapping("Media")
 		public Flux<Book> book(@Argument List<Integer> idList, List<Map<String, Object>> representations) {
 			return Flux.fromIterable(this.batchService.book(idList, representations));
 		}
@@ -334,7 +334,7 @@ public DataLoaderBookController(BatchLoaderRegistry batchLoaderRegistry) {
 		}
 
 		@Nullable
-		@EntityMapping
+		@EntityMapping("Media")
 		public Future<Book> book(@Argument int id, DataLoader<Integer, Book> dataLoader) {
 			return dataLoader.load(id);
 		}

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

@@ -1,6 +1,10 @@
 extend schema @link(url: "https://specs.apollo.dev/federation/v2.9", import: ["@key", "@extends", "@external"] )
 
-type Book @key(fields: "id") @extends {
+interface Media @key(fields: "id") {
+    id: ID! @external
+}
+
+type Book implements Media @key(fields: "id") @extends {
     id: ID! @external
     author: Author
     publisher: Publisher