From 2f1c4b1b99464608d04b213a6875887e71a9dd3e Mon Sep 17 00:00:00 2001
From: rstoyanchev <rossen.stoyanchev@broadcom.com>
Date: Thu, 20 Feb 2025 11:23:45 +0000
Subject: [PATCH] FederationSchemaFactory checks for unmapped entities

Closes gh-1088
---
 .../federation/FederationSchemaFactory.java   | 14 +++++++++++-
 .../EntityMappingInvocationTests.java         | 22 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

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 444d592f..269d7731 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -25,6 +25,7 @@ import java.util.stream.Collectors;
 
 import com.apollographql.federation.graphqljava.Federation;
 import com.apollographql.federation.graphqljava.SchemaTransformer;
+import graphql.language.TypeDefinition;
 import graphql.schema.DataFetcher;
 import graphql.schema.GraphQLSchema;
 import graphql.schema.TypeResolver;
@@ -177,12 +178,23 @@ public final class FederationSchemaFactory
 	 * @param wiring the existing runtime wiring
 	 */
 	public SchemaTransformer createSchemaTransformer(TypeDefinitionRegistry registry, RuntimeWiring wiring) {
+		checkEntityMappings(registry);
 		Assert.state(this.typeResolver != null, "afterPropertiesSet not called");
 		return Federation.transform(registry, wiring)
 				.fetchEntities(new EntitiesDataFetcher(this.handlerMethods, getExceptionResolver()))
 				.resolveEntityType(this.typeResolver);
 	}
 
+	private void checkEntityMappings(TypeDefinitionRegistry registry) {
+		for (TypeDefinition<?> type : registry.types().values()) {
+			type.getDirectives().forEach((directive) -> {
+				boolean isEntityType = directive.getName().equalsIgnoreCase("key");
+				Assert.state(!isEntityType || this.handlerMethods.containsKey(type.getName()),
+						"No EntityMapping method for federated type: '" + type.getName() + "'");
+			});
+		}
+	}
+
 
 	public record EntityMappingInfo(String typeName, HandlerMethod handlerMethod) {
 
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 448ef2e0..0f26ad20 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -52,6 +52,7 @@ import org.springframework.lang.Nullable;
 import org.springframework.stereotype.Controller;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 
 /**
  * Tests for requests handled through {@code @EntityMapping} methods.
@@ -175,6 +176,18 @@ public class EntityMappingInvocationTests {
 		assertError(helper, 2, "INTERNAL_ERROR", "Entity fetcher returned null or completed empty");
 	}
 
+	@Test
+	void unmappedEntity() {
+		Map<String, Object> variables =
+				Map.of("representations", List.of(
+						Map.of("__typename", "Book", "id", "-99"),
+						Map.of("__typename", "Book", "id", "4"),
+						Map.of("__typename", "Book", "id", "5")));
+
+		assertThatIllegalStateException().isThrownBy(() -> executeWith(EmptyController.class, variables))
+				.withMessage("No EntityMapping method for federated type: 'Book'");
+	}
+
 	private static ResponseHelper executeWith(Class<?> controllerClass, Map<String, Object> variables) {
 		ExecutionGraphQlRequest request = TestExecutionRequest.forDocumentAndVars(document, variables);
 		Mono<ExecutionGraphQlResponse> responseMono = graphQlService(controllerClass).execute(request);
@@ -299,6 +312,13 @@ public class EntityMappingInvocationTests {
 	}
 
 
+	@SuppressWarnings("unused")
+	@Controller
+	private static class EmptyController {
+
+	}
+
+
 	private static class BookBatchService {
 
 		public List<Book> book(List<Integer> idList, List<Map<String, Object>> representations) {