From 01bc66acb16f4ce82fa1227b8b68d24240c10386 Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@vmware.com>
Date: Tue, 25 Apr 2023 17:00:13 +0200
Subject: [PATCH] Add support for GraphQlExceptionHandler in AOT mode

This commit adds the relevant reflection metadata for supporting
`@GraphQlExceptionHandler` annotated methods in controllers and
`@ControllerAdvice` beans.

Closes gh-677
---
 ...BeanFactoryInitializationAotProcessor.java |  54 +++++--
 .../spring-graphql/reflect-config.json        |   7 +
 ...actoryInitializationAotProcessorTests.java | 136 ++++++++++++++----
 3 files changed, 160 insertions(+), 37 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 68e330c6..bc4676b9 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
@@ -20,7 +20,9 @@ import java.lang.reflect.AnnotatedElement;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
 import java.lang.reflect.Type;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
 import org.springframework.aop.SpringProxy;
 import org.springframework.aot.generate.GenerationContext;
@@ -45,10 +47,12 @@ import org.springframework.graphql.data.ArgumentValue;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.data.method.annotation.BatchMapping;
+import org.springframework.graphql.data.method.annotation.GraphQlExceptionHandler;
 import org.springframework.graphql.data.method.annotation.SchemaMapping;
 import org.springframework.stereotype.Controller;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.ReflectionUtils;
+import org.springframework.web.bind.annotation.ControllerAdvice;
 
 import static org.springframework.core.annotation.MergedAnnotations.SearchStrategy.TYPE_HIERARCHY;
 
@@ -58,6 +62,8 @@ import static org.springframework.core.annotation.MergedAnnotations.SearchStrate
  * <ul>
  * <li>invocation reflection on {@code @SchemaMapping} and {@code @BatchMapping}
  * 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
  * by the GraphQL Java engine itself
  * <li>reflection for SpEL support and JDK proxy creation for
@@ -85,27 +91,42 @@ class SchemaMappingBeanFactoryInitializationAotProcessor implements BeanFactoryI
 
 	@Override
 	public BeanFactoryInitializationAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
-		Class<?>[] controllerTypes = Arrays.stream(beanFactory.getBeanDefinitionNames())
+		List<Class<?>> controllers = new ArrayList<>();
+		List<Class<?>> controllerAdvices = new ArrayList<>();
+		Arrays.stream(beanFactory.getBeanDefinitionNames())
 				.map(beanName -> RegisteredBean.of(beanFactory, beanName).getBeanClass())
-				.filter(this::isController)
-				.toArray(Class<?>[]::new);
-		return new SchemaMappingBeanFactoryInitializationAotContribution(controllerTypes);
+				.forEach(beanClass -> {
+					if (isController(beanClass)) {
+						controllers.add(beanClass);
+					}
+					else if (isControllerAdvice(beanClass)) {
+						controllerAdvices.add(beanClass);
+					}
+				});
+		return new SchemaMappingBeanFactoryInitializationAotContribution(controllers, controllerAdvices);
 	}
 
 	private boolean isController(AnnotatedElement element) {
 		return MergedAnnotations.from(element, TYPE_HIERARCHY).isPresent(Controller.class);
 	}
 
