Added new feature API for find by id and partition using read document() call (#433)

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java

@@ -31,6 +31,8 @@ public interface DocumentDbOperations {
 
     <T> T findById(String collectionName, Object id, Class<T> entityClass);
 
+    <T> T findById(Object id, Class<T> entityClass, PartitionKey partitionKey);
+
     <T> T insert(T objectToSave, PartitionKey partitionKey);
 
     <T> T insert(String collectionName, T objectToSave, PartitionKey partitionKey);

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java

@@ -153,6 +153,32 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont
         }
     }
 
+    @Override
+    public <T> T findById(Object id, Class<T> entityClass, PartitionKey partitionKey) {
+        Assert.notNull(entityClass, "entityClass should not be null");
+        Assert.notNull(partitionKey, "partitionKey should not be null");
+        assertValidId(id);
+
+        final com.azure.data.cosmos.PartitionKey pk = toCosmosPartitionKey(partitionKey);
+
+        try {
+            final String collectionName = getCollectionName(entityClass);
+            return cosmosClient
+                .getDatabase(databaseName)
+                .getContainer(collectionName)
+                .getItem(id.toString(), pk)
+                .read()
+                .flatMap(cosmosItemResponse -> Mono.justOrEmpty(toDomainObject(entityClass,
+                    cosmosItemResponse.properties())))
+                .onErrorResume(Mono::error)
+                .block();
+
+        } catch (Exception e) {
+            throw new DocumentDBAccessException("findById exception", e);
+        }
+
+    }
+
     public <T> void upsert(T object, PartitionKey partitionKey) {
         Assert.notNull(object, "Upsert object should not be null");
 

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/ReactiveCosmosOperations.java

@@ -8,6 +8,7 @@ package com.microsoft.azure.spring.data.cosmosdb.core;
 
 import com.azure.data.cosmos.CosmosContainerResponse;
 import com.azure.data.cosmos.PartitionKey;
+import com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter;
 import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
 import com.microsoft.azure.spring.data.cosmosdb.repository.support.DocumentDbEntityInformation;
 import reactor.core.publisher.Flux;
@@ -25,6 +26,8 @@ public interface ReactiveCosmosOperations {
 
     <T> Mono<T> findById(String collectionName, Object id, Class<T> entityClass);
 
+    <T> Mono<T> findById(Object id, Class<T> entityClass, PartitionKey partitionKey);
+
     <T> Mono<T> insert(T objectToSave, PartitionKey partitionKey);
 
     <T> Mono<T> insert(String collectionName, Object objectToSave, PartitionKey partitionKey);
@@ -50,4 +53,6 @@ public interface ReactiveCosmosOperations {
     Mono<Long> count(String collectionName);
 
     Mono<Long> count(DocumentQuery query, String containerName);
+
+    MappingDocumentDbConverter getConverter();
 }

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/ReactiveCosmosTemplate.java

@@ -43,6 +43,8 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
 
     private final String databaseName;
 
+    private final MappingDocumentDbConverter mappingDocumentDbConverter;
+
     private final CosmosClient cosmosClient;
 
     private final List<String> collectionCache;
@@ -62,6 +64,7 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
 
         this.databaseName = dbName;
         this.collectionCache = new ArrayList<>();
+        this.mappingDocumentDbConverter = mappingDocumentDbConverter;
 
         this.cosmosClient = cosmosDbFactory.getCosmosClient();
     }
@@ -74,6 +77,11 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
         //  NOTE: When application context instance variable gets introduced, assign it here.
     }
 
