Polishing.

Add reactive type translation to Coroutine methods to return the expected type for AOP processing.

See spring-projects#2926

Author: mp911de

src/main/java/org/springframework/data/repository/core/support/RepositoryMethodInvoker.java

@@ -15,12 +15,15 @@
  */
 package org.springframework.data.repository.core.support;
 
+import kotlin.Unit;
 import kotlin.reflect.KFunction;
+import kotlinx.coroutines.flow.Flow;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Collection;
 import java.util.stream.Stream;
 
 import org.reactivestreams.Publisher;
@@ -54,25 +57,63 @@ abstract class RepositoryMethodInvoker {
 	private final Class<?> returnedType;
 	private final Invokable invokable;
 	private final boolean suspendedDeclaredMethod;
-	private final boolean returnsReactiveType;
 
+	@SuppressWarnings("ReactiveStreamsUnusedPublisher")
 	protected RepositoryMethodInvoker(Method method, Invokable invokable) {
 
 		this.method = method;
-		this.invokable = invokable;
 
 		if (KotlinDetector.isKotlinReflectPresent()) {
 
 			this.suspendedDeclaredMethod = KotlinReflectionUtils.isSuspend(method);
 			this.returnedType = this.suspendedDeclaredMethod ? KotlinReflectionUtils.getReturnType(method)
 					: method.getReturnType();
+
+			// special case for most query methods: These can return Flux but we don't want to fail later on if the method
+			// is void.
+			if (suspendedDeclaredMethod) {
+
+				this.invokable = args -> {
+
+					Object result = invokable.invoke(args);
+
+					if (returnedType == Unit.class) {
+
+						if (result instanceof Mono<?> m) {
+							return m.then();
+						}
+
+						Flux<?> flux = result instanceof Flux<?> f ? f : ReactiveWrapperConverters.toWrapper(result, Flux.class);
+						return flux.then();
+					}
+
+					if (returnedType != Flow.class) {
+
+						if (result instanceof Mono<?> m) {
+							return m;
+						}
+
+						Flux<?> flux = result instanceof Flux<?> f ? f : ReactiveWrapperConverters.toWrapper(result, Flux.class);
+
+						if (Collection.class.isAssignableFrom(returnedType)) {
+							return flux.collectList();
+						}
+
+						return flux.singleOrEmpty();
+					}
+
+					return result;
+				};
+			} else {
+				this.invokable = invokable;
+			}
+
 		} else {
 
 			this.suspendedDeclaredMethod = false;
 			this.returnedType = method.getReturnType();
+			this.invokable = invokable;
 		}
-
-		this.returnsReactiveType = ReactiveWrappers.supports(returnedType);
 	}
 
 	static RepositoryQueryMethodInvoker forRepositoryQuery(Method declaredMethod, RepositoryQuery query) {
@@ -154,7 +195,7 @@ private RepositoryMethodInvocation computeInvocationResult(RepositoryMethodInvoc
 	interface Invokable {
 
 		@Nullable
-		Object invoke(Object[] args) throws ReflectiveOperationException;
+		Object invoke(Object[] args) throws Exception;
 	}
 
 	/**
@@ -214,8 +255,6 @@ Publisher<Object> decorate(Class<?> repositoryInterface, RepositoryInvocationMul
 	 */
 	private static class RepositoryFragmentMethodInvoker extends RepositoryMethodInvoker {
 
-		private final CoroutineAdapterInformation adapterInformation;
-
 		public RepositoryFragmentMethodInvoker(Method declaredMethod, Object instance, Method baseClassMethod) {
 			this(CoroutineAdapterInformation.create(declaredMethod, baseClassMethod), declaredMethod, instance,
 					baseClassMethod);
@@ -236,13 +275,12 @@ public RepositoryFragmentMethodInvoker(CoroutineAdapterInformation adapterInform
 						return AopUtils.invokeJoinpointUsingReflection(instance, baseClassMethod, invocationArguments);
 					}
 					return AopUtils.invokeJoinpointUsingReflection(instance, baseClassMethod, args);
-				} catch (RuntimeException e) {
+				} catch (Exception e) {
 					throw e;
 				} catch (Throwable e) {
 					throw new RuntimeException(e);
 				}
 			});
-			this.adapterInformation = adapterInformation;
 		}
 
 		/**

src/test/kotlin/org/springframework/data/repository/kotlin/CoroutineCrudRepositoryUnitTests.kt

@@ -210,6 +210,22 @@ class CoroutineCrudRepositoryUnitTests {
 		Mockito.verify(invocationListener).afterInvocation(captor.capture())
 	}
 
+	@Test // DATACMNS-1508, DATACMNS-1764
+	fun shouldBridgeFluxQueryMethod() {
+
+		val sample = User()
+
+		Mockito.`when`(factory.queryOne.execute(arrayOf("foo", null, any()))).thenReturn(Flux.just(sample))
+
+		val result = runBlocking {
+			coRepository.findOne("foo")
+		}
+
+		assertThat(result).isNotNull().isEqualTo(sample)
+		val captor = ArgumentCaptor.forClass(RepositoryMethodInvocationListener.RepositoryMethodInvocation::class.java)
+		Mockito.verify(invocationListener).afterInvocation(captor.capture())
+	}
+
 	@Test // DATACMNS-1508
 	fun shouldBridgeRxJavaQueryMethod() {
 
@@ -283,13 +299,20 @@ class CoroutineCrudRepositoryUnitTests {
 
 		val sample = User()
 
-		Mockito.`when`(factory.queryOne.execute(arrayOf("foo", null, any()))).thenReturn(Mono.just(listOf(sample)), Mono.empty<User>())
+		Mockito.`when`(factory.queryOne.execute(arrayOf("foo", null, any())))
+			.thenReturn(Mono.just(listOf(sample)), Flux.just(sample, sample), Mono.empty<User>())
 
-		val result = runBlocking {
+		val result1 = runBlocking {
 			coRepository.findSuspendedAsList("foo")
 		}
 
-		assertThat(result).hasSize(1).containsOnly(sample)
+		assertThat(result1).hasSize(1).containsOnly(sample)
+
+		val result2 = runBlocking {
+			coRepository.findSuspendedAsList("foo")
+		}
+
+		assertThat(result2).hasSize(2).contains(sample)
 
 		val emptyResult = runBlocking {
 			coRepository.findSuspendedAsList("foo")
@@ -298,6 +321,18 @@ class CoroutineCrudRepositoryUnitTests {
 		assertThat(emptyResult).isNull()
 	}
 
+	@Test // DATACMNS-1802
+	fun shouldDiscardResult() {
+
+		Mockito.`when`(factory.queryOne.execute(any())).thenReturn(Flux.empty<User>())
+
+		val result = runBlocking {
+			coRepository.someDelete("foo")
+		}
+
+		assertThat(result).isInstanceOf(Unit.javaClass)
+	}
+
 	interface MyCoRepository : CoroutineCrudRepository<User, String> {
 
 		suspend fun findOne(id: String): User
@@ -307,5 +342,7 @@ class CoroutineCrudRepositoryUnitTests {
 		suspend fun findSuspendedMultiple(id: String): Flow<User>
 
 		suspend fun findSuspendedAsList(id: String): List<User>
+
+		suspend fun someDelete(id: String)
 	}
 }