Improve SimpleMongoRepository.saveAll to use batch updates

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SimpleMongoRepository.java

@@ -20,7 +20,10 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
@@ -37,7 +40,10 @@
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Window;
+import org.springframework.data.mongodb.core.BulkOperations;
+import org.springframework.data.mongodb.core.BulkOperations.BulkMode;
 import org.springframework.data.mongodb.core.ExecutableFindOperation;
+import org.springframework.data.mongodb.core.FindAndReplaceOptions;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.MongoTemplate;
 import org.springframework.data.mongodb.core.query.Criteria;
@@ -63,6 +69,7 @@
  * @author Mehran Behnam
  * @author Jens Schauder
  * @author Kirill Egorov
+ * @author Sangyeop Jeong
  */
 public class SimpleMongoRepository<T, ID> implements MongoRepository<T, ID> {
 
@@ -106,16 +113,49 @@ public <S extends T> List<S> saveAll(Iterable<S> entities) {
 
 		Assert.notNull(entities, "The given Iterable of entities not be null");
 
-		Streamable<S> source = Streamable.of(entities);
-		boolean allNew = source.stream().allMatch(entityInformation::isNew);
+		List<S> source = Streamable.of(entities).stream().collect(Collectors.toList());
 
-		if (allNew) {
+		if (source.stream().allMatch(entityInformation::isNew)) {
+			return new ArrayList<>(mongoOperations.insert(source, entityInformation.getCollectionName()));
+		}
+
+		if (entityInformation.isVersioned()) {
+			return source.stream().map(this::save).collect(Collectors.toList());
+		}
+
+		List<S> newEntities = new ArrayList<>();
+		List<S> existingEntities = new ArrayList<>();
+
+		for (S entity : source) {
+			(entityInformation.isNew(entity) ? newEntities : existingEntities).add(entity);
+		}
+
+		if (!newEntities.isEmpty()) {
+
+			Iterator<S> inserted = mongoOperations.insert(newEntities, entityInformation.getCollectionName()).iterator();
+			Map<S, S> insertedEntities = new IdentityHashMap<>(newEntities.size());
+
+			for (S entity : newEntities) {
+				insertedEntities.put(entity, inserted.next());
+			}
+
+			source.replaceAll(entity -> insertedEntities.getOrDefault(entity, entity));
+		}
+
+		if (!existingEntities.isEmpty()) {
+
+			BulkOperations bulkOps = mongoOperations.bulkOps(BulkMode.ORDERED, entityInformation.getJavaType(),
+					entityInformation.getCollectionName());
+
+			for (S entity : existingEntities) {
+				bulkOps.replaceOne(getIdQuery(entityInformation.getId(entity)), entity,
+						FindAndReplaceOptions.options().upsert());
+			}
 
-			List<S> result = source.stream().collect(Collectors.toList());
-			return new ArrayList<>(mongoOperations.insert(result, entityInformation.getCollectionName()));
+			bulkOps.execute();
 		}
 
-		return source.stream().map(this::save).collect(Collectors.toList());
+		return source;
 	}
 
 	@Override

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SimpleMongoRepositoryTests.java

@@ -437,6 +437,26 @@ void saveAllUsesEntityCollection() {
 		assertThat(repository.findAll()).containsExactlyInAnyOrder(first, second);
 	}
 
+	@Test // GH-5220
+	@DirtiesState
+	void saveAllUpdatesExistingAndInsertsNewEntities() {
+
+		dave.setFirstname("David");
+
+		Person person = new Person("Dino", "Johnson");
+		person.setId(null);
+
+		List<Person> saved = repository.saveAll(asList(dave, person));
+
+		assertThat(saved).containsExactly(dave, person);
+		assertThat(person.getId()).isNotNull();
+
+		assertThat(repository.findById(dave.getId())) //
+				.hasValueSatisfying(it -> assertThat(it.getFirstname()).isEqualTo("David"));
+		assertThat(repository.findById(person.getId())).contains(person);
+		assertThat(repository.count()).isEqualTo(all.size() + 1);
+	}
+
 	@Test // DATAMONGO-2130
 	@EnableIfReplicaSetAvailable
 	@EnableIfMongoServerVersion(isGreaterThanEqual = "4.0")

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SimpleMongoRepositoryUnitTests.java

@@ -15,10 +15,12 @@
  */
 package org.springframework.data.mongodb.repository.support;
 
+import static java.util.Arrays.*;
 import static org.assertj.core.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
+import java.util.List;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
@@ -34,7 +36,10 @@
 import org.springframework.data.domain.Example;
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Sort;
+import org.springframework.data.mongodb.core.BulkOperations;
+import org.springframework.data.mongodb.core.BulkOperations.BulkMode;
 import org.springframework.data.mongodb.core.ExecutableFindOperation.ExecutableFind;
+import org.springframework.data.mongodb.core.FindAndReplaceOptions;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.query.Collation;
 import org.springframework.data.mongodb.core.query.Query;
@@ -198,6 +203,73 @@ void shouldAddReadPreferenceToFluentFetchable() throws NoSuchMethodException {
 		assertThat(query.getValue().getReadPreference()).isEqualTo(com.mongodb.ReadPreference.secondaryPreferred());
 	}
 
+	@Test // GH-5220
+	void saveAllUsesBulkOperationsForExistingNonVersionedEntities() {
+
+		when(entityInformation.isVersioned()).thenReturn(false);
+		when(entityInformation.getJavaType()).thenReturn(Object.class);
+		when(entityInformation.getCollectionName()).thenReturn("persons");
+		when(entityInformation.getIdAttribute()).thenReturn("id");
+
+		Object existing = new Object();
+		Object fresh = new Object();
+
+		when(entityInformation.isNew(existing)).thenReturn(false);
+		when(entityInformation.isNew(fresh)).thenReturn(true);
+		when(entityInformation.getId(existing)).thenReturn("id-1");
+
+		BulkOperations bulkOperations = mock(BulkOperations.class);
+		when(mongoOperations.bulkOps(BulkMode.ORDERED, Object.class, "persons")).thenReturn(bulkOperations);
+		when(mongoOperations.insert(List.of(fresh), "persons")).thenReturn(List.of(fresh));
+
+		repository.saveAll(asList(existing, fresh));
+
+		verify(mongoOperations).insert(List.of(fresh), "persons");
+		verify(bulkOperations).replaceOne(any(Query.class), eq(existing), any(FindAndReplaceOptions.class));
+		verify(bulkOperations).execute();
+		verify(mongoOperations, never()).save(any(), anyString());
+	}
+
+	@Test // GH-5220
+	void saveAllReturnsInsertedInstancesInOriginalOrder() {
+
+		when(entityInformation.isVersioned()).thenReturn(false);
+		when(entityInformation.getJavaType()).thenReturn(Object.class);
+		when(entityInformation.getCollectionName()).thenReturn("persons");
+		when(entityInformation.getIdAttribute()).thenReturn("id");
+
+		Object existing = new Object();
+		Object fresh = new Object();
+		Object insertedFresh = new Object();
+
+		when(entityInformation.isNew(existing)).thenReturn(false);
+		when(entityInformation.isNew(fresh)).thenReturn(true);
+		when(entityInformation.getId(existing)).thenReturn("id-1");
+
+		BulkOperations bulkOperations = mock(BulkOperations.class);
+		when(mongoOperations.bulkOps(BulkMode.ORDERED, Object.class, "persons")).thenReturn(bulkOperations);
+		when(mongoOperations.insert(List.of(fresh), "persons")).thenReturn(List.of(insertedFresh));
+
+		List<Object> saved = repository.saveAll(asList(existing, fresh));
+
+		assertThat(saved).containsExactly(existing, insertedFresh);
+	}
+
+	@Test // GH-5220
+	void saveAllFallsBackToPerEntitySaveForVersionedEntities() {
+
+		when(entityInformation.isVersioned()).thenReturn(true);
+		when(entityInformation.getCollectionName()).thenReturn("persons");
+
+		Object existing = new Object();
+		when(entityInformation.isNew(existing)).thenReturn(false);
+
+		repository.saveAll(List.of(existing));
+
+		verify(mongoOperations).save(existing, "persons");
+		verify(mongoOperations, never()).bulkOps(any(BulkMode.class), any(Class.class), anyString());
+	}
+
 	private static Stream<Arguments> findAllCalls() {
 
 		Consumer<SimpleMongoRepository<Object, Object>> findAll = SimpleMongoRepository::findAll;