Support annotated exception handler methods

See gh-160

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/GraphQlExceptionHandler.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2002-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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.data.method.annotation;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.List;
+
+/**
+ * Declares a method as a handler of exceptions raised while fetching data
+ * for a field. When declared in an
+ * {@link org.springframework.stereotype.Controller @Controller}, it applies to
+ * {@code @SchemaMapping} methods of that controller only. When declared in an
+ * {@link org.springframework.web.bind.annotation.ControllerAdvice @ControllerAdvice}
+ * it applies across controllers.
+ *
+ * <p>You can also use annotated exception handler methods in
+ * {@code @ControllerAdvice} beans to handle exceptions from non-controller
+ * {@link graphql.schema.DataFetcher}s by obtaining
+ * {@link org.springframework.graphql.data.method.annotation.support.AnnotatedControllerConfigurer#getExceptionResolver()}
+ * and registering it with
+ * {@link org.springframework.graphql.execution.GraphQlSource.Builder#exceptionResolvers(List)
+ * GraphQlSource.Builder}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.2
+ */
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface GraphQlExceptionHandler {
+
+	/**
+	 * Exceptions handled by the annotated method. If empty, defaults to
+	 * exception types declared in the method signature.
+	 */
+	Class<? extends Throwable>[] value() default {};
+
+}

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java

@@ -23,6 +23,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
@@ -30,6 +31,7 @@ import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
+import graphql.execution.DataFetcherResult;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.FieldCoordinates;
@@ -38,6 +40,7 @@ import graphql.schema.idl.RuntimeWiring;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.dataloader.DataLoader;
+import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -62,7 +65,9 @@ import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComp
 import org.springframework.graphql.data.method.annotation.BatchMapping;
 import org.springframework.graphql.data.method.annotation.SchemaMapping;
 import org.springframework.graphql.execution.BatchLoaderRegistry;
+import org.springframework.graphql.execution.DataFetcherExceptionResolver;
 import org.springframework.graphql.execution.RuntimeWiringConfigurer;
+import org.springframework.graphql.execution.SubscriptionPublisherException;
 import org.springframework.lang.Nullable;
 import org.springframework.stereotype.Controller;
 import org.springframework.util.Assert;
@@ -127,6 +132,9 @@ public class AnnotatedControllerConfigurer
 	@Nullable
 	private ValidationHelper validationHelper;
 
