Apply sort to replace and bulk operation updates

Allow using sort parameter from the query for template replace as well as bulk update & replace operations.
We now also mapped fields used in sort to the domain type considering field annotations.
Also updated javadoc and reference documentation.

Original Pull Request: #4888

Author: christophstrobl

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

@@ -17,6 +17,7 @@
 
 import java.util.List;
 
+import org.springframework.data.domain.Sort;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.data.mongodb.core.query.UpdateDefinition;
@@ -81,7 +82,8 @@ enum BulkMode {
 	/**
 	 * Add a single update to the bulk operation. For the update request, only the first matching document is updated.
 	 *
-	 * @param query update criteria, must not be {@literal null}.
+	 * @param query update criteria, must not be {@literal null}. The {@link Query} may define a {@link Query#with(Sort)
+	 *          sort order} to influence which document to update when potentially matching multiple candidates.
 	 * @param update {@link Update} operation to perform, must not be {@literal null}.
 	 * @return the current {@link BulkOperations} instance with the update added, will never be {@literal null}.
 	 */
@@ -92,7 +94,8 @@ default BulkOperations updateOne(Query query, Update update) {
 	/**
 	 * Add a single update to the bulk operation. For the update request, only the first matching document is updated.
 	 *
-	 * @param query update criteria, must not be {@literal null}.
+	 * @param query update criteria, must not be {@literal null}. The {@link Query} may define a {@link Query#with(Sort)
+	 *          sort order} to influence which document to update when potentially matching multiple candidates.
 	 * @param update {@link Update} operation to perform, must not be {@literal null}.
 	 * @return the current {@link BulkOperations} instance with the update added, will never be {@literal null}.
 	 * @since 4.1
@@ -187,7 +190,8 @@ default BulkOperations upsert(Query query, Update update) {
 	/**
 	 * Add a single replace operation to the bulk operation.
 	 *
-	 * @param query Update criteria.
+	 * @param query Replace criteria. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence
+	 *          which document to replace when potentially matching multiple candidates.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @return the current {@link BulkOperations} instance with the replacement added, will never be {@literal null}.
 	 * @since 2.2
@@ -199,7 +203,8 @@ default BulkOperations replaceOne(Query query, Object replacement) {
 	/**
 	 * Add a single replace operation to the bulk operation.
 	 *
-	 * @param query Update criteria.
+	 * @param query Replace criteria. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence
+	 *          which document to replace when potentially matching multiple candidates.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @param options the {@link FindAndModifyOptions} holding additional information. Must not be {@literal null}.
 	 * @return the current {@link BulkOperations} instance with the replacement added, will never be {@literal null}.

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

@@ -106,6 +106,11 @@ WriteModel<Document> mapWriteModel(Object source, WriteModel<Document> writeMode
 
 		if (writeModel instanceof UpdateOneModel<Document> model) {
 
+			Bson sort = model.getOptions().getSort();
+			if (sort instanceof Document sortDocument) {
+				model.getOptions().sort(updateMapper().getMappedSort(sortDocument, entity().orElse(null)));
+			}
+
 			if (source instanceof AggregationUpdate aggregationUpdate) {
 
 				List<Document> pipeline = mapUpdatePipeline(aggregationUpdate);
@@ -136,6 +141,17 @@ WriteModel<Document> mapWriteModel(Object source, WriteModel<Document> writeMode
 			return new DeleteManyModel<>(getMappedQuery(model.getFilter()), model.getOptions());
 		}
 
+		if (writeModel instanceof ReplaceOneModel<Document> model) {
+
+			Bson sort = model.getReplaceOptions().getSort();
+
+			if (sort instanceof Document sortDocument) {
+				model.getReplaceOptions().sort(updateMapper().getMappedSort(sortDocument, entity().orElse(null)));
+			}
+			return new ReplaceOneModel<>(getMappedQuery(model.getFilter()), model.getReplacement(),
+					model.getReplaceOptions());
+		}
+
 		return writeModel;
 	}
 
@@ -192,9 +208,11 @@ protected static BulkWriteOptions getBulkWriteOptions(BulkMode bulkMode) {
 	 * @param filterQuery The {@link Query} to read a potential {@link Collation} from. Must not be {@literal null}.
 	 * @param update The {@link Update} to apply
 	 * @param upsert flag to indicate if document should be upserted.
+	 * @param multi flag to indicate if update might affect multiple documents.
 	 * @return new instance of {@link UpdateOptions}.
 	 */
-	protected static UpdateOptions computeUpdateOptions(Query filterQuery, UpdateDefinition update, boolean upsert) {
+	protected UpdateOptions computeUpdateOptions(Query filterQuery, UpdateDefinition update, boolean upsert,
+			boolean multi) {
 
 		UpdateOptions options = new UpdateOptions();
 		options.upsert(upsert);
@@ -207,6 +225,10 @@ protected static UpdateOptions computeUpdateOptions(Query filterQuery, UpdateDef
 			options.arrayFilters(list);
 		}
 
+		if (!multi && filterQuery.isSorted()) {
+			options.sort(filterQuery.getSortObject());
+		}
+
 		filterQuery.getCollation().map(Collation::toMongoCollation).ifPresent(options::collation);
 		return options;
 	}

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

@@ -229,12 +229,14 @@ public BulkOperations replaceOne(Query query, Object replacement, FindAndReplace
 
 		ReplaceOptions replaceOptions = new ReplaceOptions();
 		replaceOptions.upsert(options.isUpsert());
+		if (query.isSorted()) {
+			replaceOptions.sort(query.getSortObject());
+		}
 		query.getCollation().map(Collation::toMongoCollation).ifPresent(replaceOptions::collation);
 
 		maybeEmitEvent(new BeforeConvertEvent<>(replacement, collectionName));
 		Object source = maybeInvokeBeforeConvertCallback(replacement);
-		addModel(source,
-				new ReplaceOneModel<>(getMappedQuery(query.getQueryObject()), getMappedObject(source), replaceOptions));
+		addModel(source, new ReplaceOneModel<>(query.getQueryObject(), getMappedObject(source), replaceOptions));
 
 		return this;
 	}
@@ -315,7 +317,7 @@ private BulkOperations update(Query query, UpdateDefinition update, boolean upse
 		Assert.notNull(query, "Query must not be null");
 		Assert.notNull(update, "Update must not be null");
 
-		UpdateOptions options = computeUpdateOptions(query, update, upsert);
+		UpdateOptions options = computeUpdateOptions(query, update, upsert, multi);
 
 		if (multi) {
 			addModel(update, new UpdateManyModel<>(query.getQueryObject(), update.getUpdateObject(), options));

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

@@ -189,13 +189,16 @@ public ReactiveBulkOperations replaceOne(Query query, Object replacement, FindAn
 
 		ReplaceOptions replaceOptions = new ReplaceOptions();
 		replaceOptions.upsert(options.isUpsert());
+		if (query.isSorted()) {
+			replaceOptions.sort(query.getSortObject());
+		}
 		query.getCollation().map(Collation::toMongoCollation).ifPresent(replaceOptions::collation);
 
 		this.models.add(Mono.just(replacement).flatMap(it -> {
 			maybeEmitEvent(new BeforeConvertEvent<>(it, collectionName));
 			return maybeInvokeBeforeConvertCallback(it);
 		}).map(it -> new SourceAwareWriteModelHolder(it,
-				new ReplaceOneModel<>(getMappedQuery(query.getQueryObject()), getMappedObject(it), replaceOptions))));
+				new ReplaceOneModel<>(query.getQueryObject(), getMappedObject(it), replaceOptions))));
 
 		return this;
 	}
@@ -218,13 +221,13 @@ private Mono<BulkWriteResult> bulkWriteTo(MongoCollection<Document> collection)
 
 		Flux<SourceAwareWriteModelHolder> concat = Flux.concat(models).flatMapSequential(it -> {
 
-			if (it.model()instanceof InsertOneModel<Document> iom) {
+			if (it.model() instanceof InsertOneModel<Document> iom) {
 
 				Document target = iom.getDocument();
 				maybeEmitBeforeSaveEvent(it);
 				return maybeInvokeBeforeSaveCallback(it.source(), target)
 						.map(afterCallback -> new SourceAwareWriteModelHolder(afterCallback, mapWriteModel(afterCallback, iom)));
-			} else if (it.model()instanceof ReplaceOneModel<Document> rom) {
+			} else if (it.model() instanceof ReplaceOneModel<Document> rom) {
 
 				Document target = rom.getReplacement();
 				maybeEmitBeforeSaveEvent(it);
@@ -265,7 +268,7 @@ private ReactiveBulkOperations update(Query query, UpdateDefinition update, bool
 		Assert.notNull(query, "Query must not be null");
 		Assert.notNull(update, "Update must not be null");
 
-		UpdateOptions options = computeUpdateOptions(query, update, upsert);
+		UpdateOptions options = computeUpdateOptions(query, update, upsert, multi);
 
 		this.models.add(Mono.just(update).map(it -> {
 			if (multi) {

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

@@ -25,6 +25,7 @@
 
 import org.bson.Document;
 import org.springframework.data.domain.KeysetScrollPosition;
+import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Window;
 import org.springframework.data.geo.GeoResults;
 import org.springframework.data.mongodb.core.BulkOperations.BulkMode;
@@ -1604,8 +1605,9 @@ default long exactCount(Query query, String collectionName) {
 	 * A potential {@link org.springframework.data.annotation.Version} property of the {@literal entityClass} will be
 	 * auto-incremented if not explicitly specified in the update.
 	 *
-	 * @param query the query document that specifies the criteria used to select a document to be updated. Must not be
-	 *          {@literal null}.
+	 * @param query the query document that specifies the criteria used to select a document to be updated. The
+	 *          {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to update when
+	 *          potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param update the {@link UpdateDefinition} that contains the updated object or {@code $} operators to manipulate
 	 *          the existing. Must not be {@literal null}.
 	 * @param entityClass class that determines the collection to use.
@@ -1623,12 +1625,11 @@ default long exactCount(Query query, String collectionName) {
 	 * the provided updated document. <br />
 	 * <strong>NOTE:</strong> Any additional support for field mapping, versions, etc. is not available due to the lack of
 	 * domain type information. Use {@link #updateFirst(Query, UpdateDefinition, Class, String)} to get full type specific
-	 * support. <br />
-	 * <strong>NOTE:</strong> {@link Query#getSortObject() sorting} is not supported by {@code db.collection.updateOne}.
-	 * Use {@link #findAndModify(Query, UpdateDefinition, Class, String)} instead.
+	 * support.
 	 *
-	 * @param query the query document that specifies the criteria used to select a document to be updated. Must not be
-	 *          {@literal null}.
+	 * @param query the query document that specifies the criteria used to select a document to be updated. The
+	 *          {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to update when
+	 *          potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param update the {@link UpdateDefinition} that contains the updated object or {@code $} operators to manipulate
 	 *          the existing. Must not be {@literal null}.
 	 * @param collectionName name of the collection to update the object in. Must not be {@literal null}.
@@ -1646,8 +1647,9 @@ default long exactCount(Query query, String collectionName) {
 	 * A potential {@link org.springframework.data.annotation.Version} property of the {@literal entityClass} will be auto
 	 * incremented if not explicitly specified in the update.
 	 *
-	 * @param query the query document that specifies the criteria used to select a document to be updated. Must not be
-	 *          {@literal null}.
+	 * @param query the query document that specifies the criteria used to select a document to be updated. The
+	 *          {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to update when
+	 *          potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param update the {@link UpdateDefinition} that contains the updated object or {@code $} operators to manipulate
 	 *          the existing. Must not be {@literal null}.
 	 * @param entityClass class of the pojo to be operated on. Must not be {@literal null}.
@@ -1833,7 +1835,8 @@ default long exactCount(Query query, String collectionName) {
 	 *
 	 * @param query the {@link Query} class that specifies the {@link Criteria} used to find a document. The query may
 	 *          contain an index {@link Query#withHint(String) hint} or the {@link Query#collation(Collation) collation}
-	 *          to use. Must not be {@literal null}.
+	 *          to use. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to
+	 *          replace when potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @return the {@link UpdateResult} which lets you access the results of the previous replacement.
 	 * @throws org.springframework.data.mapping.MappingException if the collection name cannot be
@@ -1850,7 +1853,8 @@ default <T> UpdateResult replace(Query query, T replacement) {
 	 *
 	 * @param query the {@link Query} class that specifies the {@link Criteria} used to find a document. The query may
 	 *          contain an index {@link Query#withHint(String) hint} or the {@link Query#collation(Collation) collation}
-	 *          to use. Must not be {@literal null}.
+	 *          to use. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to
+	 *          replace when potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @param collectionName the collection to query. Must not be {@literal null}.
 	 * @return the {@link UpdateResult} which lets you access the results of the previous replacement.
@@ -1866,7 +1870,8 @@ default <T> UpdateResult replace(Query query, T replacement, String collectionNa
 	 *
 	 * @param query the {@link Query} class that specifies the {@link Criteria} used to find a document.The query may
 	 *          contain an index {@link Query#withHint(String) hint} or the {@link Query#collation(Collation) collation}
-	 *          to use. Must not be {@literal null}.
+	 *          to use. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to
+	 *          replace when potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @param options the {@link ReplaceOptions} holding additional information. Must not be {@literal null}.
 	 * @return the {@link UpdateResult} which lets you access the results of the previous replacement.
@@ -1884,7 +1889,8 @@ default <T> UpdateResult replace(Query query, T replacement, ReplaceOptions opti
 	 *
 	 * @param query the {@link Query} class that specifies the {@link Criteria} used to find a document. The query may *
 	 *          contain an index {@link Query#withHint(String) hint} or the {@link Query#collation(Collation) collation}
-	 *          to use. Must not be {@literal null}.
+	 *          to use. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence which document to
+	 *          replace when potentially matching multiple candidates. Must not be {@literal null}.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @param options the {@link ReplaceOptions} holding additional information. Must not be {@literal null}.
 	 * @return the {@link UpdateResult} which lets you access the results of the previous replacement.

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

@@ -765,7 +765,7 @@ UpdateOptions getUpdateOptions(@Nullable Class<?> domainType, @Nullable Query qu
 			}
 
 			if (query != null && query.isSorted()) {
-				options.sort(query.getSortObject());
+				options.sort(getMappedSort(domainType != null ? mappingContext.getPersistentEntity(domainType) : null));
 			}
 
 			HintFunction.from(getQuery().getHint()).ifPresent(codecRegistryProvider, options::hintString, options::hint);
@@ -799,6 +799,9 @@ ReplaceOptions getReplaceOptions(@Nullable Class<?> domainType, @Nullable Consum
 			options.collation(updateOptions.getCollation());
 			options.upsert(updateOptions.isUpsert());
 			applyHint(options::hintString, options::hint);
+			if (!isMulti() && getQuery().isSorted()) {
+				options.sort(getMappedSort(domainType != null ? mappingContext.getPersistentEntity(domainType) : null));
+			}
 
 			if (callback != null) {
 				callback.accept(options);

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

@@ -19,6 +19,7 @@
 
 import java.util.List;
 
+import org.springframework.data.domain.Sort;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.UpdateDefinition;
 
@@ -67,7 +68,8 @@ public interface ReactiveBulkOperations {
 	/**
 	 * Add a single update to the bulk operation. For the update request, only the first matching document is updated.
 	 *
-	 * @param query update criteria, must not be {@literal null}.
+	 * @param query update criteria, must not be {@literal null}. The {@link Query} may define a {@link Query#with(Sort)
+	 *          sort order} to influence which document to update when potentially matching multiple candidates.
 	 * @param update {@link UpdateDefinition} operation to perform, must not be {@literal null}.
 	 * @return the current {@link ReactiveBulkOperations} instance with the update added, will never be {@literal null}.
 	 */
@@ -111,8 +113,11 @@ public interface ReactiveBulkOperations {
 	/**
 	 * Add a single replace operation to the bulk operation.
 	 *
-	 * @param query Update criteria.
-	 * @param replacement the replacement document. Must not be {@literal null}.
+	 * @param query Replace criteria. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence
+	 *          which document to replace when potentially matching multiple candidates.
+	 * @param replacement the replacement document. Must not be {@literal null}. The {@link Query} may define a
+	 *          {@link Query#with(Sort) sort order} to influence which document to replace when potentially matching
+	 *          multiple candidates.
 	 * @return the current {@link ReactiveBulkOperations} instance with the replace added, will never be {@literal null}.
 	 */
 	default ReactiveBulkOperations replaceOne(Query query, Object replacement) {
@@ -122,7 +127,8 @@ default ReactiveBulkOperations replaceOne(Query query, Object replacement) {
 	/**
 	 * Add a single replace operation to the bulk operation.
 	 *
-	 * @param query Update criteria.
+	 * @param query Replace criteria. The {@link Query} may define a {@link Query#with(Sort) sort order} to influence
+	 *          which document to replace when potentially matching multiple candidates.
 	 * @param replacement the replacement document. Must not be {@literal null}.
 	 * @param options the {@link FindAndModifyOptions} holding additional information. Must not be {@literal null}.
 	 * @return the current {@link ReactiveBulkOperations} instance with the replace added, will never be {@literal null}.