Apply ReactiveAdapterRegistry to return values

CoroutinesUtils#invokeSuspendingFunction in adapts return values to Mono
and Flux in Kotlin. However, Flow methods don't have to be suspending
functions. This change ensures that ReactiveAdapterRegistry is applied
wherever return values are handled.

Closes gh-988

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java

@@ -18,14 +18,14 @@
 
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 
 import graphql.schema.DataFetchingEnvironment;
 import reactor.core.publisher.Mono;
 
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.data.method.annotation.support.DataFetcherHandlerMethodSupport;
+import org.springframework.graphql.execution.ReactiveAdapterRegistryHelper;
 import org.springframework.lang.Nullable;
 
 /**
@@ -80,18 +80,8 @@ Mono<Object> getEntities(DataFetchingEnvironment env, List<Map<String, Object>>
 	}
 
 	private Mono<Object> doInvoke(DataFetchingEnvironment env, Object[] args) {
-
 		Object result = doInvoke(env.getGraphQlContext(), args);
-
-		if (result instanceof Mono<?> mono) {
-			return mono.cast(Object.class);
-		}
-		else if (result instanceof CompletableFuture<?> future) {
-			return Mono.fromFuture(future);
-		}
-		else {
-			return Mono.justOrEmpty(result);
-		}
+		return ReactiveAdapterRegistryHelper.toMono(result);
 	}
 
 }

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

@@ -78,6 +78,7 @@
 import org.springframework.graphql.data.pagination.CursorStrategy;
 import org.springframework.graphql.data.query.SortStrategy;
 import org.springframework.graphql.execution.BatchLoaderRegistry;
+import org.springframework.graphql.execution.ReactiveAdapterRegistryHelper;
 import org.springframework.graphql.execution.RuntimeWiringConfigurer;
 import org.springframework.graphql.execution.SelfDescribingDataFetcher;
 import org.springframework.graphql.execution.SubscriptionPublisherException;
@@ -528,18 +529,24 @@ public Object get(DataFetchingEnvironment environment) throws Exception {
 		}
 
 		@SuppressWarnings({"unchecked", "ReactiveStreamsUnusedPublisher"})
+		@Nullable
 		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));
+			if (this.subscription) {
+				return ReactiveAdapterRegistryHelper.toSubscriptionFlux(result)
+						.onErrorResume((ex) -> handleSubscriptionError(ex, env, handlerMethod));
 			}
-			else if (result instanceof Mono) {
+
+			result = ReactiveAdapterRegistryHelper.toMonoOrFluxIfReactive(result);
+
+			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;
 		}
 

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

@@ -21,9 +21,7 @@
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.Callable;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
-import java.util.function.Function;
 
 import graphql.GraphQLContext;
 import org.dataloader.BatchLoaderEnvironment;
@@ -37,6 +35,7 @@
 import org.springframework.graphql.data.method.HandlerMethod;
 import org.springframework.graphql.data.method.InvocableHandlerMethodSupport;
 import org.springframework.graphql.data.method.annotation.ContextValue;
+import org.springframework.graphql.execution.ReactiveAdapterRegistryHelper;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
@@ -97,11 +96,11 @@ public <K, V> Mono<Map<K, V>> invokeForMap(Collection<K> keys, BatchLoaderEnviro
 		Object[] args = getMethodArgumentValues(keys, environment);
 		if (doesNotHaveAsyncArgs(args)) {
 			Object result = doInvoke(environment.getContext(), args);
-			return toMonoMap(result);
+			return ReactiveAdapterRegistryHelper.toMono(result);
 		}
 		return toArgsMono(args).flatMap((argValues) -> {
 			Object result = doInvoke(environment.getContext(), argValues);
-			return toMonoMap(result);
+			return ReactiveAdapterRegistryHelper.toMono(result);
 		});
 	}
 
@@ -117,11 +116,11 @@ public <V> Flux<V> invokeForIterable(Collection<?> keys, BatchLoaderEnvironment
 		Object[] args = getMethodArgumentValues(keys, environment);
 		if (doesNotHaveAsyncArgs(args)) {
 			Object result = doInvoke(environment.getContext(), args);
-			return toFlux(result);
+			return ReactiveAdapterRegistryHelper.toFluxFromCollection(result);
 		}
 		return toArgsMono(args).flatMapMany((resolvedArgs) -> {
 			Object result = doInvoke(environment.getContext(), resolvedArgs);
-			return toFlux(result);
+			return ReactiveAdapterRegistryHelper.toFluxFromCollection(result);
 		});
 	}
 
@@ -185,33 +184,4 @@ private boolean doesNotHaveAsyncArgs(Object[] args) {
 		return Arrays.stream(args).noneMatch((arg) -> arg instanceof Mono);
 	}
 
-	@SuppressWarnings("unchecked")
-	private static <K, V> Mono<Map<K, V>> toMonoMap(@Nullable Object result) {
-		if (result instanceof Map) {
-			return Mono.just((Map<K, V>) result);
-		}
-		else if (result instanceof Mono) {
-			return (Mono<Map<K, V>>) result;
-		}
-		else if (result instanceof CompletableFuture) {
-			return Mono.fromFuture((CompletableFuture<? extends Map<K, V>>) result);
-		}
-		return Mono.error(new IllegalStateException("Unexpected return value: " + result));
-	}
-
-	@SuppressWarnings("unchecked")
-	private static <V> Flux<V> toFlux(@Nullable Object result) {
-		if (result instanceof Collection) {
-			return Flux.fromIterable((Collection<V>) result);
-		}
-		else if (result instanceof Flux) {
-			return (Flux<V>) result;
-		}
-		else if (result instanceof CompletableFuture) {
-			return Mono.fromFuture((CompletableFuture<? extends Collection<V>>) result)
-					.flatMapIterable(Function.identity());
-		}
-		return Flux.error(new IllegalStateException("Unexpected return value: " + result));
-	}
-
 }

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

@@ -18,18 +18,16 @@
 
 import java.util.Arrays;
 import java.util.concurrent.Callable;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.function.BiConsumer;
 
 import graphql.schema.DataFetchingEnvironment;
-import org.reactivestreams.Publisher;
-import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 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.execution.ReactiveAdapterRegistryHelper;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
@@ -112,6 +110,7 @@ public Object invoke(DataFetchingEnvironment environment) {
 	 * @param providedArgs additional arguments to be matched by their type
 	 * @since 1.2.0
 	 */
