Improve SimpleMongoRepository.saveAll to use batch updates

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/BulkWriter.java

@@ -47,6 +47,7 @@
  * Internal API wrapping a {@link MongoTemplate} to encapsulate {@link Bulk} handling.
  *
  * @author Christoph Strobl
+ * @author Sangyeop Jeong
  * @since 5.1
  */
 class BulkWriter extends BulkWriterSupport {
@@ -82,11 +83,8 @@ private BulkWriteResult writeToSingleCollection(String defaultDatabase, Bulk bul
 					collection -> collection.bulkWrite(collector.getWriteModels(), new com.mongodb.client.model.BulkWriteOptions()
 							.ordered(options.getOrder().equals(BulkWriteOptions.Order.ORDERED))));
 
-			collector.getAfterSaveCallables().forEach(callable -> {
-				template
-						.maybeEmitEvent(new AfterSaveEvent<>(callable.source(), callable.document(), callable.collectionName()));
-				template.maybeCallAfterSave(callable.source(), callable.document(), callable.collectionName());
-			});
+			collector.getAfterSaveCallables().forEach(this::completeSave);
+
 			return BulkWriteResult.from(bulkWriteResult);
 		} catch (MongoBulkWriteException e) {
 			DataAccessException dataAccessException = template.getExceptionTranslator().translateExceptionIfPossible(e);
@@ -110,11 +108,8 @@ private BulkWriteResult writeToMultipleCollections(String defaultDatabase, Bulk
 					.doWithClient(client -> client.bulkWrite(collector.getWriteModels(), ClientBulkWriteOptions
 							.clientBulkWriteOptions().ordered(options.getOrder().equals(BulkWriteOptions.Order.ORDERED))));
 
-			collector.getAfterSaveCallables().forEach(callable -> {
-				template
-						.maybeEmitEvent(new AfterSaveEvent<>(callable.source(), callable.document(), callable.collectionName()));
-				template.maybeCallAfterSave(callable.source(), callable.document(), callable.collectionName());
-			});
+			collector.getAfterSaveCallables().forEach(this::completeSave);
+
 			return BulkWriteResult.from(clientBulkWriteResult);
 		} catch (MongoBulkWriteException e) {
 			DataAccessException dataAccessException = template.getExceptionTranslator().translateExceptionIfPossible(e);
@@ -174,4 +169,21 @@ private void buildWriteModels(Bulk bulk, WriteModelCollector collector) {
 		}
 	}
 
+	/**
+	 * Completes the save lifecycle after an entity has been written through an {@literal insert} or {@literal replace}
+	 * operation by propagating a generated identifier back to the entity and emitting the after save event and
+	 * callbacks.
+	 *
+	 * @param written the entity along with the document handed to the driver, carrying an identifier generated during
+	 *          the write.
+	 */
+	private void completeSave(SourceAwareDocument<Object> written) {
+
+		Object entity = populateIdIfNecessary(written.source(), written.document(),
+				template.getConverter().getConversionService());
+
+		template.maybeEmitEvent(new AfterSaveEvent<>(entity, written.document(), written.collectionName()));
+		template.maybeCallAfterSave(entity, written.document(), written.collectionName());
+	}
+
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/BulkWriterSupport.java

@@ -23,6 +23,7 @@
 import org.bson.Document;
 import org.jspecify.annotations.Nullable;
 
+import org.springframework.core.convert.ConversionService;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.bulk.Bulk;
 import org.springframework.data.mongodb.core.bulk.BulkOperation;
@@ -44,6 +45,7 @@
  *
  * @author Mark Paluch
  * @author Christoph Strobl
+ * @author Sangyeop Jeong
  * @since 5.1
  */
 abstract class BulkWriterSupport {
@@ -75,6 +77,22 @@ String resolveCollectionName(TypedNamespace namespace) {
 		return entityOperations.determineCollectionName(namespace.type());
 	}
 
+	/**
+	 * Propagates an identifier generated during the write back to the entity the write was issued for. Entities that
+	 * already carry an identifier remain untouched.
+	 *
+	 * @param source the entity the write was issued for.
+	 * @param document the document handed to the server, potentially carrying a generated {@literal _id}.
+	 * @param conversionService used to adapt the identifier to the id property type.
+	 * @return the entity carrying the identifier. Can be a different instance for immutable types.
+	 */
+	<T> T populateIdIfNecessary(T source, Document document, ConversionService conversionService) {
+
+		Object id = MappedDocument.of(document).getId();
+
+		return id != null ? entityOperations.forEntity(source, conversionService).populateIdIfNecessary(id) : source;
+	}
+
 	@Nullable
 	@SuppressWarnings("unchecked")
 	MongoPersistentEntity<?> getPersistentEntity(BulkOperationContext context) {

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveBulkWriter.java

@@ -50,6 +50,7 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Sangyeop Jeong
  * @since 5.1
  */
 class ReactiveBulkWriter extends BulkWriterSupport {
@@ -80,21 +81,13 @@ private Mono<BulkWriteResult> writeToSingleCollection(String defaultDatabase, Bu
 		return buildWriteModelsReactive(bulk, collector).then(Mono.defer(() -> {
 
 			String collectionName = collector.getNamespace().getCollectionName();
-			List<SourceAwareDocument<Object>> afterSaveCallables = collector.getAfterSaveCallables();
 
 			return template
 					.createMono(collectionName,
 							col -> col.bulkWrite(collector.getWriteModels(),
 									new com.mongodb.client.model.BulkWriteOptions()
 											.ordered(options.getOrder().equals(BulkWriteOptions.Order.ORDERED))))
-					.map(BulkWriteResult::from)
-					.doOnSuccess(
-							v -> afterSaveCallables
-									.forEach(callable -> template.maybeEmitEvent(new AfterSaveEvent<>(callable.source(),
-											callable.document(), callable.collectionName()))))
-					.flatMap(result -> Flux.concat(afterSaveCallables.stream().map(callable -> template
-							.maybeCallAfterSave(callable.source(), callable.document(), callable.collectionName())).toList())
-							.then(Mono.just(result)));
+					.flatMap(result -> completeSaves(collector).thenReturn(BulkWriteResult.from(result)));
 		}));
 	}
 
@@ -106,20 +99,12 @@ private Mono<BulkWriteResult> writeToMultipleCollections(String defaultDatabase,
 		return buildWriteModelsReactive(bulk, collector).then(Mono.defer(() -> {
 
 			List<ClientNamespacedWriteModel> writeModels = collector.getWriteModels();
-			List<SourceAwareDocument<Object>> afterSaveCallables = collector.getAfterSaveCallables();
 
 			return template
 					.doWithCluster(client -> client.bulkWrite(writeModels,
 							ClientBulkWriteOptions
 									.clientBulkWriteOptions().ordered(options.getOrder().equals(BulkWriteOptions.Order.ORDERED))))
-					.map(BulkWriteResult::from)
-					.doOnSuccess(
-							v -> afterSaveCallables
-									.forEach(callable -> template.maybeEmitEvent(new AfterSaveEvent<>(callable.source(),
-											callable.document(), callable.collectionName()))))
-					.flatMap(result -> Flux.concat(afterSaveCallables.stream().map(callable -> template
-							.maybeCallAfterSave(callable.source(), callable.document(), callable.collectionName())).toList())
-							.then(Mono.just(result)));
+					.flatMap(result -> completeSaves(collector).thenReturn(BulkWriteResult.from(result)));
 		}));
 	}
 
@@ -186,6 +171,28 @@ private Mono<Void> addOperationReactive(BulkOperation bulkOp, WriteModelCollecto
 		return Mono.error(new IllegalStateException("Unknown bulk operation type: " + bulkOp.getClass()));
 	}
 
+	private Mono<Void> completeSaves(WriteModelCollector collector) {
+		return Flux.fromIterable(collector.getAfterSaveCallables()).concatMap(this::completeSave).then();
+	}
+
+	/**
+	 * Completes the save lifecycle after an entity has been written through an {@literal insert} or {@literal replace}
+	 * operation by propagating a generated identifier back to the entity and emitting the after save event and
+	 * callbacks.
+	 *
+	 * @param written the entity along with the document handed to the driver, carrying an identifier generated during
+	 *          the write.
+	 * @return the entity as returned by the after save callbacks.
+	 */
+	private Mono<Object> completeSave(SourceAwareDocument<Object> written) {
+
+		Object entity = populateIdIfNecessary(written.source(), written.document(),
+				template.getConverter().getConversionService());
+
+		template.maybeEmitEvent(new AfterSaveEvent<>(entity, written.document(), written.collectionName()));
+		return template.maybeCallAfterSave(entity, written.document(), written.collectionName());
+	}
+
 	@SuppressWarnings("unchecked")
 	private static SourceAwareDocument<Object> toObject(SourceAwareDocument<?> sad) {
 		return (SourceAwareDocument<Object>) sad;

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

@@ -40,6 +40,9 @@
 import org.springframework.data.mongodb.core.ExecutableFindOperation;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.bulk.Bulk;
+import org.springframework.data.mongodb.core.bulk.BulkWriteOptions;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.query.Criteria;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.repository.MongoRepository;
@@ -63,6 +66,7 @@
  * @author Mehran Behnam
  * @author Jens Schauder
  * @author Kirill Egorov
+ * @author Sangyeop Jeong
  */
 public class SimpleMongoRepository<T, ID> implements MongoRepository<T, ID> {
 
@@ -106,16 +110,52 @@ 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.isEmpty()) {
+			return source;
+		}
 
-			List<S> result = source.stream().collect(Collectors.toList());
-			return new ArrayList<>(mongoOperations.insert(result, entityInformation.getCollectionName()));
+		// bulk writes re-initialize rather than increment @Version and cannot attribute a conflict to a single entity
+		if (source.stream().anyMatch(this::isVersionedEntity)) {
+			return source.stream().map(this::save).collect(Collectors.toList());
 		}
 
-		return source.stream().map(this::save).collect(Collectors.toList());
+		mongoOperations.bulkWrite(createSaveBulk(source), BulkWriteOptions.ordered());
+
+		return source;
+	}
+
+	/**
+	 * Returns whether the given entity declares a version property. Resolves the {@link MongoPersistentEntity} for the
+	 * actual entity type as a repository can be declared for a supertype that does not declare a version property while
+	 * the instance at hand does.
+	 *
+	 * @param entity the entity to inspect.
+	 * @return {@literal true} if the entity type declares a version property.
+	 */
+	private boolean isVersionedEntity(Object entity) {
+
+		MongoPersistentEntity<?> persistentEntity = mongoOperations.getConverter().getMappingContext()
+				.getPersistentEntity(entity.getClass());
+
+		return persistentEntity != null && persistentEntity.hasVersionProperty();
+	}
+
+	private <S extends T> Bulk createSaveBulk(List<S> source) {
+
+		return Bulk.create(builder -> builder.inCollection(entityInformation.getJavaType(),
+				entityInformation.getCollectionName(), spec -> {
+
+					for (S entity : source) {
+
+						if (entityInformation.isNew(entity)) {
+							spec.insert(entity);
+						} else {
+							spec.replaceOne(getIdQuery(entityInformation.getId(entity)), entity);
+						}
+					}
+				}));
 	}
 
 	@Override

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

@@ -45,6 +45,9 @@
 import org.springframework.data.mongodb.core.MongoTemplate;
 import org.springframework.data.mongodb.core.ReactiveFindOperation;
 import org.springframework.data.mongodb.core.ReactiveMongoOperations;
+import org.springframework.data.mongodb.core.bulk.Bulk;
+import org.springframework.data.mongodb.core.bulk.BulkWriteOptions;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.query.Criteria;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.repository.ReactiveMongoRepository;
@@ -65,6 +68,7 @@
  * @author Jens Schauder
  * @author Clément Petit
  * @author Kirill Egorov
+ * @author Sangyeop Jeong
  * @since 2.0
  */
 public class SimpleReactiveMongoRepository<T, ID extends Serializable> implements ReactiveMongoRepository<T, ID> {
@@ -112,8 +116,50 @@ public <S extends T> Flux<S> saveAll(Iterable<S> entities) {
 		Assert.notNull(entities, "The given Iterable of entities must not be null");
 
 		List<S> source = toList(entities);
-		return source.stream().allMatch(entityInformation::isNew) ? //
-				insert(source) : concatMapSequentially(source, this::save);
+
+		if (source.isEmpty()) {
+			return Flux.empty();
+		}
+
+		// bulk writes re-initialize rather than increment @Version and cannot attribute a conflict to a single entity
+		if (source.stream().anyMatch(this::isVersionedEntity)) {
+			return concatMapSequentially(source, this::save);
+		}
+
+		return mongoOperations.bulkWrite(createSaveBulk(source), BulkWriteOptions.ordered())
+				.thenMany(Flux.fromIterable(source));
+	}
+
+	/**
+	 * Returns whether the given entity declares a version property. Resolves the {@link MongoPersistentEntity} for the
+	 * actual entity type as a repository can be declared for a supertype that does not declare a version property while
+	 * the instance at hand does.
+	 *
+	 * @param entity the entity to inspect.
+	 * @return {@literal true} if the entity type declares a version property.
+	 */
+	private boolean isVersionedEntity(Object entity) {
+
+		MongoPersistentEntity<?> persistentEntity = mongoOperations.getConverter().getMappingContext()
+				.getPersistentEntity(entity.getClass());
+
+		return persistentEntity != null && persistentEntity.hasVersionProperty();
+	}
+
+	private <S extends T> Bulk createSaveBulk(List<S> source) {
+
+		return Bulk.create(builder -> builder.inCollection(entityInformation.getJavaType(),
+				entityInformation.getCollectionName(), spec -> {
+
+					for (S entity : source) {
+
+						if (entityInformation.isNew(entity)) {
+							spec.insert(entity);
+						} else {
+							spec.replaceOne(getIdQuery(entityInformation.getId(entity)), entity);
+						}
+					}
+				}));
 	}
 
 	@Override

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateBulkTests.java

@@ -23,6 +23,7 @@
 import java.util.List;
 
 import org.bson.Document;
+import org.bson.types.ObjectId;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -48,6 +49,7 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Sangyeop Jeong
  */
 @EnableIfMongoServerVersion(isGreaterThanEqual = "8.0")
 class MongoTemplateBulkTests {
@@ -88,6 +90,21 @@ void bulkWriteMultipleCollections() {
 		assertThat(inSpecialCollection).isOne();
 	}
 
+	@Test // GH-5220
+	void bulkInsertPropagatesGeneratedIdToEntity() {
+
+		BaseDoc doc = new BaseDoc();
+		doc.value = "value-doc";
+
+		operations.bulkWrite(Bulk.create(builder -> builder.inCollection(BaseDoc.class, ops -> ops.insert(doc))),
+				BulkWriteOptions.ordered());
+
+		ObjectId storedId = operations.execute(BaseDoc.class,
+				collection -> collection.find().first().getObjectId("_id"));
+
+		assertThat(doc.id).isNotNull().isEqualTo(storedId.toHexString());
+	}
+
 	@Test // GH-5087
 	void bulkWriteRawDocument() {
 

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")