Added new upsert API which returns updated entity (#501)

microsoft · Feb 26, 2020 · e9db481 · e9db481
1 parent 7defda2
commit e9db481
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 38 deletions.
diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosOperations.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosOperations.java
@@ -58,6 +58,8 @@ public interface CosmosOperations {
 
     <T> void upsert(String containerName, T object, PartitionKey partitionKey);
 
+    <T> T upsertAndReturnEntity(String containerName, T object, PartitionKey partitionKey);
+
     void deleteById(String containerName, Object id, PartitionKey partitionKey);
 
     void deleteAll(String containerName, Class<?> domainType);

diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosTemplate.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosTemplate.java
@@ -180,16 +180,23 @@ public <T> void upsert(T object, PartitionKey partitionKey) {
     }
 
     public <T> void upsert(String containerName, T object, PartitionKey partitionKey) {
+        upsertAndReturnEntity(containerName, object,  partitionKey);
+    }
+
+    public <T> T upsertAndReturnEntity(String containerName, T object, PartitionKey partitionKey) {
         Assert.hasText(containerName, "containerName should not be null, empty or only whitespaces");
         Assert.notNull(object, "Upsert object should not be null");
 
         final CosmosItemProperties originalItem = mappingCosmosConverter.writeCosmosItemProperties(object);
 
         log.debug("execute upsert item in database {} container {}", this.databaseName, containerName);
 
+        @SuppressWarnings("unchecked")
+        final Class<T> domainType = (Class<T>) object.getClass();
+
         final CosmosItemRequestOptions options = new CosmosItemRequestOptions();
         options.partitionKey(partitionKey);
-        applyVersioning(object.getClass(), originalItem, options);
+        applyVersioning(domainType, originalItem, options);
 
         final CosmosItemResponse cosmosItemResponse = cosmosClient
             .getDatabase(this.databaseName)
@@ -201,6 +208,7 @@ public <T> void upsert(String containerName, T object, PartitionKey partitionKey
             .block();
 
         assert cosmosItemResponse != null;
+        return mappingCosmosConverter.read(domainType, cosmosItemResponse.properties());
     }
 
     public <T> List<T> findAll(Class<T> domainType) {

diff --git a/...a/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleCosmosRepository.java b/...a/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleCosmosRepository.java
@@ -22,6 +22,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;
@@ -72,11 +73,9 @@ public <S extends T> S save(S entity) {
                     entity,
                     createKey(information.getPartitionKeyFieldValue(entity)));
         } else {
-            operation.upsert(information.getContainerName(),
+            return operation.upsertAndReturnEntity(information.getContainerName(),
                     entity, createKey(information.getPartitionKeyFieldValue(entity)));
         }
-
-        return entity;
     }
 
     private PartitionKey createKey(String partitionKeyValue) {
@@ -98,9 +97,13 @@ private PartitionKey createKey(String partitionKeyValue) {
     public <S extends T> Iterable<S> saveAll(Iterable<S> entities) {
         Assert.notNull(entities, "Iterable entities should not be null");
 
-        entities.forEach(this::save);
+        final List<S> 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/CosmosTemplateIT.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosTemplateIT.java
@@ -177,21 +177,31 @@ public void testUpsertNewDocument() {
         final Person newPerson = new Person(TEST_PERSON.getId(), firstName,
                 NEW_FIRST_NAME, null, null);
 
-        cosmosTemplate.upsert(Person.class.getSimpleName(), newPerson,
-                new PartitionKey(personInfo.getPartitionKeyFieldValue(newPerson)));
+        final Person person = cosmosTemplate.upsertAndReturnEntity(Person.class.getSimpleName(),
+            newPerson,
+            new PartitionKey(personInfo.getPartitionKeyFieldValue(newPerson)));
 
         assertThat(responseDiagnosticsTestUtils.getCosmosResponseDiagnostics()).isNotNull();
         assertThat(responseDiagnosticsTestUtils.getFeedResponseDiagnostics()).isNull();
         assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNull();
 
-        final List<Person> result = cosmosTemplate.findAll(Person.class);
+        assertEquals(person.getFirstName(), firstName);
+    }
 
-        assertThat(responseDiagnosticsTestUtils.getFeedResponseDiagnostics()).isNotNull();
-        assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
-        assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
+    @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());
 
-        assertThat(result.size()).isEqualTo(1);
-        assertEquals(result.get(0).getFirstName(), firstName);
+        final Person updatedPerson = cosmosTemplate.upsertAndReturnEntity(Person.class.getSimpleName(),
+            updated, null);
+
+        final Person findPersonById = cosmosTemplate.findById(Person.class.getSimpleName(),
+            updatedPerson.getId(), Person.class);
+
+        assertEquals(updatedPerson, updated);
+        assertThat(updatedPerson.get_etag()).isEqualTo(findPersonById.get_etag());
     }
 
     @Test
@@ -200,20 +210,14 @@ public void testUpdate() {
                 TEST_PERSON.getLastName(), TEST_PERSON.getHobbies(), TEST_PERSON.getShippingAddresses());
         updated.set_etag(insertedPerson.get_etag());
 
-        cosmosTemplate.upsert(Person.class.getSimpleName(), updated, null);
+        final Person person = cosmosTemplate.upsertAndReturnEntity(Person.class.getSimpleName(),
+            updated, null);
 
         assertThat(responseDiagnosticsTestUtils.getCosmosResponseDiagnostics()).isNotNull();
         assertThat(responseDiagnosticsTestUtils.getFeedResponseDiagnostics()).isNull();
         assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNull();
 
-        final Person result = cosmosTemplate.findById(Person.class.getSimpleName(),
-                updated.getId(), Person.class);
-
-        assertThat(responseDiagnosticsTestUtils.getFeedResponseDiagnostics()).isNotNull();
-        assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
-        assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
-
-        assertEquals(result, updated);
+        assertEquals(person, updated);
     }
 
     @Test

diff --git a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosTemplatePartitionIT.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/CosmosTemplatePartitionIT.java
@@ -41,7 +41,6 @@
 import static com.microsoft.azure.spring.data.cosmosdb.core.query.CriteriaType.IS_EQUAL;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
 
 @RunWith(SpringJUnit4ClassRunner.class)
 @ContextConfiguration(classes = TestRepositoryConfig.class)
@@ -135,28 +134,25 @@ public void testUpsertNewDocumentPartition() {
             null, null);
 
         final String partitionKeyValue = newPerson.getLastName();
-        cosmosTemplate.upsert(PartitionPerson.class.getSimpleName(), newPerson, new PartitionKey(partitionKeyValue));
+        final PartitionPerson partitionPerson =
+            cosmosTemplate.upsertAndReturnEntity(PartitionPerson.class.getSimpleName(), newPerson
+                , new PartitionKey(partitionKeyValue));
 
         final List<PartitionPerson> result = cosmosTemplate.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());
-        cosmosTemplate.upsert(PartitionPerson.class.getSimpleName(), updated, new PartitionKey(updated.getLastName()));
-
-        final List<PartitionPerson> result = cosmosTemplate.findAll(PartitionPerson.class);
-        final PartitionPerson person = result.stream().filter(
-                p -> TEST_PERSON.getId().equals(p.getId())).findFirst().get();
+        final PartitionPerson partitionPerson =
+            cosmosTemplate.upsertAndReturnEntity(PartitionPerson.class.getSimpleName(), updated,
+                new PartitionKey(updated.getLastName()));
 
-        assertTrue(person.equals(updated));
+        assertEquals(partitionPerson, updated);
     }
 
     @Test
@@ -174,7 +170,7 @@ public void testDeleteByIdPartition() {
 
         final List<PartitionPerson> result = cosmosTemplate.findAll(PartitionPerson.class);
         assertThat(result.size()).isEqualTo(1);
-        assertTrue(result.get(0).equals(TEST_PERSON_2));
+        assertEquals(result.get(0), TEST_PERSON_2);
     }
 
     @Test

diff --git a/.../com/microsoft/azure/spring/data/cosmosdb/repository/integration/ContactRepositoryIT.java b/.../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,7 +116,11 @@ public void testCountAndDeleteEntity() {
     public void testUpdateEntity() {
         final Contact updatedContact = new Contact(TEST_CONTACT.getLogicId(), "updated");
 
-        repository.save(updatedContact);
+        final Contact savedContact = repository.save(updatedContact);
+
+        //  Test save operation return saved entity
+        assertThat(savedContact.getLogicId()).isEqualTo(updatedContact.getLogicId());
+        assertThat(savedContact.getTitle()).isEqualTo(updatedContact.getTitle());
 
         final Contact contact = repository.findById(TEST_CONTACT.getLogicId()).get();
 
@@ -131,7 +136,15 @@ public void testBatchOperations() {
         final ArrayList<Contact> contacts = new ArrayList<Contact>();
         contacts.add(contact1);
         contacts.add(contact2);
-        repository.saveAll(contacts);
+        final Iterable<Contact> 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<String> ids = new ArrayList<String>();
         ids.add(contact1.getLogicId());