+	@SuppressWarnings("ReactiveStreamsUnusedPublisher")
 	@Nullable
 	public Object invoke(DataFetchingEnvironment environment, Object... providedArgs) {
 		Object[] args;
@@ -129,23 +128,11 @@ public Object invoke(DataFetchingEnvironment environment, Object... providedArgs
 		return this.subscription ?
 				toArgsMono(args).flatMapMany((argValues) -> {
 					Object result = validateAndInvoke(argValues, environment);
-					Assert.state(result instanceof Publisher, "Expected a Publisher from a Subscription response");
-					return Flux.from((Publisher<?>) result);
+					return ReactiveAdapterRegistryHelper.toSubscriptionFlux(result);
 				}) :
 				toArgsMono(args).flatMap((argValues) -> {
 					Object result = validateAndInvoke(argValues, environment);
-					if (result instanceof Mono<?> mono) {
-						return mono;
-					}
-					else if (result instanceof Flux<?> flux) {
-						return Flux.from(flux).collectList();
-					}
-					else if (result instanceof CompletableFuture<?> future) {
-						return Mono.fromFuture(future);
-					}
-					else {
-						return Mono.justOrEmpty(result);
-					}
+					return ReactiveAdapterRegistryHelper.toMono(result);
 				});
 	}
 

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

@@ -73,6 +73,7 @@ private ContextDataFetcherDecorator(
 	}
 
 
+	@SuppressWarnings("ReactiveStreamsUnusedPublisher")
 	@Override
 	public Object get(DataFetchingEnvironment env) throws Exception {
 
@@ -86,24 +87,22 @@ public Object get(DataFetchingEnvironment env) throws Exception {
 		Object value = snapshot.wrap(() -> this.delegate.get(env)).call();
 
 		if (this.subscription) {
-			Assert.state(value instanceof Publisher, "Expected Publisher for a subscription");
-			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);
+			return ReactiveAdapterRegistryHelper.toSubscriptionFlux(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)));
+					})
+					.contextWrite(snapshot::updateContext);
 		}
 
-		if (value instanceof Flux) {
-			value = ((Flux<?>) value).collectList();
-		}
+		value = ReactiveAdapterRegistryHelper.toMonoIfReactive(value);
 
