From 14baa4e32a906ad415a0d48c0665e79b6ee8f6a4 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rossen.stoyanchev@broadcom.com>
Date: Thu, 29 May 2025 17:07:52 +0100
Subject: [PATCH] Polishing contribution

Closes: gh-1225
---
 .../federation/FederationSchemaFactory.java   | 33 +++++++++----------
 1 file changed, 16 insertions(+), 17 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 095d55fd..3e2835e1 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
@@ -193,9 +193,11 @@ public final class FederationSchemaFactory
 	private void checkEntityMappings(TypeDefinitionRegistry registry) {
 		List<String> unmappedEntities = new ArrayList<>();
 		for (TypeDefinition<?> type : registry.types().values()) {
-			if (isEntityMappingExpected(type) && !this.handlerMethods.containsKey(type.getName())) {
-				unmappedEntities.add(type.getName());
-			}
+			type.getDirectives().forEach((directive) -> {
+				if (isResolvableKeyDirective(directive) && !this.handlerMethods.containsKey(type.getName())) {
+					unmappedEntities.add(type.getName());
+				}
+			});
 		}
 		if (!unmappedEntities.isEmpty()) {
 			throw new IllegalStateException("Unmapped entity types: " +
@@ -203,22 +205,19 @@ public final class FederationSchemaFactory
 		}
 	}
 
-	/**
-	 * Determine if a handler method is expected for this type: there is at least one '@key' directive
-	 * whose 'resolvable' argument resolves to true (either explicitly, or if the argument is not set).
-	 * @param type the type to inspect.
-	 * @return true if a handler method is expected for this type
-	 */
-	private boolean isEntityMappingExpected(TypeDefinition<?> type) {
-		List<Directive> keyDirectives = type.getDirectives("key");
-		return !keyDirectives.isEmpty() && keyDirectives.stream()
-			.anyMatch((keyDirective) -> {
-				Argument resolvableArg = keyDirective.getArgument("resolvable");
-				return resolvableArg == null ||
-					(resolvableArg.getValue() instanceof BooleanValue) && ((BooleanValue) resolvableArg.getValue()).isValue();
-			});
+	private boolean isResolvableKeyDirective(Directive directive) {
+		if (!directive.getName().equalsIgnoreCase("key")) {
+			return false;
+		}
+		Argument argument = directive.getArgument("resolvable");
+		if (argument != null) {
+			Object value = argument.getValue();
+			return (value instanceof BooleanValue bv && bv.isValue());
+		}
+		return true;
 	}
 
+
 	public record EntityMappingInfo(String typeName, HandlerMethod handlerMethod) {
 
 		public boolean isBatchHandlerMethod() {