From 1ec063432120f6d59f2f37dae47259621759dd08 Mon Sep 17 00:00:00 2001 From: Kushagra Thapar Date: Tue, 25 Feb 2020 17:26:48 -0800 Subject: [PATCH] Added upsert and return entity api which fixes save and save all API in documentDBRepository (#502) --- .../cosmosdb/core/DocumentDbOperations.java | 2 ++ .../cosmosdb/core/DocumentDbTemplate.java | 21 +++++++++---- .../support/SimpleDocumentDbRepository.java | 13 +++++--- .../cosmosdb/core/DocumentDbTemplateIT.java | 31 +++++++++++++------ .../core/DocumentDbTemplatePartitionIT.java | 18 +++++------ .../integration/ContactRepositoryIT.java | 19 ++++++++---- 6 files changed, 67 insertions(+), 37 deletions(-) diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java index e36007c3..14e4072d 100644 --- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java +++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java @@ -46,6 +46,8 @@ public interface DocumentDbOperations { void upsert(String collectionName, T object, PartitionKey partitionKey); + T upsertAndReturnEntity(String collectionName, T object, PartitionKey partitionKey); + void deleteById(String collectionName, Object id, PartitionKey partitionKey); void deleteAll(String collectionName, Class domainClass); diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java index 81e30133..802ad103 100644 --- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java +++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java @@ -204,6 +204,10 @@ public void upsert(T object, PartitionKey partitionKey) { } public void upsert(String collectionName, T object, PartitionKey partitionKey) { + upsertAndReturnEntity(collectionName, object, partitionKey); + } + + public T upsertAndReturnEntity(String collectionName, T object, PartitionKey partitionKey) { Assert.hasText(collectionName, "collectionName should not be null, empty or only whitespaces"); Assert.notNull(object, "Upsert object should not be null"); @@ -212,19 +216,24 @@ public void upsert(String collectionName, T object, PartitionKey partitionKe log.debug("execute upsert document in database {} collection {}", this.databaseName, collectionName); + @SuppressWarnings("unchecked") + final Class domainClass = (Class) object.getClass(); + final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); options.partitionKey(toCosmosPartitionKey(partitionKey)); - applyVersioning(object.getClass(), originalItem, options); + applyVersioning(domainClass, originalItem, options); - final CosmosItemResponse cosmosItemResponse = cosmosClient.getDatabase(this.databaseName) - .getContainer(collectionName) - .upsertItem(originalItem, options) - .onErrorResume(this::databaseAccessExceptionHandler) - .block(); + final CosmosItemResponse cosmosItemResponse = cosmosClient + .getDatabase(this.databaseName) + .getContainer(collectionName) + .upsertItem(originalItem, options) + .onErrorResume(this::databaseAccessExceptionHandler) + .block(); if (cosmosItemResponse == null) { throw new DocumentDBAccessException("Failed to upsert item"); } + return toDomainObject(domainClass, cosmosItemResponse.properties()); } catch (Exception ex) { throw new DocumentDBAccessException("Failed to upsert document to database.", ex); } diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java index a500fdd8..cdf9f50b 100644 --- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java +++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java @@ -23,6 +23,7 @@ import org.springframework.util.StringUtils; import java.io.Serializable; +import java.util.ArrayList; import java.util.List; import java.util.Optional; import java.util.stream.StreamSupport; @@ -78,11 +79,9 @@ public S save(S entity) { entity, createKey(information.getPartitionKeyFieldValue(entity))); } else { - operation.upsert(information.getCollectionName(), + return operation.upsertAndReturnEntity(information.getCollectionName(), entity, createKey(information.getPartitionKeyFieldValue(entity))); } - - return entity; } private PartitionKey createKey(String partitionKeyValue) { @@ -104,9 +103,13 @@ private PartitionKey createKey(String partitionKeyValue) { public Iterable saveAll(Iterable entities) { Assert.notNull(entities, "Iterable entities should not be null"); - entities.forEach(this::save); + final List savedEntities = new ArrayList<>(); + entities.forEach(entity -> { + final S savedEntity = this.save(entity); + savedEntities.add(savedEntity); + }); - return entities; + return savedEntities; } /** diff --git a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplateIT.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplateIT.java index 9f1e2636..905b4cdf 100644 --- a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplateIT.java +++ b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplateIT.java @@ -172,13 +172,26 @@ public void testUpsertNewDocument() { final Person newPerson = new Person(TEST_PERSON.getId(), firstName, NEW_FIRST_NAME, null, null); - dbTemplate.upsert(Person.class.getSimpleName(), newPerson, - new PartitionKey(personInfo.getPartitionKeyFieldValue(newPerson))); + final Person person = dbTemplate.upsertAndReturnEntity(Person.class.getSimpleName(), newPerson, + new PartitionKey(personInfo.getPartitionKeyFieldValue(newPerson))); - final List result = dbTemplate.findAll(Person.class); + assertEquals(person.getFirstName(), firstName); + } - assertThat(result.size()).isEqualTo(1); - assertEquals(result.get(0).getFirstName(), firstName); + @Test + public void testUpdateWithReturnEntity() { + final Person updated = new Person(TEST_PERSON.getId(), UPDATED_FIRST_NAME, + TEST_PERSON.getLastName(), TEST_PERSON.getHobbies(), TEST_PERSON.getShippingAddresses()); + updated.set_etag(insertedPerson.get_etag()); + + final Person updatedPerson = dbTemplate.upsertAndReturnEntity(Person.class.getSimpleName(), + updated, null); + + final Person findPersonById = dbTemplate.findById(Person.class.getSimpleName(), + updatedPerson.getId(), Person.class); + + assertEquals(updatedPerson, updated); + assertThat(updatedPerson.get_etag()).isEqualTo(findPersonById.get_etag()); } @Test @@ -187,12 +200,10 @@ public void testUpdate() { TEST_PERSON.getLastName(), TEST_PERSON.getHobbies(), TEST_PERSON.getShippingAddresses()); updated.set_etag(insertedPerson.get_etag()); - dbTemplate.upsert(Person.class.getSimpleName(), updated, null); - - final Person result = dbTemplate.findById(Person.class.getSimpleName(), - updated.getId(), Person.class); + final Person person = dbTemplate.upsertAndReturnEntity(Person.class.getSimpleName(), updated, + null); - assertEquals(result, updated); + assertEquals(person, updated); } @Test diff --git a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java index e5d5afb9..f622a01e 100644 --- a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java +++ b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java @@ -138,28 +138,26 @@ public void testUpsertNewDocumentPartition() { null, null); final String partitionKeyValue = newPerson.getLastName(); - dbTemplate.upsert(PartitionPerson.class.getSimpleName(), newPerson, new PartitionKey(partitionKeyValue)); + final PartitionPerson partitionPerson = + dbTemplate.upsertAndReturnEntity(PartitionPerson.class.getSimpleName(), newPerson, + new PartitionKey(partitionKeyValue)); final List result = dbTemplate.findAll(PartitionPerson.class); assertThat(result.size()).isEqualTo(2); - final PartitionPerson person = result.stream() - .filter(p -> p.getLastName().equals(partitionKeyValue)).findFirst().get(); - assertThat(person.getFirstName()).isEqualTo(firstName); + assertThat(partitionPerson.getFirstName()).isEqualTo(firstName); } @Test public void testUpdatePartition() { final PartitionPerson updated = new PartitionPerson(TEST_PERSON.getId(), UPDATED_FIRST_NAME, TEST_PERSON.getLastName(), TEST_PERSON.getHobbies(), TEST_PERSON.getShippingAddresses()); - dbTemplate.upsert(PartitionPerson.class.getSimpleName(), updated, new PartitionKey(updated.getLastName())); + final PartitionPerson partitionPerson = + dbTemplate.upsertAndReturnEntity(PartitionPerson.class.getSimpleName(), updated, + new PartitionKey(updated.getLastName())); - final List result = dbTemplate.findAll(PartitionPerson.class); - final PartitionPerson person = result.stream().filter( - p -> TEST_PERSON.getId().equals(p.getId())).findFirst().get(); - - assertTrue(person.equals(updated)); + assertTrue(partitionPerson.equals(updated)); } @Test diff --git a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/ContactRepositoryIT.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/ContactRepositoryIT.java index fe67b5ea..bb4401c1 100644 --- a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/ContactRepositoryIT.java +++ b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/ContactRepositoryIT.java @@ -26,6 +26,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -115,12 +116,10 @@ public void testCountAndDeleteEntity() { public void testUpdateEntity() { final Contact updatedContact = new Contact(TEST_CONTACT.getLogicId(), "updated"); - repository.save(updatedContact); + final Contact savedContact = repository.save(updatedContact); - final Contact contact = repository.findById(TEST_CONTACT.getLogicId()).get(); - - assertThat(contact.getLogicId()).isEqualTo(updatedContact.getLogicId()); - assertThat(contact.getTitle()).isEqualTo(updatedContact.getTitle()); + assertThat(savedContact.getLogicId()).isEqualTo(updatedContact.getLogicId()); + assertThat(savedContact.getTitle()).isEqualTo(updatedContact.getTitle()); } @Test @@ -131,7 +130,15 @@ public void testBatchOperations() { final ArrayList contacts = new ArrayList(); contacts.add(contact1); contacts.add(contact2); - repository.saveAll(contacts); + final Iterable savedContacts = repository.saveAll(contacts); + + final AtomicInteger savedCount = new AtomicInteger(); + savedContacts.forEach(se -> { + savedCount.incrementAndGet(); + assertThat(contacts.contains(se)).isTrue(); + }); + + assertThat(savedCount.get()).isEqualTo(contacts.size()); final ArrayList ids = new ArrayList(); ids.add(contact1.getLogicId());