+    @Override
+    public MappingDocumentDbConverter getConverter() {
+        return this.mappingDocumentDbConverter;
+    }
+
     /**
      * Creates a collection if it doesn't already exist
      *
@@ -157,12 +165,36 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
                 .flatMap(cosmosItemFeedResponse -> Mono.justOrEmpty(cosmosItemFeedResponse
                     .results()
                     .stream()
-                    .map(cosmosItem -> cosmosItem.toObject(entityClass))
+                    .map(cosmosItem -> toDomainObject(entityClass, cosmosItem))
                     .findFirst()))
                 .onErrorResume(Mono::error)
                 .next();
     }
 
+    /**
+     * Find by id
+     *
+     * @param id            the id
+     * @param entityClass   the entity class
+     * @param partitionKey  partition Key
+     * @return Mono with the item or error
+     */
+    @Override
+    public <T> Mono<T> findById(Object id, Class<T> entityClass, PartitionKey partitionKey) {
+        Assert.notNull(entityClass, "entityClass should not be null");
+        assertValidId(id);
+
+        final String containerName = getContainerName(entityClass);
+        return cosmosClient.getDatabase(databaseName)
+                           .getContainer(containerName)
+                           .getItem(id.toString(), partitionKey)
+                           .read()
+                           .flatMap(cosmosItemResponse -> Mono.justOrEmpty(toDomainObject(entityClass,
+                               cosmosItemResponse.properties())))
+                           .onErrorResume(Mono::error);
+    }
+
+
     /**
      * Insert
      *
@@ -190,7 +222,7 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
                 .getContainer(getContainerName(objectToSave.getClass()))
                 .createItem(objectToSave, new CosmosItemRequestOptions())
                 .onErrorResume(Mono::error)
-                .flatMap(cosmosItemResponse -> Mono.just(cosmosItemResponse.properties().toObject(domainClass)));
+                .flatMap(cosmosItemResponse -> Mono.just(toDomainObject(domainClass, cosmosItemResponse.properties())));
     }
 
     /**
@@ -215,7 +247,7 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
                 .getContainer(containerName)
                 .createItem(objectToSave, options)
                 .onErrorResume(Mono::error)
-                .flatMap(cosmosItemResponse -> Mono.just(cosmosItemResponse.properties().toObject(domainClass)));
+                .flatMap(cosmosItemResponse -> Mono.just(toDomainObject(domainClass, cosmosItemResponse.properties())));
     }
 
     /**
@@ -249,7 +281,7 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
         return cosmosClient.getDatabase(this.databaseName)
                 .getContainer(containerName)
                 .upsertItem(object, options)
-                .flatMap(cosmosItemResponse -> Mono.just(cosmosItemResponse.properties().toObject(domainClass)))
+                .flatMap(cosmosItemResponse -> Mono.just(toDomainObject(domainClass, cosmosItemResponse.properties())))
                 .onErrorResume(Mono::error);
     }
 
@@ -341,7 +373,7 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
     @Override
     public <T> Flux<T> find(DocumentQuery query, Class<T> entityClass, String containerName) {
         return findDocuments(query, entityClass, containerName)
-                .map(cosmosItemProperties -> cosmosItemProperties.toObject(entityClass));
+                .map(cosmosItemProperties -> toDomainObject(entityClass, cosmosItemProperties));
     }
 
     /**
@@ -498,4 +530,9 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
                 .map(cosmosItemResponse -> cosmosItemProperties);
     }
 
+    private <T> T toDomainObject(@NonNull Class<T> domainClass, CosmosItemProperties cosmosItemProperties) {
+        return this.mappingDocumentDbConverter.read(domainClass, cosmosItemProperties);
+    }
+
+
 }

src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/DocumentDbRepository.java

@@ -6,12 +6,25 @@
 
 package com.microsoft.azure.spring.data.cosmosdb.repository;
 
+import com.microsoft.azure.documentdb.PartitionKey;
 import org.springframework.data.repository.NoRepositoryBean;
 import org.springframework.data.repository.PagingAndSortingRepository;
 
 import java.io.Serializable;
+import java.util.Optional;
 
 @NoRepositoryBean
 public interface DocumentDbRepository<T, ID extends Serializable> extends PagingAndSortingRepository<T, ID> {
+
+    /**
+     * Retrieves an entity by its id.
+     *
+     * @param id must not be {@literal null}.
+     * @param partitionKey partition key value of entity
+     * @return the entity with the given id or {@literal Optional#empty()} if none found
+     * @throws IllegalArgumentException if {@code id} is {@literal null}.
+     */
+    Optional<T> findById(ID id, PartitionKey partitionKey);
+
 }
 

src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/ReactiveCosmosRepository.java

@@ -5,9 +5,20 @@
  */
 package com.microsoft.azure.spring.data.cosmosdb.repository;
 
+import com.azure.data.cosmos.PartitionKey;
 import org.springframework.data.repository.NoRepositoryBean;
 import org.springframework.data.repository.reactive.ReactiveSortingRepository;