-		if (value instanceof Mono<?> valueMono) {
-			value = valueMono.contextWrite(snapshot::updateContext).toFuture();
+		if (value instanceof Mono<?> mono) {
+			value = mono.contextWrite(snapshot::updateContext).toFuture();
 		}
 
 		return value;

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

@@ -0,0 +1,135 @@
+/*
+ * 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.
+ * 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.execution;
+
+import java.util.Collection;
+
+import io.micrometer.context.Nullable;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.core.ReactiveAdapter;
+import org.springframework.core.ReactiveAdapterRegistry;
+import org.springframework.util.Assert;
+
+/**
+ * Helper to adapt a result Object to {@link Mono} or {@link Flux} through
+ * {@link ReactiveAdapterRegistry}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.3.1
+ */
+@SuppressWarnings({"ReactiveStreamsUnusedPublisher", "unchecked"})
+public abstract class ReactiveAdapterRegistryHelper {
+
+	private static final ReactiveAdapterRegistry registry = ReactiveAdapterRegistry.getSharedInstance();
+
+
+	/**
+	 * Return a {@link Mono} for the given Object by delegating to
+	 * {@link #toMonoIfReactive}, and then applying {@link Mono#justOrEmpty}
+	 * if necessary.
+	 * @param result the result Object to adapt
+	 * @param <T> the type of element in the Mono to cast to
+	 * @return a {@code Mono} that represents the result
+	 */
+	public static <T> Mono<T> toMono(@Nullable Object result) {
+		result = toMonoIfReactive(result);
+		return (Mono<T>) ((result instanceof Mono<?> mono) ? mono : Mono.justOrEmpty(result));
+	}
+
+	/**
+	 * Return a {@link Mono} for the given result Object if it can be adapted
+	 * to a {@link Publisher} via {@link ReactiveAdapterRegistry}. Multivalued
+	 * publishers are collected to a List.
+	 * @param result the result Object to adapt
+	 * @return the same instance or a {@code Mono} if the object is known to
+	 * {@code ReactiveAdapterRegistry}
+	 */
+	@Nullable
+	public static Object toMonoIfReactive(@Nullable Object result) {
+		ReactiveAdapter adapter = ((result != null) ? registry.getAdapter(result.getClass()) : null);
+		if (adapter == null) {
+			return result;
+		}
+		Publisher<?> publisher = adapter.toPublisher(result);
+		return (adapter.isMultiValue() ? Flux.from(publisher).collectList() : Mono.from(publisher));
+	}
+
+	/**
+	 * Adapt the given result Object to {@link Mono} or {@link Flux} if it can
+	 * be adapted to a single or multi-value {@link Publisher} respectively
+	 * via {@link ReactiveAdapterRegistry}.
+	 * @param result the result Object to adapt
+	 * @return the same instance, a {@code Mono}, or a {@code Flux}
+	 */
+	@Nullable
+	public static Object toMonoOrFluxIfReactive(@Nullable Object result) {
+		ReactiveAdapter adapter = ((result != null) ? registry.getAdapter(result.getClass()) : null);
+		if (adapter == null) {
+			return result;
+		}
+		Publisher<Object> publisher = adapter.toPublisher(result);
+		return (adapter.isMultiValue() ? Flux.from(publisher) : Mono.from(publisher));
+	}
+
+	/**
+	 * Return a {@link Flux} for the given result Object, adapting to a
+	 * {@link Publisher} first if necessary via {@link ReactiveAdapterRegistry}.
+	 * @param result the result Object to adapt
+	 * @return a {@link Flux}, possibly empty if the result is {@code null}
+	 */
+	public static Flux<?> toSubscriptionFlux(@Nullable Object result) {
+		if (result == null) {
+			return Flux.empty();
+		}
+		if (result instanceof Publisher<?> publisher) {
+			return Flux.from(publisher);
+		}
+		ReactiveAdapter adapter = registry.getAdapter(result.getClass());
+		Assert.state(adapter != null, "Expected Publisher for a subscription");
+		return Flux.from(adapter.toPublisher(result));
+	}
+
+	/**
+	 * Return a {@link Flux} for the given result Object that represents a
+	 * logical collection of values. The Object must be a {@link Collection}
+	 * or a publisher of a {@code Collection}, which is flattened with
+	 * {@link Flux#fromIterable(Iterable)}, or a multi-value publisher.
+	 * @param result the result Object to adapt
+	 * @param <T> the type of element in the collection to cast to
+	 * @return a {@code Flux} that represents the collection
+	 */
+	public static <T> Flux<T> toFluxFromCollection(@Nullable Object result) {
+		if (result instanceof Collection) {
+			return Flux.fromIterable((Collection<T>) result);
+		}
+		ReactiveAdapter adapter = ((result != null) ? registry.getAdapter(result.getClass()) : null);
+		if (adapter == null) {
+			return Flux.error(new IllegalStateException("Unexpected return value: " + result));
+		}
+		Publisher<?> publisher = adapter.toPublisher(result);
+		if (adapter.isMultiValue()) {
+			return (Flux<T>) Flux.from(publisher);
+		}
+		else {
+			return Mono.from(publisher).flatMapMany((c) -> Flux.fromIterable((Collection<T>) c));
+		}
+	}
+
+}