Single result queries

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,24 @@ 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) {
+
+    private CosmosQueryExecution getExecution(CosmosParameterAccessor accessor, ReturnedType returnedType) {
         if (isDeleteQuery()) {
             return new CosmosQueryExecution.DeleteExecution(operations);
         } else if (method.isPageQuery()) {
             return new CosmosQueryExecution.PagedExecution(operations, accessor.getPageable());
         } else if (isExistsQuery()) {
             return new CosmosQueryExecution.ExistsExecution(operations);
-        } else {
+        } else if (method.isCollectionQuery()) {
             return new CosmosQueryExecution.MultiEntityExecution(operations);
+        } else {
+            return new CosmosQueryExecution.SingleEntityExecution(operations, returnedType);
         }
     }
 

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,21 @@ public Object execute(Object[] parameters) {
         final String containerName =
             ((ReactiveCosmosEntityMetadata) method.getEntityInformation()).getContainerName();
 
-        final ReactiveCosmosQueryExecution execution = getExecution(accessor);
+        final ReactiveCosmosQueryExecution execution = getExecution(accessor, processor.getReturnedType());
         return execution.execute(query, processor.getReturnedType().getDomainType(), containerName);
     }
 
-
-    private ReactiveCosmosQueryExecution getExecution(ReactiveCosmosParameterAccessor accessor) {
+    private ReactiveCosmosQueryExecution getExecution(ReactiveCosmosParameterAccessor accessor,
+                                                      ReturnedType returnedType) {
         if (isDeleteQuery()) {
             return new ReactiveCosmosQueryExecution.DeleteExecution(operations);
         } else if (method.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 +80,8 @@ public ReactiveCosmosQueryMethod getQueryMethod() {
 
     protected abstract boolean isExistsQuery();
 
+    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.findByName(COURSE_1.getName());
+        StepVerifier.create(find).expectError(CosmosAccessException.class).verify();
+    }
+
+    @Test
+    public void testShouldFindSingleEntity() {
+        final Mono<Course> find = repository.findByName(COURSE_1.getName());
+        StepVerifier.create(find).expectNext(COURSE_1).expectComplete().verify();
+    }
+
+    @Test
+    public void testShouldReturnEmptyMonoWhenNoResults() {
+        final Mono<Course> find = repository.findByName("unusedName");
+        StepVerifier.create(find).verifyComplete();
+    }
+
     @Test
     public void testInsert() {
         final Mono<Course> save = repository.save(COURSE_5);

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

@@ -0,0 +1,99 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.spring.data.cosmos.repository.integration;
+
+import com.azure.spring.data.cosmos.core.CosmosTemplate;
+import com.azure.spring.data.cosmos.domain.Contact;
+import com.azure.spring.data.cosmos.exception.CosmosAccessException;
+import com.azure.spring.data.cosmos.repository.TestRepositoryConfig;
+import com.azure.spring.data.cosmos.repository.repository.ContactRepository;
+import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import java.util.Iterator;
+import java.util.List;
+import java.util.Optional;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(classes = TestRepositoryConfig.class)
+public class SingleResponseRepositoryIT {
+
+    private static final Contact TEST_CONTACT = new Contact("testId", "faketitle");
+
+    private static final CosmosEntityInformation<Contact, String> entityInformation =
+        new CosmosEntityInformation<>(Contact.class);
+
+    private static CosmosTemplate staticTemplate;
+    private static boolean isSetupDone;
+
+    @Autowired
+    private ContactRepository repository;
+
+    @Autowired
+    private CosmosTemplate template;
+
+    @Before
+    public void setUp() {
+        if (!isSetupDone) {
+            staticTemplate = template;
+            template.createContainerIfNotExists(entityInformation);
+        }
+        repository.save(TEST_CONTACT);
+        isSetupDone = true;
+    }
+
+    @After
+    public void cleanup() {
+        repository.deleteAll();
+    }
+
+    @AfterClass
+    public static void afterClassCleanup() {
+        staticTemplate.deleteContainer(entityInformation.getContainerName());
+    }
+
+    @Test
+    public void testShouldFindSingleEntity() {
+        final Contact contact = repository.findOneByTitle(TEST_CONTACT.getTitle());
+
+        Assert.assertEquals(TEST_CONTACT, contact);
+    }
+
+    @Test
+    public void testShouldFindSingleOptionalEntity() {
+        final Optional<Contact> contact = repository.findOptionallyByTitle(TEST_CONTACT.getTitle());
+        Assert.assertTrue(contact.isPresent());
+        Assert.assertEquals(TEST_CONTACT, contact.get());
+
+        Assert.assertFalse(repository.findOptionallyByTitle("not here").isPresent());
+    }
+
+    @Test(expected = CosmosAccessException.class)
+    public void testShouldFailIfMultipleResultsReturned() {
+        repository.save(new Contact("testId2", TEST_CONTACT.getTitle()));
+
+        repository.findOneByTitle(TEST_CONTACT.getTitle());
+    }
+
+    @Test
+    public void testShouldAllowListAndIterableResponses() {
+        final List<Contact> contactList = repository.findByTitle(TEST_CONTACT.getTitle());
+        Assert.assertEquals(TEST_CONTACT, contactList.get(0));
+        Assert.assertEquals(1, contactList.size());
+
+        final Iterator<Contact> contactIterator = repository.findByLogicId(TEST_CONTACT.getLogicId()).iterator();
+        Assert.assertTrue(contactIterator.hasNext());
+        Assert.assertEquals(TEST_CONTACT, contactIterator.next());
+        Assert.assertFalse(contactIterator.hasNext());
+    }
+
+}

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

@@ -7,8 +7,15 @@
 import org.springframework.stereotype.Repository;
 
 import java.util.List;
+import java.util.Optional;
 
 @Repository
 public interface ContactRepository extends CosmosRepository<Contact, String> {
     List<Contact> findByTitle(String title);
+
+    Iterable<Contact> findByLogicId(String title);
+
+    Contact findOneByTitle(String title);
+
+    Optional<Contact> findOptionallyByTitle(String title);
 }