+import reactor.core.publisher.Mono;
 
 @NoRepositoryBean
 public interface ReactiveCosmosRepository<T, K> extends ReactiveSortingRepository<T, K> {
+
+    /**
+     * Retrieves an entity by its id and partition key.
+     * @param id must not be {@literal null}.
+     * @param partitionKey partition key value of the entity.
+     * @return {@link Mono} emitting the entity with the given id or {@link Mono#empty()} if none found.
+     * @throws IllegalArgumentException in case the given {@code id} is {@literal null}.
+     */
+    Mono<T> findById(K id, PartitionKey partitionKey);
 }

src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java

@@ -141,6 +141,24 @@ public class SimpleDocumentDbRepository<T, ID extends Serializable> implements D
         return Optional.ofNullable(operation.findById(information.getCollectionName(), id, information.getJavaType()));
     }
 
+    /**
+     * find one entity per id without partitions
+     *
+     * @param id
+     * @return
+     */
+    @Override
+    public Optional<T> findById(ID id, PartitionKey partitionKey) {
+        Assert.notNull(id, "id must not be null");
+
+        if (id instanceof String && !StringUtils.hasText((String) id)) {
+            return Optional.empty();
+        }
+
+        return Optional.ofNullable(operation.findById(id, information.getJavaType(), partitionKey));
+    }
+
+
     /**
      * return count of documents in one collection without partitions
      *

src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleReactiveCosmosRepository.java

@@ -83,6 +83,13 @@ public class SimpleReactiveCosmosRepository<T, K extends Serializable> implement
                         id, entityInformation.getJavaType()));
     }
 
+    @Override
+    public Mono<T> findById(K id, PartitionKey partitionKey) {
+        Assert.notNull(id, "The given id must not be null!");
+        return cosmosOperations.findById(id,
+            entityInformation.getJavaType(), partitionKey);
+    }
+
     @Override
     public Mono<Boolean> existsById(K id) {
         Assert.notNull(id, "The given id must not be null!");

src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java

@@ -111,6 +111,15 @@ public class DocumentDbTemplatePartitionIT {
         assertEquals(TEST_PERSON, result.get(0));
     }
 
+    @Test
+    public void testFindByIdWithPartition() {
+        final PartitionPerson partitionPersonById = dbTemplate.findById(TEST_PERSON.getId(),
+            PartitionPerson.class,
+            new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON)));
+
+        assertEquals(TEST_PERSON, partitionPersonById);
+    }
+
     @Test
     public void testFindByNonExistIdWithPartition() {
         final Criteria criteria = Criteria.getInstance(IS_EQUAL, PROPERTY_ID, Arrays.asList(NOT_EXIST_ID));

src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/ReactiveCosmosTemplatePartitionIT.java

@@ -5,7 +5,6 @@
  */
 package com.microsoft.azure.spring.data.cosmosdb.core;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
 import com.azure.data.cosmos.PartitionKey;
 import com.microsoft.azure.spring.data.cosmosdb.CosmosDbFactory;
 import com.microsoft.azure.spring.data.cosmosdb.config.DocumentDBConfig;
@@ -80,14 +79,13 @@ public class ReactiveCosmosTemplatePartitionIT {
             final CosmosDbFactory dbFactory = new CosmosDbFactory(dbConfig);
 
             final DocumentDbMappingContext mappingContext = new DocumentDbMappingContext();
-            final ObjectMapper objectMapper = new ObjectMapper();
 
             personInfo = new DocumentDbEntityInformation<>(PartitionPerson.class);
             containerName = personInfo.getCollectionName();
 
             mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
 
-            final MappingDocumentDbConverter dbConverter = new MappingDocumentDbConverter(mappingContext, objectMapper);
+            final MappingDocumentDbConverter dbConverter = new MappingDocumentDbConverter(mappingContext, null);
             cosmosTemplate = new ReactiveCosmosTemplate(dbFactory, dbConverter, DB_NAME);
             cosmosTemplate.createCollectionIfNotExists(personInfo).block();
 
@@ -113,6 +111,17 @@ public class ReactiveCosmosTemplatePartitionIT {
         }).verifyComplete();
     }
 
+    @Test
+    public void testFindByIdWithPartition() {
+        final Mono<PartitionPerson> partitionPersonMono = cosmosTemplate.findById(TEST_PERSON.getId(),
+            PartitionPerson.class,
+            new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON)));
+        StepVerifier.create(partitionPersonMono).consumeNextWith(actual -> {
+            Assert.assertThat(actual.getFirstName(), is(equalTo(TEST_PERSON.getFirstName())));
+            Assert.assertThat(actual.getLastName(), is(equalTo(TEST_PERSON.getLastName())));
+        }).verifyComplete();
+    }
+
 //    @Test
 //    public void testFindByNonExistIdWithPartition() {
 //

