From 5977b2d8b7bb4f296e148d8026f0fa7c6b29cced Mon Sep 17 00:00:00 2001
From: Brian Clozel <brian.clozel@broadcom.com>
Date: Fri, 19 Apr 2024 18:18:09 +0200
Subject: [PATCH] Register runtime hints for @EntityMapping methods

Since the introduction of `@EntityMapping` support on controller
handlers, we need to support such variants for the GraalVM Native case.

This commit registers the relevant hints for runtime reflection support
with GraalVM Native.

Closes gh-928
---
 ...BeanFactoryInitializationAotProcessor.java |  10 +-
 ...actoryInitializationAotProcessorTests.java | 110 +++++++++++++++++-
 2 files changed, 115 insertions(+), 5 deletions(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessor.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessor.java
index 397e9daa..518c6fbc 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessor.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessor.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2023 the original author or authors.
+ * Copyright 2020-2024 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.
@@ -44,6 +44,7 @@ import org.springframework.core.annotation.MergedAnnotations;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
 import org.springframework.data.projection.TargetAware;
 import org.springframework.graphql.data.ArgumentValue;
+import org.springframework.graphql.data.federation.EntityMapping;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.data.method.annotation.BatchMapping;
@@ -58,8 +59,8 @@ import org.springframework.web.bind.annotation.ControllerAdvice;
  * {@link BeanFactoryInitializationAotProcessor} implementation for registering
  * runtime hints discoverable through GraphQL controllers, such as:
  * <ul>
- * <li>invocation reflection on {@code @SchemaMapping} and {@code @BatchMapping}
- * annotated controllers methods
+ * <li>invocation reflection on {@code @SchemaMapping}, {@code @BatchMapping}
+ * and {@code @EntityMapping} annotated controllers methods
  * <li>invocation reflection on {@code @GraphQlExceptionHandler} methods
  * in {@code @Controller} and {@code @ControllerAdvice} beans
  * <li>binding reflection on controller method arguments, needed for binding or
@@ -167,7 +168,8 @@ class SchemaMappingBeanFactoryInitializationAotProcessor implements BeanFactoryI
 
 		private boolean isGraphQlHandlerMethod(AnnotatedElement element) {
 			MergedAnnotations annotations = MergedAnnotations.from(element, MergedAnnotations.SearchStrategy.TYPE_HIERARCHY);
-			return annotations.isPresent(SchemaMapping.class) || annotations.isPresent(BatchMapping.class);
+			return annotations.isPresent(SchemaMapping.class) || annotations.isPresent(BatchMapping.class)
+					|| annotations.isPresent(EntityMapping.class);
 		}
 
 		private boolean isExceptionHandlerMethod(AnnotatedElement element) {
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessorTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessorTests.java
index 198619a7..295aa563 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessorTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingBeanFactoryInitializationAotProcessorTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2023 the original author or authors.
+ * Copyright 2020-2024 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.
@@ -59,6 +59,7 @@ import org.springframework.data.web.ProjectedPayload;
 import org.springframework.graphql.Author;
 import org.springframework.graphql.Book;
 import org.springframework.graphql.data.ArgumentValue;
+import org.springframework.graphql.data.federation.EntityMapping;
 import org.springframework.graphql.data.method.annotation.Argument;
 import org.springframework.graphql.data.method.annotation.BatchMapping;
 import org.springframework.graphql.data.method.annotation.ContextValue;
@@ -451,6 +452,113 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 		}
 	}
 
+	@Nested
+	class EntityMappingTests {
+
+		@Test
+		void registerBindingReflectionOnReturnType() {
+			processBeanClasses(ReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(ReturnTypeController.class, "bookById");
+			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class);
+		}
+
+		@Test
+		void registerBindingReflectionOnOnNamedValue() {
+			processBeanClasses(NamedValueController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(NamedValueController.class);
+			assertThatInvocationHintRegisteredForMethods(NamedValueController.class, "book");
+			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, Identifier.class);
+		}
+
+		@Test
+		void registerBindingReflectionOnOnNamedValues() {
+			processBeanClasses(NamedValuesController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(NamedValuesController.class);
+			assertThatInvocationHintRegisteredForMethods(NamedValuesController.class, "books");
+			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, Identifier.class);
+		}
+
+		@Test
+		void registerBindingReflectionOnAsyncReturnType() {
+			processBeanClasses(AsyncReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(AsyncReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(AsyncReturnTypeController.class, "author");
+			assertThatHintsForJavaBeanBindingRegisteredForTypes(Author.class);
+		}
+
+		@Test
+		void doNotRegisterBindingForContextArguments() {
+			processBeanClasses(ContextArgumentsController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ContextArgumentsController.class);
+			assertThatInvocationHintRegisteredForMethods(ContextArgumentsController.class, "dataFetchingEnvironment");
+			assertThatHintsAreNotRegisteredForTypes(GraphQLContext.class, DataFetchingFieldSelectionSet.class, Locale.class);
+		}
+
+		@Test
+		void doNotRegisterBindingForAnnotatedContextArguments() {
+			processBeanClasses(AnnotatedContextArgumentController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(AnnotatedContextArgumentController.class);
+			assertThatInvocationHintRegisteredForMethods(AnnotatedContextArgumentController.class, "contextValue");
+			assertThatHintsAreNotRegisteredForTypes(Book.class);
+		}
+
+
+		@Controller
+		static class ReturnTypeController {
+			@EntityMapping
+			public Book bookById(@Argument Long id) {
+				return null;
+			}
+
+		}
+
+		@Controller
+		static class NamedValueController {
+			@EntityMapping
+			public Book book(@Argument Identifier id) {
+				return null;
+			}
+
+		}
+
+		@Controller
+		static class NamedValuesController {
+			@EntityMapping
+			public List<Book> books(@Argument List<Identifier> idList) {
+				return null;
+			}
+
+		}
+
+		@Controller
+		static class AsyncReturnTypeController {
+			@EntityMapping
+			public CompletableFuture<Author> author(Long bookId) {
+				return null;
+			}
+
+		}
+
+		@Controller
+		static class ContextArgumentsController {
+			@EntityMapping
+			public void dataFetchingEnvironment(GraphQLContext context, DataFetchingFieldSelectionSet selectionSet, Locale locale) {
+			}
+		}
+
+		@Controller
+		class AnnotatedContextArgumentController {
+			@EntityMapping
+			public void contextValue(@ContextValue Book book, @LocalContextValue Book localBook) {
+			}
+		}
+
+		record Identifier(String id) {
+
+		}
+	}
+
 
 	private void processBeanClasses(Class<?>... beanClasses) {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();