Consistent catch handling in InvocableHandlerMethodSupport

Closes gh-973

Author: rstoyanchev

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-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.
@@ -87,23 +87,13 @@ protected Object doInvoke(GraphQLContext graphQLContext, Object... argValues) {
 				return invokeSuspendingFunction(getBean(), method, argValues);
 			}
 			Object result = method.invoke(getBean(), argValues);
-			return handleReturnValue(graphQLContext, result);
+			return handleReturnValue(graphQLContext, result, method, argValues);
 		}
 		catch (IllegalArgumentException ex) {
-			assertTargetBean(method, getBean(), argValues);
-			String text = (ex.getMessage() != null) ? ex.getMessage() : "Illegal argument";
-			return Mono.error(new IllegalStateException(formatInvokeError(text, argValues), ex));
+			return Mono.error(processIllegalArgumentException(argValues, ex, method));
 		}
 		catch (InvocationTargetException ex) {
-			// Unwrap for DataFetcherExceptionResolvers ...
-			Throwable targetException = ex.getTargetException();
-			if (targetException instanceof Error || targetException instanceof Exception) {
-				return Mono.error(targetException);
-			}
-			else {
-				return Mono.error(new IllegalStateException(
-						formatInvokeError("Invocation failure", argValues), targetException));
-			}
+			return Mono.error(processInvocationTargetException(argValues, ex));
 		}
 		catch (Throwable ex) {
 			return Mono.error(ex);
@@ -124,24 +114,51 @@ private static Object invokeSuspendingFunction(Object bean, Method method, Objec
 	}
 
 	@Nullable
-	@SuppressWarnings("deprecation")
-	private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object result) {
+	@SuppressWarnings({"deprecation", "DataFlowIssue"})
+	private Object handleReturnValue(
+			GraphQLContext graphQLContext, @Nullable Object result, Method method, Object[] argValues) {
+
 		if (this.hasCallableReturnValue && result != null) {
-			return CompletableFuture.supplyAsync(
-					() -> {
-						try {
-							return ContextSnapshot.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
-						}
-						catch (Exception ex) {
-							throw new IllegalStateException(
-									"Failure in Callable returned from " + getBridgedMethod().toGenericString(), ex);
-						}
-					},
-					this.executor);
+			CompletableFuture<Object> future = new CompletableFuture<>();
+			this.executor.execute(() -> {
+				try {
+					ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
+					Object value = snapshot.wrap((Callable<?>) result).call();
+					future.complete(value);
+				}
+				catch (IllegalArgumentException ex) {
+					future.completeExceptionally(processIllegalArgumentException(argValues, ex, method));
+				}
+				catch (InvocationTargetException ex) {
+					future.completeExceptionally(processInvocationTargetException(argValues, ex));
+				}
+				catch (Exception ex) {
+					future.completeExceptionally(ex);
+				}
+			});
+			return future;
 		}
 		return result;
 	}
 
+	private IllegalStateException processIllegalArgumentException(
+			Object[] argValues, IllegalArgumentException ex, Method method) {
+
+		assertTargetBean(method, getBean(), argValues);
+		String text = (ex.getMessage() != null) ? ex.getMessage() : "Illegal argument";
+		return new IllegalStateException(formatInvokeError(text, argValues), ex);
+	}
+
+	private Throwable processInvocationTargetException(Object[] argValues, InvocationTargetException ex) {
+		// Unwrap for DataFetcherExceptionResolvers ...
+		Throwable targetException = ex.getTargetException();
+		if (targetException instanceof Error || targetException instanceof Exception) {
+			return targetException;
+		}
+		String message = formatInvokeError("Invocation failure", argValues);
+		return new IllegalStateException(message, targetException);
+	}
+
 	/**
 	 * Use this method to resolve the arguments asynchronously. This is only
 	 * useful when at least one of the values is a {@link Mono}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-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.
@@ -19,20 +19,20 @@
 
 import java.lang.reflect.Method;
 import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 
 import graphql.GraphQLContext;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.DataFetchingEnvironmentImpl;
 import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
 import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.graphql.data.GraphQlArgumentBinder;
 import org.springframework.graphql.data.method.HandlerMethod;
-import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.data.method.annotation.Argument;
 import org.springframework.graphql.data.method.annotation.QueryMapping;
@@ -74,14 +74,15 @@ void annotatedMethodsOnInterface() {
 	void callableReturnValue() throws Exception {
 
 		HandlerMethodArgumentResolverComposite resolvers = new HandlerMethodArgumentResolverComposite();
-		resolvers.addResolver(Mockito.mock(HandlerMethodArgumentResolver.class));
+		resolvers.addResolver(new ArgumentMethodArgumentResolver(new GraphQlArgumentBinder()));
 
 		DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
 				handlerMethodFor(new TestController(), "handleAndReturnCallable"), resolvers, null,
 				new SimpleAsyncTaskExecutor(), false);
 
 		DataFetchingEnvironment environment = DataFetchingEnvironmentImpl
 				.newDataFetchingEnvironment()
+				.arguments(Map.of("raiseError", false))
 				.graphQLContext(GraphQLContext.newContext().build())
 				.build();
 
@@ -92,6 +93,29 @@ void callableReturnValue() throws Exception {
 		assertThat(future.get()).isEqualTo("A");
 	}
 
+	@Test // gh-973
+	void callableReturnValueWithError() {
+
+		HandlerMethodArgumentResolverComposite resolvers = new HandlerMethodArgumentResolverComposite();
+		resolvers.addResolver(new ArgumentMethodArgumentResolver(new GraphQlArgumentBinder()));
+
+		DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
+				handlerMethodFor(new TestController(), "handleAndReturnCallable"), resolvers, null,
+				new SimpleAsyncTaskExecutor(), false);
+
+		DataFetchingEnvironment environment = DataFetchingEnvironmentImpl
+				.newDataFetchingEnvironment()
+				.arguments(Map.of("raiseError", true))
+				.graphQLContext(GraphQLContext.newContext().build())
+				.build();
+
+		Object result = handlerMethod.invoke(environment);
+
+		assertThat(result).isInstanceOf(CompletableFuture.class);
+		CompletableFuture<String> future = (CompletableFuture<String>) result;
+		StepVerifier.create(Mono.fromFuture(future)).expectErrorMessage("simulated exception").verify();
+	}
+
 	@Test
 	void completableFutureReturnValue() {
 
@@ -137,8 +161,13 @@ public String hello(String name) {
 		}
 
 		@Nullable
-		public Callable<String> handleAndReturnCallable() {
-			return () -> "A";
+		public Callable<String> handleAndReturnCallable(@Argument boolean raiseError) {
+			return () -> {
+				if (raiseError) {
+					throw new IllegalStateException("simulated exception");
+				}
+				return "A";
+			};
 		}
 
 		public CompletableFuture<String> handleAndReturnFuture(@AuthenticationPrincipal User user) {