+	@Nullable
+	private AnnotatedControllerExceptionResolver exceptionResolver;
+
 
 	/**
 	 * Add a {@code FormatterRegistrar} to customize the {@link ConversionService}
@@ -165,6 +173,25 @@ public class AnnotatedControllerConfigurer
 		this.applicationContext = applicationContext;
 	}
 
+	/**
+	 * Return a {@link DataFetcherExceptionResolver} that resolves exceptions with
+	 * {@code @GraphQlExceptionHandler} methods in {@code @ControllerAdvice}
+	 * classes declared in Spring configuration. This is useful primarily for
+	 * exceptions from non-controller {@link DataFetcher}s since exceptions from
+	 * {@code @SchemaMapping} controller methods are handled automatically at
+	 * the point of invocation.
+	 *
+	 * @return a resolver instance that can be plugged into
+	 * {@link org.springframework.graphql.execution.GraphQlSource.Builder#exceptionResolvers(List)
+	 * GraphQlSource.Builder}
+	 *
+	 * @since 1.2
+	 */
+	public DataFetcherExceptionResolver getExceptionResolver() {
+		Assert.notNull(this.exceptionResolver, "ExceptionResolver is not initialized, was afterPropertiesSet called?");
+		return (ex, env) -> this.exceptionResolver.resolveException(ex, env, null);
+	}
+
 	@Nullable
 	HandlerMethodArgumentResolverComposite getArgumentResolvers() {
 		return this.argumentResolvers;
@@ -175,6 +202,11 @@ public class AnnotatedControllerConfigurer
 
 		this.argumentResolvers = initArgumentResolvers();
 
+		this.exceptionResolver = new AnnotatedControllerExceptionResolver(this.argumentResolvers);
+		if (this.applicationContext != null) {
+			this.exceptionResolver.registerControllerAdvice(this.applicationContext);
+		}
+
 		if (beanValidationPresent) {
 			this.validationHelper = ValidationHelper.createIfValidatorPresent(obtainApplicationContext());
 		}
@@ -222,12 +254,13 @@ public class AnnotatedControllerConfigurer
 	@Override
 	public void configure(RuntimeWiring.Builder runtimeWiringBuilder) {
 		Assert.state(this.argumentResolvers != null, "`argumentResolvers` is not initialized");
+		Assert.state(this.exceptionResolver != null, "`exceptionResolver` is not initialized");
 
 		findHandlerMethods().forEach((info) -> {
 			DataFetcher<?> dataFetcher;
 			if (!info.isBatchMapping()) {
 				dataFetcher = new SchemaMappingDataFetcher(
-						info, this.argumentResolvers, this.validationHelper, this.executor);
+						info, this.argumentResolvers, this.validationHelper, this.exceptionResolver, this.executor);
 			}
 			else {
 				String dataLoaderKey = registerBatchLoader(info);
@@ -493,19 +526,30 @@ public class AnnotatedControllerConfigurer
 		@Nullable
 		private final Consumer<Object[]> methodValidationHelper;
 
+		private final AnnotatedControllerExceptionResolver exceptionResolver;
+
 		@Nullable
 		private final Executor executor;
 
 		private final boolean subscription;
 
 		SchemaMappingDataFetcher(
-				MappingInfo info, HandlerMethodArgumentResolverComposite resolvers,
-				@Nullable ValidationHelper validationHelper, @Nullable Executor executor) {
+				MappingInfo info, HandlerMethodArgumentResolverComposite argumentResolvers,
+				@Nullable ValidationHelper helper, AnnotatedControllerExceptionResolver exceptionResolver,
+				@Nullable Executor executor) {
 
 			this.info = info;
-			this.argumentResolvers = resolvers;
-			this.methodValidationHelper = (validationHelper != null ?
-					validationHelper.getValidationHelperFor(info.getHandlerMethod()) : null);
+			this.argumentResolvers = argumentResolvers;
+
+			this.methodValidationHelper =
+					(helper != null ? helper.getValidationHelperFor(info.getHandlerMethod()) : null);
+
+			// Register controllers early to validate exception handler return types
+			Class<?> controllerType = info.getHandlerMethod().getBeanType();
+			exceptionResolver.registerController(controllerType);
+
+			this.exceptionResolver = exceptionResolver;
+
 			this.executor = executor;
 			this.subscription = this.info.getCoordinates().getTypeName().equalsIgnoreCase("Subscription");
 		}
@@ -517,17 +561,53 @@ public class AnnotatedControllerConfigurer
 			return this.info.getHandlerMethod();
 		}
 
-
 		@Override
-		@SuppressWarnings("ConstantConditions")
+		@SuppressWarnings({"ConstantConditions", "ReactiveStreamsUnusedPublisher"})
 		public Object get(DataFetchingEnvironment environment) throws Exception {
 
 			DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
 					getHandlerMethod(), this.argumentResolvers, this.methodValidationHelper,
 					this.executor, this.subscription);
 
-			return handlerMethod.invoke(environment);
+			try {
+				Object result = handlerMethod.invoke(environment);
+				return applyExceptionHandling(environment, handlerMethod, result);
+			}
+			catch (Throwable ex) {
+				return handleException(ex, environment, handlerMethod);
+			}
 		}
+
+		@SuppressWarnings({"unchecked", "ReactiveStreamsUnusedPublisher"})
+		private <T> Object applyExceptionHandling(
+				DataFetchingEnvironment env, DataFetcherHandlerMethod handlerMethod, Object result) {
+
+			if (this.subscription && result instanceof Publisher<?> publisher) {
+				result = Flux.from(publisher).onErrorResume(ex -> handleSubscriptionError(ex, env, handlerMethod));
+			}
+			else if (result instanceof Mono) {
+				result = ((Mono<T>) result).onErrorResume(ex -> (Mono<T>) handleException(ex, env, handlerMethod));
+			}
+			else if (result instanceof Flux<?>) {
+				result = ((Flux<T>) result).onErrorResume(ex -> (Mono<T>) handleException(ex, env, handlerMethod));
+			}
+			return result;
+		}
+
+		private Mono<DataFetcherResult<?>> handleException(
+				Throwable ex, DataFetchingEnvironment env, DataFetcherHandlerMethod handlerMethod) {
+
+			return this.exceptionResolver.resolveException(ex, env, handlerMethod.getBean())
+					.map(errors -> DataFetcherResult.newResult().errors(errors).build());
+		}
+
+		private <T> Publisher<T> handleSubscriptionError(
+				Throwable ex, DataFetchingEnvironment env, DataFetcherHandlerMethod handlerMethod) {
+
+			return this.exceptionResolver.resolveException(ex, env, handlerMethod.getBean())
+					.flatMap(errors -> Mono.error(new SubscriptionPublisherException(errors, ex)));
+		}
+
 	}
 
 

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerExceptionResolver.java

@@ -0,0 +1,450 @@
+/*
+ * Copyright 2002-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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.data.method.annotation.support;
+
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import graphql.GraphQLError;
+import graphql.schema.DataFetchingEnvironment;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import reactor.core.publisher.Mono;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.core.ExceptionDepthComparator;
+import org.springframework.core.MethodIntrospector;
+import org.springframework.core.MethodParameter;
+import org.springframework.core.annotation.AnnotatedElementUtils;
+import org.springframework.graphql.data.method.HandlerMethod;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
+import org.springframework.graphql.data.method.annotation.GraphQlExceptionHandler;
+import org.springframework.graphql.execution.DataFetcherExceptionResolver;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+import org.springframework.util.ConcurrentReferenceHashMap;
+import org.springframework.util.ObjectUtils;
+import org.springframework.util.ReflectionUtils;
+import org.springframework.web.method.ControllerAdviceBean;
+
+/**
+ * Resolves exceptions via {@link GraphQlExceptionHandler @GraphQlExceptionHandler}
+ * handler methods, which can be either local to a controller, or applicable
+ * across controllers and {@link graphql.schema.DataFetcher}s when declared in
+ * an {@link org.springframework.web.bind.annotation.ControllerAdvice} bean.
+ *
+ * <p>This {@link #resolveException(Throwable, DataFetchingEnvironment, Object)}
+ * method is similar to the {@link DataFetcherExceptionResolver} contract except
+ * it takes an additional, optional argument with the controller that raised the
+ * exception, for finding exception handler methods relative to the controller.
+ *
+ * <p>{@code AnnotatedControllerExceptionResolver} is package private and
+ * automatically applied from {@link AnnotatedControllerConfigurer} to controller
+ * method invocations. In addition, you can access it as a
+ * {@link DataFetcherExceptionResolver} via
+ * {@link AnnotatedControllerConfigurer#getExceptionResolver()} to extend
+ * exception handling with {@code @ControllerAdvice} exception handlers to
+ * non-controller {@link graphql.schema.DataFetcher}s.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.2
+ */
+final class AnnotatedControllerExceptionResolver {
+
+	private static final Log logger = LogFactory.getLog(AnnotatedControllerExceptionResolver.class);
+
+
+	private final HandlerMethodArgumentResolverComposite argumentResolvers;
+
+	private final Map<Class<?>, MethodResolver> controllerCache = new ConcurrentHashMap<>(64);
+
+	private final Map<ControllerAdviceBean, MethodResolver> controllerAdviceCache = new ConcurrentHashMap<>(64);
+
+
+	AnnotatedControllerExceptionResolver(HandlerMethodArgumentResolverComposite resolvers) {
+		Assert.notNull(resolvers, "'resolvers' are required");
+		this.argumentResolvers = resolvers;
+	}
+
+
+	/**
+	 * Detect {@link GraphQlExceptionHandler} methods in the given controller
+	 * class, and save this information for use at runtime. Method return types
+	 * are validated to ensure they are within a range of supported types.
+	 * @param controllerType the controller type to register
+	 */
+	public void registerController(Class<?> controllerType) {
+		this.controllerCache.computeIfAbsent(
+				controllerType, type -> new MethodResolver(findExceptionHandlers(controllerType)));
+	}
+
+	/**
+	 * Find {@link org.springframework.web.bind.annotation.ControllerAdvice}
+	 * beans in the given {@code ApplicationContext}, and detect
+	 * {@link GraphQlExceptionHandler} methods in them, saving this information
+	 * for use at runtime.
+	 * @param context the context to look into
+	 */
+	public void registerControllerAdvice(ApplicationContext context) {
+		for (ControllerAdviceBean bean : ControllerAdviceBean.findAnnotatedBeans(context)) {
+			Class<?> beanType = bean.getBeanType();
+			if (beanType != null) {
+				Map<Class<? extends Throwable>, Method> methods = findExceptionHandlers(beanType);
+				if (!methods.isEmpty()) {
+					this.controllerAdviceCache.put(bean, new MethodResolver(methods));
+				}
+			}
+		}
+		if (logger.isDebugEnabled()) {
+			logger.debug("@GraphQlException methods in ControllerAdvice beans: " +
+					(this.controllerAdviceCache.size() == 0 ? "none" : this.controllerAdviceCache.size()));
+		}
+	}
+
+	@SuppressWarnings("unchecked")
+	private static Map<Class<? extends Throwable>, Method> findExceptionHandlers(Class<?> handlerType) {
+
+		Map<Method, GraphQlExceptionHandler> handlerMap = MethodIntrospector.selectMethods(
+				handlerType, (MethodIntrospector.MetadataLookup<GraphQlExceptionHandler>) method ->
+						AnnotatedElementUtils.findMergedAnnotation(method, GraphQlExceptionHandler.class));
+
+		Map<Class<? extends Throwable>, Method> mappings = new HashMap<>(handlerMap.size());
+		handlerMap.forEach((method, annotation) -> {
+			List<Class<? extends Throwable>> exceptionTypes = new ArrayList<>();
+			if (!ObjectUtils.isEmpty(annotation.value())) {
+				exceptionTypes.addAll(Arrays.asList(annotation.value()));
+			}
+			else {
+				for (Class<?> parameterType : method.getParameterTypes()) {
+					if (Throwable.class.isAssignableFrom(parameterType)) {
+						exceptionTypes.add((Class<? extends Throwable>) parameterType);
+					}
+				}
+			}
+			Assert.state(!exceptionTypes.isEmpty(), () -> "No exception types for " + method);
+			for (Class<? extends Throwable> type : exceptionTypes) {
+				Method oldMethod = mappings.put(type, method);
+				Assert.state(oldMethod == null || oldMethod.equals(method), () ->
+						"Ambiguous @GraphQlExceptionHandler for [" + type + "]: {" + oldMethod + ", " + method + "}");
+			}
+		});
+		return mappings;
+	}
+
+
+	/**
+	 * Resolve the exception with an {@code @GraphQlExceptionHandler} method.
+	 * If a controller is provided, look for a matching exception handler in the
+	 * controller first, and then in any applicable {@code @ControllerAdvice}.
+	 * If a controller is not provided, look in all {@code @ControllerAdvice}.
+	 * @param ex the exception to resolve
+	 * @param environment the environment for the invoked {@code DataFetcher}
+	 * @param controller the controller that raised the exception, if applicable
+	 * @return a {@code Mono} with errors as specified in
+	 * {@link DataFetcherExceptionResolver#resolveException(Throwable, DataFetchingEnvironment)}
+	 */
+	public Mono<List<GraphQLError>> resolveException(
+			Throwable ex, DataFetchingEnvironment environment, @Nullable Object controller) {
+
+		Object controllerOrAdvice = null;
+		MethodHolder methodHolder = null;
+
+		if (controller != null) {
+			MethodResolver methodResolver = this.controllerCache.get(controller.getClass());
+			if (methodResolver != null) {
+				controllerOrAdvice = controller;
+				methodHolder = methodResolver.resolveMethod(ex);
+			}
+			else if (logger.isWarnEnabled()) {
+				logger.warn("No registration for controller type: " + controller.getClass().getName());
+			}
+		}
+
+		if (methodHolder == null) {
+			for (Map.Entry<ControllerAdviceBean, MethodResolver> entry : this.controllerAdviceCache.entrySet()) {
+				ControllerAdviceBean advice = entry.getKey();
+				if (controller == null || advice.isApplicableToBeanType(controller.getClass())) {
+					methodHolder = entry.getValue().resolveMethod(ex);
+					if (methodHolder != null) {
+						controllerOrAdvice = advice.resolveBean();
+						break;
+					}
+				}
+			}
+		}
+
+		if (methodHolder == null) {
+			return Mono.error(ex);
+		}
+
+		return invokeExceptionHandler(ex, environment, controllerOrAdvice, methodHolder);
+	}
+
+	private Mono<List<GraphQLError>> invokeExceptionHandler(
+			Throwable exception, DataFetchingEnvironment env, Object controllerOrAdvice, MethodHolder methodHolder) {
+
+		DataFetcherHandlerMethod exceptionHandler = new DataFetcherHandlerMethod(
+				new HandlerMethod(controllerOrAdvice, methodHolder.getMethod()), this.argumentResolvers,
+				null, null, false);
+
+		List<Throwable> exceptions = new ArrayList<>();
+		try {
+			if (logger.isDebugEnabled()) {
+				logger.debug("Handling exception with " + exceptionHandler);
+			}
+
+			// Expose causes as provided arguments as well
+			Throwable exToExpose = exception;
+			while (exToExpose != null) {
+				exceptions.add(exToExpose);
+				Throwable cause = exToExpose.getCause();
+				exToExpose = (cause != exToExpose ? cause : null);
+			}
+			Object[] arguments = new Object[exceptions.size() + 1];
+			exceptions.toArray(arguments);  // efficient arraycopy call in ArrayList
+			arguments[arguments.length - 1] = exceptionHandler;
+
+			Object result = exceptionHandler.invoke(env, arguments);
+
+			return methodHolder.adapt(result, exception);
+		}
+		catch (Throwable invocationEx) {
+			// Any other than the original exception (or a cause) is unintended here,
+			// probably an accident (e.g. failed assertion or the like).
+			if (!exceptions.contains(invocationEx) && logger.isWarnEnabled()) {
+				logger.warn("Failure while handling exception with " + exceptionHandler, invocationEx);
+			}
+			// Continue with processing of the original exception...
+			return Mono.error(exception);
+		}
+	}
+
+
+	/**
+	 * Helps to resolve Exception instances to handler methods.
+	 */
+	private static final class MethodResolver {
+
+		@SuppressWarnings("DataFlowIssue")
+		private static final MethodHolder NO_MATCH =
+				new MethodHolder(ReflectionUtils.findMethod(MethodResolver.class, "noMatch"));
+
+
+		private final Map<Class<? extends Throwable>, MethodHolder> exceptionMappings = new HashMap<>(16);
+
+		private final Map<Class<? extends Throwable>, MethodHolder> resolvedExceptionCache = new ConcurrentReferenceHashMap<>(16);
+
+		MethodResolver(Map<Class<? extends Throwable>, Method> methodMap) {
+			methodMap.forEach((exceptionType, method) ->
+					this.exceptionMappings.put(exceptionType, new MethodHolder(method)));
+		}
+
+		/**
+		 * Find an exception handler method mapped to the given exception, using
+		 * {@link ExceptionDepthComparator} if more than one match is found.
+		 * @param exception the exception
+		 * @return the exception handler to use, or {@code null} if no match
+		 */
+		@Nullable
+		public MethodHolder resolveMethod(Throwable exception) {
+			MethodHolder method = resolveMethodByExceptionType(exception.getClass());
+			if (method == null) {
+				Throwable cause = exception.getCause();
+				if (cause != null) {
+					method = resolveMethod(cause);
+				}
+			}
+			return method;
+		}
+
+		@Nullable
+		private MethodHolder resolveMethodByExceptionType(Class<? extends Throwable> exceptionType) {
+			MethodHolder method = this.resolvedExceptionCache.get(exceptionType);
+			if (method == null) {
+				method = getMappedMethod(exceptionType);
+				this.resolvedExceptionCache.put(exceptionType, method);
+			}
+			return (method != NO_MATCH ? method : null);
+		}
+
+		private MethodHolder getMappedMethod(Class<? extends Throwable> exceptionType) {
+			List<Class<? extends Throwable>> matches = new ArrayList<>();
+			for (Class<? extends Throwable> mappedException : this.exceptionMappings.keySet()) {
+				if (mappedException.isAssignableFrom(exceptionType)) {
+					matches.add(mappedException);
+				}
+			}
+			if (!matches.isEmpty()) {
+				if (matches.size() > 1) {
+					matches.sort(new ExceptionDepthComparator(exceptionType));
+				}
+				return this.exceptionMappings.get(matches.get(0));
+			}
+			else {
+				return NO_MATCH;
+			}
+		}
+
+		@SuppressWarnings("unused")
+		private void noMatch() {
+		}
+
+	}
+
+
+	/**
+	 * Container for an exception handler method, and an adapter for its return values.
+	 */
+	private static class MethodHolder {
+
+		private final Method method;
+
+		private final MethodParameter returnType;
+
+		private final ReturnValueAdapter adapter;
+
+		MethodHolder(Method method) {
+			Assert.notNull(method, "Method is required");
+			this.method = method;
+			this.returnType = new MethodParameter(method, -1);
+			this.adapter = ReturnValueAdapter.createFor(this.returnType);
+		}
+
+		public Method getMethod() {
+			return this.method;
+		}
+
+		public Mono<List<GraphQLError>> adapt(@Nullable Object result, Throwable ex) {
+			return this.adapter.adapt(result, this.returnType, ex);
+		}
+
+	}
+
+
+	/**
+	 * Contract to adapt the value returned from a {@code @GraphQlExceptionHandler}.
+	 */
+	@SuppressWarnings("unchecked")
+	private interface ReturnValueAdapter {
+
+		/**
+		 * Adapt the given return value to {@code Mono<List<GraphQLError>>}.
+		 * @param result the return value
+		 * @param returnType the return type of the method, mainly used for error logging
+		 * @param ex the exception being handled
+		 * @return the adapted result according to the contact for
+		 * {@link DataFetcherExceptionResolver#resolveException(Throwable, DataFetchingEnvironment)}
+		 */
+		Mono<List<GraphQLError>> adapt(@Nullable Object result, MethodParameter returnType, Throwable ex);
+
+		/**
+		 * Verify the method return type is supported and can be adapted to
+		 * {@code Mono<List<GraphQLError>>}, and create a suitable adapter.
+		 * @param returnType the return type of the method
+		 * @return the chosen adapter
+		 * @throws IllegalStateException if the return value type that cannot be
+		 * adapted to {@code Mono<List<GraphQLError>>} and is not supported
+		 */
+		static ReturnValueAdapter createFor(MethodParameter returnType) {
+			Class<?> parameterType = returnType.getParameterType();
+			if (parameterType == void.class || parameterType == Void.class) {
+				return forVoid;
+			}
+			else if (parameterType.equals(GraphQLError.class)) {
+				return forSingleError;
+			}
+			else if (Collection.class.isAssignableFrom(parameterType)) {
+				if (returnType.nested().getNestedParameterType().equals(GraphQLError.class)) {
+					return forCollection;
+				}
+			}
+			else if (Mono.class.isAssignableFrom(parameterType)) {
+				returnType = returnType.nested();
+				Class<?> nestedType = returnType.getNestedParameterType();
+				if (nestedType == void.class || nestedType == Void.class) {
+					return forMonoVoid;
+				}
+				if (Collection.class.isAssignableFrom(nestedType)) {
+					returnType = returnType.nested();
+					nestedType = returnType.getNestedParameterType();
+				}
+				if (nestedType.equals(GraphQLError.class) || nestedType.equals(Object.class)) {
+					return forMono;
+				}
+			}
+			else if (parameterType.equals(Object.class)) {
+				return forObject;
+			}
+			throw new IllegalStateException(
+					"Invalid return type for @GraphQlExceptionHandler method: " + returnType);
+		}
+
+		/** Adapter for void */
+		ReturnValueAdapter forVoid = (result, returnType, ex) -> Mono.just(Collections.emptyList());
+
+		/** Adapter for a single GraphQLError */
+		ReturnValueAdapter forSingleError = (result, returnType, ex) ->
+				(result == null ?
+						Mono.error(ex) :
+						Mono.just(Collections.singletonList((GraphQLError) result)));
+
+		/** Adapter for a collection of GraphQLError's */
+		ReturnValueAdapter forCollection = (result, returnType, ex) ->
+				(result == null ?
+						Mono.error(ex) :
+						Mono.just((result instanceof List ?
+								(List<GraphQLError>) result :
+								new ArrayList<>((Collection<GraphQLError>) result))));
+
+		/** Adapter for Object */
+		ReturnValueAdapter forObject = (result, returnType, ex) -> {
+			if (result == null) {
+				return Mono.error(ex);
+			}
+			else if (result instanceof GraphQLError) {
+				return forSingleError.adapt(result, returnType, ex);
+			}
+			else if (result instanceof Collection<?>) {
+				return forCollection.adapt(result, returnType, ex);
+			}
+			else {
+				if (logger.isWarnEnabled()) {
+					logger.warn("Unexpected return value of type " +
+							result.getClass().getName() + " from method " + returnType);
+				}
+				return Mono.error(ex);
+			}
+		};
+
+		/** Adapter for {@code Mono<Void>} */
+		ReturnValueAdapter forMonoVoid = (result, returnType, ex) ->
+				(result == null ? Mono.error(ex) : Mono.just(Collections.emptyList()));
+
+		/** Adapter for a {@code Mono} wrapping any of the other synchronous return value types */
+		ReturnValueAdapter forMono = (result, returnType, ex) ->
+				(result == null ?
+						Mono.error(ex) :
+						((Mono<?>) result).flatMap(o -> forObject.adapt(o, returnType, ex)).switchIfEmpty(Mono.error(ex)));
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java

@@ -84,7 +84,6 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	}
 
 
-
 	/**
 	 * Invoke the method after resolving its argument values in the context of
 	 * the given {@link DataFetchingEnvironment}.
@@ -95,8 +94,7 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	 * be used directly, i.e. without argument resolution. Provided argument
 	 * values are checked before argument resolvers.
 	 *
-	 * @param environment the GraphQL {@link DataFetchingEnvironment} to use to
-	 * resolve arguments.
+	 * @param environment the environment to resolve arguments from
 	 *
 	 * @return the raw value returned by the invoked method, possibly a
 	 * {@code Mono} in case a method argument requires asynchronous resolution;
@@ -104,9 +102,19 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	 */
 	@Nullable
 	public Object invoke(DataFetchingEnvironment environment) {
+		return invoke(environment, new Object[0]);
+	}
+
+	/**
+	 * Variant of {@link #invoke(DataFetchingEnvironment)} that also accepts
+	 * "given" arguments, which are matched by type.
+	 * @since 1.2
+	 */
+	@Nullable
+	public Object invoke(DataFetchingEnvironment environment, Object... providedArgs) {
 		Object[] args;
 		try {
-			args = getMethodArgumentValues(environment);
+			args = getMethodArgumentValues(environment, providedArgs);
 		}
 		catch (Throwable ex) {
 			return Mono.error(ex);

spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java

@@ -77,9 +77,14 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
 
 		if (this.subscription) {
 			Assert.state(value instanceof Publisher, "Expected Publisher for a subscription");
-			Flux<?> flux = Flux.from((Publisher<?>) value).onErrorResume(exception ->
-					this.subscriptionExceptionResolver.resolveException(exception)
-							.flatMap(errors -> Mono.error(new SubscriptionPublisherException(errors, exception))));
+			Flux<?> flux = Flux.from((Publisher<?>) value).onErrorResume(exception -> {
+				// Already handled, e.g. controller methods?
+				if (exception instanceof SubscriptionPublisherException) {
+					return Mono.error(exception);
+				}
+				return this.subscriptionExceptionResolver.resolveException(exception)
+						.flatMap(errors -> Mono.error(new SubscriptionPublisherException(errors, exception)));
+			});
 			return flux.contextWrite(snapshot::updateContext);
 		}
 

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerExceptionResolverTests.java

@@ -0,0 +1,232 @@
+/*
+ * Copyright 2002-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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.data.method.annotation.support;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import graphql.GraphQLError;
+import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DataFetchingEnvironmentImpl;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.support.StaticApplicationContext;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
+import org.springframework.graphql.data.method.annotation.GraphQlExceptionHandler;
+import org.springframework.lang.Nullable;
+import org.springframework.stereotype.Controller;
+import org.springframework.util.StringUtils;
+import org.springframework.web.bind.annotation.ControllerAdvice;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
+
+/**
+ * Unit tests for {@link AnnotatedControllerExceptionResolver}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.2
+ */
+public class AnnotatedControllerExceptionResolverTests {
+
+	private final DataFetchingEnvironment environment =
+			DataFetchingEnvironmentImpl.newDataFetchingEnvironment().build();
+
+
+	@Test
+	void resolveToSingleError() {
+		Exception ex = new IllegalArgumentException("Bad input");
+		testResolve(ex, new TestController(), Collections.singletonList("handleToSingleError: " + ex.getMessage()));
+	}
+
+	@Test
+	void resolveToList() {
+		Exception ex = new IllegalAccessException("No access");
+		testResolve(ex, new TestController(), Arrays.asList(
+				"handleToList[1]: " + ex.getMessage(), "handleToList[2]: " + ex.getMessage()));
+	}
+
+	@Test
+	void resolveToMono() {
+		Exception ex = new InstantiationException("Failed to instantiate");
+		testResolve(ex, new TestController(), Collections.singletonList("handleToMono: " + ex.getMessage()));
+	}
+
+	@Test
+	void resolveToObject() {
+		Exception ex = new ClassCastException("Wrong type");
+		testResolve(ex, new TestController(), Collections.singletonList("handleToObject: " + ex.getMessage()));
+	}
+
+	@Test
+	void resolveToVoid() {
+		Exception ex = new ArithmeticException();
+		testResolve(ex, new TestController(), Collections.emptyList());
+	}
+
+	@Test
+	void resolveTypeDeclaredOnAnnotation() {
+		Exception ex = new SecurityException();
+		testResolve(ex, new TestController(), Collections.singletonList("handleWithTypeOnAnnotation"));
+	}
+
+	@Test
+	void resolveFromRootCause() {
+		Exception ex = new Exception("A", new Exception("B", new IndexOutOfBoundsException(5)));
+		testResolve(ex, new TestController(), Collections.singletonList("handleRootCause: Index out of range: 5"));
+	}
+
+	@Test
+	void leaveUnresolvedViaNullReturnValue() {
+		Exception ex = new ClassNotFoundException("Not found");
+
+		TestController controller = new TestController();
+		AnnotatedControllerExceptionResolver resolver = exceptionResolver();
+		resolver.registerController(controller.getClass());
+
+		StepVerifier.create(resolver.resolveException(ex, this.environment, controller))
+				.expectErrorSatisfies(actualEx -> assertThat(actualEx).isSameAs(ex))
+				.verify();
+	}
+
+	@Test
+	void resolveWithControllerAdvice() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(TestControllerAdvice.class);
+		context.refresh();
+
+		Exception ex = new IllegalArgumentException("Bad input");
+		List<GraphQLError> actual = exceptionResolver(context).resolveException(ex, this.environment, null).block();
+
+		assertThat(actual).hasSize(1);
+		assertThat(actual.get(0).getMessage()).isEqualTo("handle: Bad input");
+	}
+
+	@Test
+	void invalidReturnType() {
+		assertThatIllegalStateException().isThrownBy(() ->
+				exceptionResolver().registerController(InvalidReturnTypeController.class));
+	}
+
+	private void testResolve(Throwable ex, TestController controller, List<String> expected) {
+
+		AnnotatedControllerExceptionResolver resolver = exceptionResolver();
+		resolver.registerController(controller.getClass());
+
+		List<GraphQLError> actual = resolver.resolveException(ex, this.environment, controller).block();
+
+		assertThat(actual).hasSize(expected.size());
+		for (int i = 0; i < expected.size(); i++) {
+			assertThat(actual.get(i).getMessage()).isEqualTo(expected.get(i));
+		}
+	}
+
+	private AnnotatedControllerExceptionResolver exceptionResolver() {
+		return exceptionResolver(new StaticApplicationContext());
+	}
+
+	private AnnotatedControllerExceptionResolver exceptionResolver(ApplicationContext applicationContext) {
+		AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
+		configurer.setApplicationContext(applicationContext);
+		configurer.afterPropertiesSet();
+
+		HandlerMethodArgumentResolverComposite argumentResolvers = configurer.getArgumentResolvers();
+		AnnotatedControllerExceptionResolver resolver = new AnnotatedControllerExceptionResolver(argumentResolvers);
+		resolver.registerControllerAdvice(applicationContext);
+		return resolver;
+	}
+
+
+	@SuppressWarnings("unused")
+	@Controller
+	private static class TestController {
+
+		@GraphQlExceptionHandler
+		GraphQLError handleToSingleError(IllegalArgumentException ex) {
+			return createError("handleToSingleError", ex);
+		}
+
+		@GraphQlExceptionHandler
+		List<GraphQLError> handleToList(IllegalAccessException ex) {
+			return Arrays.asList(createError("handleToList[1]", ex), createError("handleToList[2]", ex));
+		}
+
+		@GraphQlExceptionHandler
+		Mono<GraphQLError> handleToMono(InstantiationException ex) {
+			return Mono.just(createError("handleToMono", ex));
+		}
+
+		@GraphQlExceptionHandler
+		Object handleToObject(ClassCastException ex) {
+			return createError("handleToObject", ex);
+		}
+
+		@GraphQlExceptionHandler(ArithmeticException.class)
+		public void handleToVoid() {
+		}
+
+		@Nullable
+		@GraphQlExceptionHandler
+		GraphQLError handleAndLeaveNotHandled(ClassNotFoundException ex) {
+			return null;
+		}
+
+		@GraphQlExceptionHandler
+		public GraphQLError handleRootCause(IndexOutOfBoundsException ex) {
+			return createError("handleRootCause", ex);
+		}
+
+		@GraphQlExceptionHandler(SecurityException.class)
+		public GraphQLError handleWithTypeOnAnnotation() {
+			return createError("handleWithTypeOnAnnotation", null);
+		}
+
+		private static GraphQLError createError(String methodName, @Nullable Throwable ex) {
+			return GraphQLError.newError()
+					.message(methodName + (ex != null && StringUtils.hasText(ex.getMessage()) ? ": " + ex.getMessage() : ""))
+					.build();
+		}
+
+	}
+
+
+	@SuppressWarnings("unused")
+	@ControllerAdvice
+	private static class TestControllerAdvice {
+
+		@GraphQlExceptionHandler
+		GraphQLError handle(IllegalArgumentException ex) {
+			return GraphQLError.newError().message("handle: " + ex.getMessage()).build();
+		}
+
+	}
+
+
+	private static class InvalidReturnTypeController {
+
+		@GraphQlExceptionHandler
+		public String handle(IllegalArgumentException ex) {
+			return "Handled";
+		}
+
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingInvocationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-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.
@@ -21,6 +21,7 @@ import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 
 import graphql.GraphQLContext;
+import graphql.GraphQLError;
 import graphql.schema.DataFetchingEnvironment;
 import org.dataloader.DataLoader;
 import org.junit.jupiter.api.Test;
@@ -42,13 +43,17 @@ import org.springframework.graphql.GraphQlSetup;
 import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.TestExecutionRequest;
 import org.springframework.graphql.data.method.annotation.Argument;
+import org.springframework.graphql.data.method.annotation.GraphQlExceptionHandler;
 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.graphql.data.method.annotation.SubscriptionMapping;
 import org.springframework.graphql.execution.BatchLoaderRegistry;
 import org.springframework.graphql.execution.DefaultBatchLoaderRegistry;
+import org.springframework.graphql.execution.ErrorType;
+import org.springframework.graphql.execution.SubscriptionPublisherException;
 import org.springframework.stereotype.Controller;
+import org.springframework.util.Assert;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -206,6 +211,50 @@ public class SchemaMappingInvocationTests {
 				.verifyComplete();
 	}
 
+	@Test
+	void handleExceptionFromQuery() {
+		String document = "{ " +
+				"  booksByCriteria(criteria: {author:\"Fitzgerald\"}) { " +
+				"    id" +
+				"    name" +
+				"  }" +
+				"}";
+
+		Mono<ExecutionGraphQlResponse> responseMono =
+				graphQlService().execute(TestExecutionRequest.forDocument(document));
+
+		ResponseHelper responseHelper = ResponseHelper.forResponse(responseMono);
+		assertThat(responseHelper.errorCount()).isEqualTo(1);
+		assertThat(responseHelper.error(0).errorType()).isEqualTo("BAD_REQUEST");
+		assertThat(responseHelper.error(0).message()).isEqualTo("Rejected: Bad input");
+	}
+
+	@Test
+	void handleExceptionFromSubscription() {
+		String document = "subscription { " +
+				"  bookSearch(author:\"Fitzgerald\") { " +
+				"    id" +
+				"    name" +
+				"  }" +
+				"}";
+
+		Mono<ExecutionGraphQlResponse> responseMono =
+				graphQlService().execute(TestExecutionRequest.forDocument(document));
+
+		Flux<Book> bookFlux = ResponseHelper.forSubscription(responseMono)
+				.map(response -> response.toEntity("bookSearch", Book.class));
+
+		StepVerifier.create(bookFlux)
+				.expectErrorSatisfies(ex -> {
+					SubscriptionPublisherException theEx = (SubscriptionPublisherException) ex;
+					List<GraphQLError> errors = theEx.getErrors();
+					assertThat(errors).hasSize(1);
+					assertThat(errors.get(0).getErrorType().toString()).isEqualTo("BAD_REQUEST");
+					assertThat(errors.get(0).getMessage()).isEqualTo("Rejected: Bad input");
+				})
+				.verify();
+	}
+
 
 	private ExecutionGraphQlService graphQlService() {
 		BatchLoaderRegistry registry = new DefaultBatchLoaderRegistry();
@@ -243,6 +292,7 @@ public class SchemaMappingInvocationTests {
 
 		@QueryMapping
 		public List<Book> booksByCriteria(@Argument BookCriteria criteria) {
+			Assert.isTrue(!criteria.getAuthor().equalsIgnoreCase("Fitzgerald"), "Bad input");
 			return BookSource.findBooksByAuthor(criteria.getAuthor());
 		}
 
@@ -277,7 +327,16 @@ public class SchemaMappingInvocationTests {
 
 		@SubscriptionMapping
 		public Flux<Book> bookSearch(@Argument String author) {
-			return Flux.fromIterable(BookSource.findBooksByAuthor(author));
+			return (author.equalsIgnoreCase("Fitzgerald") ?
+					Flux.error(new IllegalArgumentException("Bad input")) :
+					Flux.fromIterable(BookSource.findBooksByAuthor(author)));
+		}
+
+		@GraphQlExceptionHandler
+		public GraphQLError handleInputError(IllegalArgumentException ex) {
+			return GraphQLError.newError().errorType(ErrorType.BAD_REQUEST)
+					.message("Rejected: " + ex.getMessage())
+					.build();
 		}
 	}