Polishing in InvocableHandlerMethodSupport

Author: rstoyanchev

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

@@ -60,10 +60,12 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
  */
 protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod, @Nullable Executor executor) {
 super(handlerMethod.createWithResolvedBean());
+
 this.hasCallableReturnValue = getReturnType().getParameterType().equals(Callable.class);
 this.executor = executor;
-Assert.isTrue(!this.hasCallableReturnValue || this.executor != null,
-"Controller method declared with Callable return value, but no Executor configured: " +
+
+Assert.isTrue(!this.hasCallableReturnValue || executor != null,
+"Controller method has Callable return value, but Executor not provided: " +
 handlerMethod.getBridgedMethod().toGenericString());
 }
 
@@ -86,8 +88,14 @@ protected Object doInvoke(GraphQLContext graphQLContext, Object... argValues) {
 if (KotlinDetector.isSuspendingFunction(method)) {
 return invokeSuspendingFunction(getBean(), method, argValues);
 }
+
 Object result = method.invoke(getBean(), argValues);
-return handleReturnValue(graphQLContext, result);
+
+if (this.hasCallableReturnValue && result != null) {
+result = adaptCallable(graphQLContext, (Callable<?>) result);
+}
+
+return result;
 }
 catch (IllegalArgumentException ex) {
 assertTargetBean(method, getBean(), argValues);
@@ -114,7 +122,7 @@ protected Object doInvoke(GraphQLContext graphQLContext, Object... argValues) {
 private static Object invokeSuspendingFunction(Object bean, Method method, Object[] argValues) {
 Object result = CoroutinesUtils.invokeSuspendingFunction(method, bean, argValues);
 
-// Support DataLoader use
+// Support use of DataLoader from suspending function
 Class<?> returnType = KotlinReflectionUtils.getReturnType(method);
 if (CompletableFuture.class.isAssignableFrom(returnType)) {
 return ((Mono<CompletableFuture<?>>) result).flatMap(Mono::fromFuture);
@@ -123,22 +131,16 @@ private static Object invokeSuspendingFunction(Object bean, Method method, Objec
 return result;
 }
 
-@Nullable
-private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object result) {
-if (this.hasCallableReturnValue && result != null) {
-return CompletableFuture.supplyAsync(
-() -> {
-try {
-return ContextSnapshotFactoryHelper.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
-}
-catch (Exception ex) {
-throw new IllegalStateException(
-"Failure in Callable returned from " + getBridgedMethod().toGenericString(), ex);
-}
-},
-this.executor);
-}
-return result;
+private CompletableFuture<?> adaptCallable(GraphQLContext graphQLContext, Callable<?> result) {
+return CompletableFuture.supplyAsync(() -> {
+try {
+return ContextSnapshotFactoryHelper.captureFrom(graphQLContext).wrap(result).call();
+}
+catch (Exception ex) {
+String msg = "Failure in Callable returned from " + getBridgedMethod().toGenericString();
+throw new IllegalStateException(msg, ex);
+}
+}, this.executor);
 }
 
 /**