src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/SimpleDocumentDbRepositoryUnitTest.java

@@ -67,7 +67,7 @@ public class SimpleDocumentDbRepositoryUnitTest {
 
     @Test
     public void testFindOne() {
-        when(dbOperations.findById(anyString(), any(), any())).thenReturn(TEST_PERSON);
+        when(dbOperations.findById(anyString(), anyString(), any())).thenReturn(TEST_PERSON);
 
         repository.save(TEST_PERSON);
 
@@ -82,7 +82,7 @@ public class SimpleDocumentDbRepositoryUnitTest {
 
         repository.save(TEST_PERSON);
 
-        when(dbOperations.findById(anyString(), any(), any()))
+        when(dbOperations.findById(anyString(), anyString(), any()))
                 .thenThrow(new UnsupportedOperationException(PARTITION_VALUE_REQUIRED_MSG));
 
         final Person result = repository.findById(TEST_PERSON.getId()).get();
@@ -98,7 +98,7 @@ public class SimpleDocumentDbRepositoryUnitTest {
                         TestConstants.UPDATED_HOBBIES, updatedAddress);
         repository.save(updatedPerson);
 
-        when(dbOperations.findById(anyString(), any(), any())).thenReturn(updatedPerson);
+        when(dbOperations.findById(anyString(), anyString(), any())).thenReturn(updatedPerson);
 
         final Person result = repository.findById(TEST_PERSON.getId()).get();
         assertEquals(updatedPerson, result);

src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/AddressRepositoryIT.java

@@ -5,6 +5,7 @@
  */
 package com.microsoft.azure.spring.data.cosmosdb.repository.integration;
 
+import com.microsoft.azure.documentdb.PartitionKey;
 import com.microsoft.azure.spring.data.cosmosdb.common.TestConstants;
 import com.microsoft.azure.spring.data.cosmosdb.common.TestUtils;
 import com.microsoft.azure.spring.data.cosmosdb.core.DocumentDbTemplate;
@@ -23,6 +24,7 @@ import javax.annotation.PreDestroy;
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Optional;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -77,6 +79,14 @@ public class AddressRepositoryIT {
         assertThat(result.size()).isEqualTo(4);
     }
 
+    @Test
+    public void testFindByIdWithPartitionKey() {
+        final Optional<Address> addressById = repository.findById(TEST_ADDRESS1_PARTITION1.getPostalCode(),
+            new PartitionKey(entityInformation.getPartitionKeyFieldValue(TEST_ADDRESS1_PARTITION1)));
+
+        assertThat(addressById.equals(TEST_ADDRESS1_PARTITION1));
+    }
+
     @Test
     public void testFindByIdForPartitionedCollection() {
         final List<Address> addresses = repository.findByPostalCode(TestConstants.POSTAL_CODE);

src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/ProjectRepositoryIT.java

@@ -5,6 +5,7 @@
  */
 package com.microsoft.azure.spring.data.cosmosdb.repository.integration;
 
+import com.microsoft.azure.documentdb.PartitionKey;
 import com.microsoft.azure.spring.data.cosmosdb.core.DocumentDbTemplate;
 import com.microsoft.azure.spring.data.cosmosdb.domain.Project;
 import com.microsoft.azure.spring.data.cosmosdb.repository.TestRepositoryConfig;
@@ -25,6 +26,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Optional;
 
 @RunWith(SpringJUnit4ClassRunner.class)
 @ContextConfiguration(classes = TestRepositoryConfig.class)
@@ -346,6 +348,16 @@ public class ProjectRepositoryIT {
         assertProjectListEquals(projects, Arrays.asList(PROJECT_3, PROJECT_4));
     }
 
+    @Test
+    public void findByIdWithPartitionKey() {
+        final Optional<Project> project = repository.findById(PROJECT_0.getId(),
+            new PartitionKey(entityInformation.getPartitionKeyFieldValue(PROJECT_0)));
+
+        Assert.assertTrue(project.isPresent());
+
+        Assert.assertEquals(project.get(), PROJECT_0);
+    }
+
     @Test
     public void testFindByIn() {
         List<Project> projects = repository.findByCreatorIn(Collections.singleton(FAKE_CREATOR));