From e9f9d66abed20e88c2f636d1b7f12101bd55c0f8 Mon Sep 17 00:00:00 2001
From: Michael Reiche <48999328+mikereiche@users.noreply.github.com>
Date: Fri, 29 Jul 2022 14:39:16 -0700
Subject: [PATCH] Fix reactive parameter handling. (#1519)

* Fix Reactive Parameter Handling.

Closes #1516.
---
 .../query/AbstractCouchbaseQueryBase.java        | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/spring-data-couchbase/src/main/java/org/springframework/data/couchbase/repository/query/AbstractCouchbaseQueryBase.java b/spring-data-couchbase/src/main/java/org/springframework/data/couchbase/repository/query/AbstractCouchbaseQueryBase.java
index 6a7a5efa..2818a101 100644
--- a/spring-data-couchbase/src/main/java/org/springframework/data/couchbase/repository/query/AbstractCouchbaseQueryBase.java
+++ b/spring-data-couchbase/src/main/java/org/springframework/data/couchbase/repository/query/AbstractCouchbaseQueryBase.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 the original author or authors
+ * Copyright 2020-2022 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.
@@ -18,6 +18,7 @@ package org.springframework.data.couchbase.repository.query;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import org.reactivestreams.Publisher;
 import org.springframework.data.couchbase.core.CouchbaseOperations;
 import org.springframework.data.couchbase.core.ExecutableFindByQueryOperation.ExecutableFindByQuery;
 import org.springframework.data.couchbase.core.query.Query;
@@ -108,15 +109,16 @@ public abstract class AbstractCouchbaseQueryBase<CouchbaseOperationsType> implem
 	 * @see org.springframework.data.repository.query.RepositoryQuery#execute(java.lang.Object[])
 	 */
 	public Object execute(Object[] parameters) {
-
 		ReactiveCouchbaseParameterAccessor accessor = new ReactiveCouchbaseParameterAccessor(getQueryMethod(), parameters);
-
-		Object result = accessor.resolveParameters().map(this::executeDeferred);
-		return ((Mono<Object>) result).block() ;
+		return method.hasReactiveWrapperParameter() ? accessor.resolveParameters().flatMapMany(this::executeDeferred)
+				: execute(accessor);
 	}
 
-	private Object executeDeferred(ReactiveCouchbaseParameterAccessor parameterAccessor) {
-		return execute(parameterAccessor);
+	private Publisher<Object> executeDeferred(ReactiveCouchbaseParameterAccessor parameterAccessor) {
+		if (getQueryMethod().isCollectionQuery()) {
+			return Flux.defer(() -> (Publisher<Object>) execute(parameterAccessor));
+		}
+		return Mono.defer(() -> (Mono<Object>) execute(parameterAccessor));
 	}
 
 	private Object execute(ParametersParameterAccessor parameterAccessor) {