+	private boolean isControllerAdvice(AnnotatedElement element) {
+		return MergedAnnotations.from(element, TYPE_HIERARCHY).isPresent(ControllerAdvice.class);
+	}
+
 
 	private static class SchemaMappingBeanFactoryInitializationAotContribution
 			implements BeanFactoryInitializationAotContribution {
 
-		private final Class<?>[] controllers;
+		private final List<Class<?>> controllers;
+
+		private final List<Class<?>> controllerAdvices;
 
 		private final HandlerMethodArgumentResolverComposite argumentResolvers;
 
-		public SchemaMappingBeanFactoryInitializationAotContribution(Class<?>[] controllers) {
+		public SchemaMappingBeanFactoryInitializationAotContribution(List<Class<?>> controllers, List<Class<?>> controllerAdvices) {
 			this.controllers = controllers;
+			this.controllerAdvices = controllerAdvices;
 			this.argumentResolvers = createArgumentResolvers();
 		}
 
@@ -120,11 +141,20 @@ class SchemaMappingBeanFactoryInitializationAotProcessor implements BeanFactoryI
 		public void applyTo(GenerationContext context, BeanFactoryInitializationCode initializationCode) {
 			RuntimeHints runtimeHints = context.getRuntimeHints();
 			registerSpringDataSpelSupport(runtimeHints);
-			Arrays.stream(this.controllers).forEach(controller -> {
-				runtimeHints.reflection().registerType(controller);
+			this.controllers.forEach(controller -> {
+				runtimeHints.reflection().registerType(controller, MemberCategory.INTROSPECT_DECLARED_METHODS);
 				ReflectionUtils.doWithMethods(controller,
 						method -> processSchemaMappingMethod(runtimeHints, method),
 						this::isGraphQlHandlerMethod);
+				ReflectionUtils.doWithMethods(controller,
+						method -> processExceptionHandlerMethod(runtimeHints, method),
+						this::isExceptionHandlerMethod);
+			});
+			this.controllerAdvices.forEach(controllerAdvice -> {
+				runtimeHints.reflection().registerType(controllerAdvice, MemberCategory.INTROSPECT_DECLARED_METHODS);
+				ReflectionUtils.doWithMethods(controllerAdvice,
+						method -> processExceptionHandlerMethod(runtimeHints, method),
+						this::isExceptionHandlerMethod);
 			});
 		}
 
@@ -143,6 +173,10 @@ class SchemaMappingBeanFactoryInitializationAotProcessor implements BeanFactoryI
 			return annotations.isPresent(SchemaMapping.class) || annotations.isPresent(BatchMapping.class);
 		}
 
+		private boolean isExceptionHandlerMethod(AnnotatedElement element) {
+			return MergedAnnotations.from(element, TYPE_HIERARCHY).isPresent(GraphQlExceptionHandler.class);
+		}
+
 		private void processSchemaMappingMethod(RuntimeHints runtimeHints, Method method) {
 			runtimeHints.reflection().registerMethod(method, ExecutableMode.INVOKE);
 			for (Parameter parameter : method.getParameters()) {
@@ -151,6 +185,10 @@ class SchemaMappingBeanFactoryInitializationAotProcessor implements BeanFactoryI
 			processReturnType(runtimeHints, MethodParameter.forExecutable(method, -1));
 		}
 
+		private void processExceptionHandlerMethod(RuntimeHints runtimeHints, Method method) {
+			runtimeHints.reflection().registerMethod(method, ExecutableMode.INVOKE);
+		}
+
 		private void processMethodParameter(RuntimeHints hints, MethodParameter parameter) {
 			MethodParameterRuntimeHintsRegistrar.fromMethodParameter(this.argumentResolvers, parameter).apply(hints);
 		}
diff --git a/spring-graphql/src/main/resources/META-INF/native-image/org.springframework.graphql/spring-graphql/reflect-config.json b/spring-graphql/src/main/resources/META-INF/native-image/org.springframework.graphql/spring-graphql/reflect-config.json
index 98b48667..4ffdc240 100644
--- a/spring-graphql/src/main/resources/META-INF/native-image/org.springframework.graphql/spring-graphql/reflect-config.json
+++ b/spring-graphql/src/main/resources/META-INF/native-image/org.springframework.graphql/spring-graphql/reflect-config.json
@@ -1,4 +1,11 @@
 [
+  {
+    "name": "org.springframework.graphql.data.method.annotation.support.AnnotatedControllerExceptionResolver$MethodResolver",
+    "allDeclaredMethods": true,
+    "condition": {
+      "typeReachable": "org.springframework.graphql.data.method.annotation.support.AnnotatedControllerExceptionResolver"
+    }
+  },
   {
     "name":"org.springframework.graphql.server.support.GraphQlWebSocketMessage",
     "allDeclaredFields":true,
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 3ffb582d..0bff594c 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-2022 the original author or authors.
+ * Copyright 2020-2023 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.
@@ -32,10 +32,14 @@ import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
 
 import graphql.GraphQLContext;
+import graphql.GraphQLError;
 import graphql.schema.DataFetchingFieldSelectionSet;
 import org.dataloader.DataLoader;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
+import org.springframework.graphql.data.method.annotation.*;
+import org.springframework.validation.BindException;
+import org.springframework.web.bind.annotation.ControllerAdvice;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -58,13 +62,6 @@ 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.method.annotation.Argument;
-import org.springframework.graphql.data.method.annotation.BatchMapping;
-import org.springframework.graphql.data.method.annotation.ContextValue;
-import org.springframework.graphql.data.method.annotation.LocalContextValue;
-import org.springframework.graphql.data.method.annotation.MutationMapping;
-import org.springframework.graphql.data.method.annotation.QueryMapping;
-import org.springframework.graphql.data.method.annotation.SchemaMapping;
 import org.springframework.stereotype.Controller;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -92,39 +89,51 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 		@Test
 		void registerBindingReflectionOnReturnType() {
-			processController(ReturnTypeController.class);
+			processBeanClasses(ReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(ReturnTypeController.class, "bookById");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnInput() {
-			processController(InputController.class);
+			processBeanClasses(InputController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(InputController.class);
+			assertThatInvocationHintRegisteredForMethods(InputController.class, "addBook");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, BookInput.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnArgumentCollection() {
-			processController(ArgumentCollectionController.class);
+			processBeanClasses(ArgumentCollectionController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ArgumentCollectionController.class);
+			assertThatInvocationHintRegisteredForMethods(ArgumentCollectionController.class, "addBooks");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnArgumentValue() {
-			processController(ArgumentValueController.class);
+			processBeanClasses(ArgumentValueController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ArgumentValueController.class);
+			assertThatInvocationHintRegisteredForMethods(ArgumentValueController.class, "addBook");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, BookInput.class);
 			assertThatHintsAreNotRegisteredForTypes(ArgumentValue.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnDataLoaderArgument() {
-			processController(DataLoaderController.class);
+			processBeanClasses(DataLoaderController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(DataLoaderController.class);
+			assertThatInvocationHintRegisteredForMethods(DataLoaderController.class, "authorWithLoader");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Author.class);
 			assertThatHintsAreNotRegisteredForTypes(DataLoader.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnAsyncReturnType() {
-			processController(AsyncReturnTypeController.class);
+			processBeanClasses(AsyncReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(AsyncReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(AsyncReturnTypeController.class, "author");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Author.class);
 		}
 
@@ -211,37 +220,49 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 		@Test
 		void registerBindingReflectionOnFluxReturnType() {
-			processController(FluxReturnTypeController.class);
+			processBeanClasses(FluxReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(FluxReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(FluxReturnTypeController.class, "batchMappingFlux");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Author.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnMonoMapReturnType() {
-			processController(MonoMapReturnTypeController.class);
+			processBeanClasses(MonoMapReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(MonoMapReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(MonoMapReturnTypeController.class, "batchMappingMono");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, Author.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnMapReturnType() {
-			processController(MapReturnTypeController.class);
+			processBeanClasses(MapReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(MapReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(MapReturnTypeController.class, "batchMappingMap");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, Author.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnCollectionReturnType() {
-			processController(CollectionReturnTypeController.class);
+			processBeanClasses(CollectionReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(CollectionReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(CollectionReturnTypeController.class, "batchMappingCollection");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Author.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnCallableCollectionReturnType() {
-			processController(CallableCollectionReturnTypeController.class);
+			processBeanClasses(CallableCollectionReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(CallableCollectionReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(CallableCollectionReturnTypeController.class, "batchMappingCallableCollection");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Author.class);
 		}
 
 		@Test
 		void registerBindingReflectionOnCallableMapReturnType() {
-			processController(CallableMapReturnTypeController.class);
+			processBeanClasses(CallableMapReturnTypeController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(CallableMapReturnTypeController.class);
+			assertThatInvocationHintRegisteredForMethods(CallableMapReturnTypeController.class, "batchMappingCallableMap");
 			assertThatHintsForJavaBeanBindingRegisteredForTypes(Book.class, Author.class);
 		}
 
@@ -302,13 +323,17 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 		@Test
 		void doNotRegisterBindingForContextArguments() {
-			processController(ContextArgumentsController.class);
+			processBeanClasses(ContextArgumentsController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ContextArgumentsController.class);
+			assertThatInvocationHintRegisteredForMethods(ContextArgumentsController.class, "dataFetchingEnvironment");
 			assertThatHintsAreNotRegisteredForTypes(GraphQLContext.class, DataFetchingFieldSelectionSet.class, Locale.class);
 		}
 
 		@Test
 		void doNotRegisterBindingForAnnotatedContextArguments() {
-			processController(AnnotatedContextArgumentController.class);
+			processBeanClasses(AnnotatedContextArgumentController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(AnnotatedContextArgumentController.class);
+			assertThatInvocationHintRegisteredForMethods(AnnotatedContextArgumentController.class, "contextValue");
 			assertThatHintsAreNotRegisteredForTypes(Book.class);
 		}
 
@@ -333,7 +358,7 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 		@Test
 		void registerSpringDataSpelSupport() {
-			processController();
+			processBeanClasses();
 			TypeReference targetWrapper = TypeReference.of("org.springframework.data.projection.SpelEvaluatingMethodInterceptor$TargetWrapper");
 			assertThat(RuntimeHintsPredicates.reflection().onType(targetWrapper)
 					.withMemberCategories(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS, MemberCategory.INVOKE_DECLARED_METHODS,
@@ -342,7 +367,9 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 		@Test
 		void registerProxyForProjection() {
-			processController(ProjectionController.class);
+			processBeanClasses(ProjectionController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ProjectionController.class);
+			assertThatInvocationHintRegisteredForMethods(ProjectionController.class, "projection");
 			assertThat(RuntimeHintsPredicates.proxies().forInterfaces(BookProjection.class, TargetAware.class,
 					SpringProxy.class, DecoratingProxy.class)).accepts(generationContext.getRuntimeHints());
 		}
@@ -350,7 +377,9 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 		@Test
 		void registerProxyForOptionalProjection() {
-			processController(OptionalProjectionController.class);
+			processBeanClasses(OptionalProjectionController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(OptionalProjectionController.class);
+			assertThatInvocationHintRegisteredForMethods(OptionalProjectionController.class, "optionalProjection");
 			assertThat(RuntimeHintsPredicates.proxies().forInterfaces(BookProjection.class, TargetAware.class,
 					SpringProxy.class, DecoratingProxy.class)).accepts(generationContext.getRuntimeHints());
 		}
@@ -379,10 +408,46 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 
 	}
 
+	@Nested
+	class ExceptionHandlerTests {
 
-	private void processController(Class<?>... controllers) {
+		@Test
+		void registerReflectionOnControllerExceptionHandler() {
+			processBeanClasses(ExceptionController.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ExceptionController.class);
+			assertThatInvocationHintRegisteredForMethods(ExceptionController.class, "handleIllegalState");
+		}
+
+		@Test
+		void registerReflectionOnControllerAdviceExceptionHandler() {
+			processBeanClasses(ExceptionHandlers.class);
+			assertThatIntrospectionOnMethodsHintRegisteredForType(ExceptionHandlers.class);
+			assertThatInvocationHintRegisteredForMethods(ExceptionHandlers.class, "handleBindException");
+		}
+
+
+		@Controller
+		class ExceptionController {
+			@GraphQlExceptionHandler
+			public GraphQLError handleIllegalState(IllegalStateException exc) {
+				return null;
+			}
+		}
+
+		@ControllerAdvice
+		class ExceptionHandlers {
+
+			@GraphQlExceptionHandler
+			public GraphQLError handleBindException(BindException exc) {
+				return null;
+			}
+		}
+	}
+
+
+	private void processBeanClasses(Class<?>... beanClasses) {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
-		for (Class<?> beanClass : controllers) {
+		for (Class<?> beanClass : beanClasses) {
 			beanFactory.registerBeanDefinition(beanClass.getName(), new RootBeanDefinition(beanClass));
 		}
 		BeanFactoryInitializationAotContribution contribution = this.processor.processAheadOfTime(beanFactory);
@@ -390,6 +455,20 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 		contribution.applyTo(this.generationContext, mock(BeanFactoryInitializationCode.class));
 	}
 
+	private void assertThatIntrospectionOnMethodsHintRegisteredForType(Class<?> type) {
+		Predicate<RuntimeHints> predicate = RuntimeHintsPredicates.reflection()
+				.onType(type).withAnyMemberCategory(MemberCategory.INTROSPECT_DECLARED_METHODS);
+		assertThat(predicate).accepts(this.generationContext.getRuntimeHints());
+	}
+
+	private void assertThatInvocationHintRegisteredForMethods(Class<?> type, String... methodNames) {
+		Predicate<RuntimeHints> predicate = Arrays.stream(methodNames)
+				.map(methodName -> (Predicate<RuntimeHints>) RuntimeHintsPredicates.reflection().onMethod(type, methodName))
+				.reduce(Predicate::and)
+				.orElseThrow(() -> new IllegalArgumentException("Could not generate predicate on type " + type + " for methods " + Arrays.toString(methodNames)));
+		assertThat(predicate).accepts(this.generationContext.getRuntimeHints());
+	}
+
 	private void assertThatHintsForJavaBeanBindingRegisteredForTypes(Class<?>... types) {
 		Predicate<RuntimeHints> predicate = Arrays.stream(types)
 				.map(this::javaBeanBindingOnType)
@@ -406,7 +485,6 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 		assertThat(predicate).rejects(this.generationContext.getRuntimeHints());
 	}
 
-
 	private Predicate<RuntimeHints> javaBeanBindingOnType(Class<?> type) {
 		Predicate<RuntimeHints> predicate = RuntimeHintsPredicates.reflection().onType(type)
 				.withMemberCategories(MemberCategory.DECLARED_FIELDS, MemberCategory.INVOKE_DECLARED_CONSTRUCTORS);
@@ -430,4 +508,4 @@ class SchemaMappingBeanFactoryInitializationAotProcessorTests {
 		return predicate;
 	}
 
-}
\ No newline at end of file
+}