Single result queries (#13251)

* add support for single result queries. Works for both normal and reactive repositories.

* fix test order issue where other contact test removes container

* fix references to renamed classes. Fix checkstyle issues

* add unit test coverae for CosmosQuery get execution methods. Address code review feedback.

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/query/AbstractCosmosQuery.java

@@ -6,6 +6,7 @@
 import com.azure.spring.data.cosmos.core.query.DocumentQuery;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.data.repository.query.ResultProcessor;
+import org.springframework.data.repository.query.ReturnedType;
 
 /**
  * Abstract class for cosmos query.
@@ -39,20 +40,30 @@ public Object execute(Object[] parameters) {
         final ResultProcessor processor = method.getResultProcessor().withDynamicProjection(accessor);
         final String container = ((CosmosEntityMetadata) method.getEntityInformation()).getContainerName();
 
-        final CosmosQueryExecution execution = getExecution(accessor);
+        final CosmosQueryExecution execution = getExecution(accessor, processor.getReturnedType());
+
         return execution.execute(query, processor.getReturnedType().getDomainType(), container);
     }
 
 
-    private CosmosQueryExecution getExecution(CosmosParameterAccessor accessor) {
+    /**
+     * Determines the appropriate execution path for a query
+     *
+     * @param returnedType The return type of the method
+     * @param accessor Object for accessing method parameters
+     * @return the execution type needed to handle the query
+     */
+    protected CosmosQueryExecution getExecution(CosmosParameterAccessor accessor, ReturnedType returnedType) {
         if (isDeleteQuery()) {
             return new CosmosQueryExecution.DeleteExecution(operations);
-        } else if (method.isPageQuery()) {
+        } else if (isPageQuery()) {
             return new CosmosQueryExecution.PagedExecution(operations, accessor.getPageable());
         } else if (isExistsQuery()) {
             return new CosmosQueryExecution.ExistsExecution(operations);
-        } else {
+        } else if (isCollectionQuery()) {
             return new CosmosQueryExecution.MultiEntityExecution(operations);
+        } else {
+            return new CosmosQueryExecution.SingleEntityExecution(operations, returnedType);
         }
     }
 
@@ -71,4 +82,12 @@ public CosmosQueryMethod getQueryMethod() {
 
     protected abstract boolean isExistsQuery();
 
+    protected boolean isPageQuery() {
+        return method.isPageQuery();
+    }
+
+    protected boolean isCollectionQuery() {
+        return method.isCollectionQuery();
+    }
+
 }

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/query/AbstractReactiveCosmosQuery.java

@@ -6,6 +6,8 @@
 import com.azure.spring.data.cosmos.core.query.DocumentQuery;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.data.repository.query.ResultProcessor;
+import org.springframework.data.repository.query.ReturnedType;
+import reactor.core.publisher.Mono;
 
 /**
  * Abstract class for reactive cosmos query.
@@ -43,19 +45,27 @@ public Object execute(Object[] parameters) {
         final String containerName =
             ((ReactiveCosmosEntityMetadata) method.getEntityInformation()).getContainerName();
 
-        final ReactiveCosmosQueryExecution execution = getExecution(accessor);
+        final ReactiveCosmosQueryExecution execution = getExecution(processor.getReturnedType());
         return execution.execute(query, processor.getReturnedType().getDomainType(), containerName);
     }
 
-
-    private ReactiveCosmosQueryExecution getExecution(ReactiveCosmosParameterAccessor accessor) {
+    /**
+     * Determines the appropriate execution path for a reactive query
+     *
+     * @throws IllegalArgumentException if execution requires paging
+     * @param returnedType The return type of the method
+     * @return the execution type needed to handle the query
+     */
+    protected ReactiveCosmosQueryExecution getExecution(ReturnedType returnedType) {
         if (isDeleteQuery()) {
             return new ReactiveCosmosQueryExecution.DeleteExecution(operations);
-        } else if (method.isPageQuery()) {
+        } else if (isPageQuery()) {
             throw new IllegalArgumentException("Paged Query is not supported by reactive cosmos "
                 + "db");
         } else if (isExistsQuery()) {
             return new ReactiveCosmosQueryExecution.ExistsExecution(operations);
+        } else if (isReactiveSingleResultQuery()) {
+            return new ReactiveCosmosQueryExecution.SingleEntityExecution(operations, returnedType);
         } else {
             return new ReactiveCosmosQueryExecution.MultiEntityExecution(operations);
         }
@@ -76,4 +86,12 @@ public ReactiveCosmosQueryMethod getQueryMethod() {
 
     protected abstract boolean isExistsQuery();
 
+    protected boolean isPageQuery() {
+        return method.isPageQuery();
+    }
+
+    private boolean isReactiveSingleResultQuery() {
+        return method.getReactiveWrapper() != null && method.getReactiveWrapper().equals(Mono.class);
+    }
+
 }

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/query/CosmosQueryExecution.java

@@ -5,7 +5,12 @@
 import com.azure.spring.data.cosmos.core.CosmosOperations;
 import com.azure.spring.data.cosmos.core.query.CosmosPageRequest;
 import com.azure.spring.data.cosmos.core.query.DocumentQuery;
+import com.azure.spring.data.cosmos.exception.CosmosAccessException;
 import org.springframework.data.domain.Pageable;
+import org.springframework.data.repository.query.ReturnedType;
+
+import java.util.List;
+import java.util.Optional;
 
 /**
  * Interface to execute cosmos query operations
@@ -40,7 +45,7 @@ public Object execute(DocumentQuery query, Class<?> type, String container) {
     }
 
     /**
-     * Find operation implementation to execute a find query
+     * Find operation implementation to execute a find query for multiple items
      */
     final class MultiEntityExecution implements CosmosQueryExecution {
 
@@ -56,6 +61,43 @@ public Object execute(DocumentQuery query, Class<?> type, String container) {
         }
     }
 
+    /**
+     * Find operation implementation to execute a find query for a single item
+     */
+    final class SingleEntityExecution implements CosmosQueryExecution {
+
+        private final CosmosOperations operations;
+        private final ReturnedType returnedType;
+
+        public SingleEntityExecution(CosmosOperations operations, ReturnedType returnedType) {
+            this.operations = operations;
+            this.returnedType = returnedType;
+        }
+
+        @Override
+        public Object execute(DocumentQuery query, Class<?> type, String collection) {
+            final List<?> results = operations.find(query, type, collection);
+            final Object result;
+            if (results == null || results.isEmpty()) {
+                result = null;
+            } else if (results.size() == 1) {
+                result = results.get(0);
+            } else {
+                throw new CosmosAccessException("Too many results - return type "
+                    + returnedType.getReturnedType()
+                    + " is not of type Iterable but find returned "
+                    + results.size()
+                    + " results");
+            }
+
+            if (returnedType.getReturnedType() == Optional.class) {
+                return result == null ? Optional.empty() : Optional.of(result);
+            } else {
+                return result;
+            }
+        }
+    }
+
     /**
      * exist operation implementation to execute a exists query
      */

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/query/ReactiveCosmosQueryExecution.java

@@ -4,6 +4,8 @@
 
 import com.azure.spring.data.cosmos.core.ReactiveCosmosOperations;
 import com.azure.spring.data.cosmos.core.query.DocumentQuery;
+import com.azure.spring.data.cosmos.exception.CosmosAccessException;
+import org.springframework.data.repository.query.ReturnedType;
 
 /**
  * Interface to execute reactive cosmos query operations
@@ -38,7 +40,7 @@ public Object execute(DocumentQuery query, Class<?> type, String container) {
     }
 
     /**
-     * Find operation implementation to execute a find query
+     * Find operation implementation to execute a find query for multiple items
      */
     final class MultiEntityExecution implements ReactiveCosmosQueryExecution {
 
@@ -54,6 +56,34 @@ public Object execute(DocumentQuery query, Class<?> type, String container) {
         }
     }
 
+    /**
+     * Find operation implementation to execute a find query for a single item
+     */
+    final class SingleEntityExecution implements ReactiveCosmosQueryExecution {
+
+        private final ReactiveCosmosOperations operations;
+        private final ReturnedType returnedType;
+
+        public SingleEntityExecution(ReactiveCosmosOperations operations, ReturnedType returnedType) {
+            this.operations = operations;
+            this.returnedType = returnedType;
+        }
+
+        @Override
+        public Object execute(DocumentQuery query, Class<?> type, String container) {
+            return operations.find(query, type, container)
+                .buffer(2)
+                .map((vals) -> {
+                    if (vals.size() > 1) {
+                        throw new CosmosAccessException("Too many results - Expected Mono<"
+                            + returnedType.getReturnedType()
+                            + "> but query returned multiple results");
+                    }
+                    return vals.iterator().next();
+                });
+        }
+    }
+
     /**
      * Exist operation implementation to execute a exist query
      */

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/query/ReactiveCosmosQueryMethod.java

@@ -7,6 +7,8 @@
 import org.springframework.data.repository.core.EntityMetadata;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.QueryMethod;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import java.lang.reflect.Method;
 
@@ -16,6 +18,7 @@
 public class ReactiveCosmosQueryMethod extends QueryMethod {
 
     private ReactiveCosmosEntityMetadata<?> metadata;
+    private final Method method;
 
     /**
      * Creates a new {@link QueryMethod} from the given parameters. Looks up the correct query to use for following
@@ -27,6 +30,7 @@ public class ReactiveCosmosQueryMethod extends QueryMethod {
      */
     public ReactiveCosmosQueryMethod(Method method, RepositoryMetadata metadata, ProjectionFactory factory) {
         super(method, metadata, factory);
+        this.method = method;
     }
 
     @Override
@@ -39,4 +43,17 @@ public EntityMetadata<?> getEntityInformation() {
         this.metadata = new SimpleReactiveCosmosEntityMetadata<Object>(domainType, entityInformation);
         return this.metadata;
     }
+
+    /**
+     * Returns the reactive wrapper class type if it exists or null otherwise
+     *
+     * @return Reactive wrapper class (Flux or Mono)
+     */
+    public Class<?> getReactiveWrapper() {
+        return isReactiveWrapperClass(method.getReturnType()) ? method.getReturnType() : null;
+    }
+
+    private static boolean isReactiveWrapperClass(Class<?> clazz) {
+        return clazz.equals(Flux.class) || clazz.equals(Mono.class);
+    }
 }

sdk/cosmos/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/integration/ReactiveCourseRepositoryIT.java

@@ -133,6 +133,28 @@ public void testFindAll() {
         StepVerifier.create(allFlux).expectNextCount(4).verifyComplete();
     }
 
+    @Test
+    public void testFindOneShouldFailIfMultipleResultsReturned() {
+        final Course course = new Course("unusedId", COURSE_1.getName(), COURSE_1.getDepartment());
+        final Mono<Course> saveSecond = repository.save(course);
+        StepVerifier.create(saveSecond).expectNext(course).verifyComplete();
+
+        final Mono<Course> find = repository.findOneByName(COURSE_1.getName());
+        StepVerifier.create(find).expectError(CosmosAccessException.class).verify();
+    }
+
+    @Test
+    public void testShouldFindSingleEntity() {
+        final Mono<Course> find = repository.findOneByName(COURSE_1.getName());
+        StepVerifier.create(find).expectNext(COURSE_1).expectComplete().verify();
+    }
+
+    @Test
+    public void testShouldReturnEmptyMonoWhenNoResults() {
+        final Mono<Course> find = repository.findOneByName("unusedName");
+        StepVerifier.create(find).verifyComplete();
+    }
+
     @Test
     public void testInsert() {
         final Mono<Course> save = repository.save(COURSE_5);