Added upsert and return entity api which fixes save and save all API …

…in documentDBRepository (#502)
microsoft · Feb 26, 2020 · 1ec0634 · 1ec0634
1 parent 4c432bd
commit 1ec0634
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 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
@@ -46,6 +46,8 @@ public interface DocumentDbOperations {
 
     <T> void upsert(String collectionName, T object, PartitionKey partitionKey);
 
+    <T> 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
@@ -204,6 +204,10 @@ public <T> void upsert(T object, PartitionKey partitionKey) {
     }
 
     public <T> void upsert(String collectionName, T object, PartitionKey partitionKey) {
+        upsertAndReturnEntity(collectionName, object, partitionKey);
+    }
+
+    public <T> 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 <T> void upsert(String collectionName, T object, PartitionKey partitionKe
 
             log.debug("execute upsert document in database {} collection {}", this.databaseName, collectionName);
 
+            @SuppressWarnings("unchecked")
+            final Class<T> domainClass = (Class<T>) 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/...m/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java b/...m/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 extends T> 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 <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/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<Person> 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/...est/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java b/...est/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<PartitionPerson> 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<PartitionPerson> 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/.../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,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<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());