Add query sort parameter for updateFirst method.

Closes: #4797
Original Pull Request: #4888

Signed-off-by: Florian Lüdiger <florian.luediger@googlemail.com>

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

@@ -184,6 +184,7 @@ import com.mongodb.client.result.UpdateResult;
  * @author Bartłomiej Mazur
  * @author Michael Krog
  * @author Jakub Zurawa
+ * @author Florian Lüdiger
  */
 public class MongoTemplate implements MongoOperations, ApplicationContextAware, IndexOperationsProvider,
 		SearchIndexOperationsProvider, ReadPreferenceAware {
@@ -1700,12 +1701,6 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 		Assert.notNull(query, "Query must not be null");
 		Assert.notNull(update, "Update must not be null");
 
-		if (query.isSorted() && LOGGER.isWarnEnabled()) {
-
-			LOGGER.warn(String.format("%s does not support sort ('%s'); Please use findAndModify() instead",
-					upsert ? "Upsert" : "UpdateFirst", serializeToJsonSafely(query.getSortObject())));
-		}
-
 		MongoPersistentEntity<?> entity = entityClass == null ? null : getPersistentEntity(entityClass);
 
 		UpdateContext updateContext = multi ? queryOperations.updateContext(update, query, upsert)
@@ -1713,7 +1708,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 		updateContext.increaseVersionForUpdateIfNecessary(entity);
 
 		Document queryObj = updateContext.getMappedQuery(entity);
-		UpdateOptions opts = updateContext.getUpdateOptions(entityClass);
+		UpdateOptions opts = updateContext.getUpdateOptions(entityClass, query);
 
 		if (updateContext.isAggregationUpdate()) {
 

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

@@ -77,6 +77,7 @@ import com.mongodb.client.model.UpdateOptions;
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Florian Lüdiger
  * @since 3.0
  */
 class QueryOperations {
@@ -740,7 +741,7 @@ class QueryOperations {
 		/**
 		 * Get the {@link UpdateOptions} applicable for the {@link Query}.
 		 *
-		 * @param domainType must not be {@literal null}.
+		 * @param domainType can be {@literal null}.
 		 * @return never {@literal null}.
 		 */
 		UpdateOptions getUpdateOptions(@Nullable Class<?> domainType) {
@@ -751,11 +752,10 @@ class QueryOperations {
 		 * Get the {@link UpdateOptions} applicable for the {@link Query}.
 		 *
 		 * @param domainType can be {@literal null}.
-		 * @param callback a callback to modify the generated options. Can be {@literal null}.
+		 * @param query can be {@literal null}
-		 * @return
+		 * @return never {@literal null}.
 		 */
-		UpdateOptions getUpdateOptions(@Nullable Class<?> domainType, @Nullable Consumer<UpdateOptions> callback) {
+		UpdateOptions getUpdateOptions(@Nullable Class<?> domainType, @Nullable Query query) {
-
 			UpdateOptions options = new UpdateOptions();
 			options.upsert(upsert);
 
@@ -764,13 +764,13 @@ class QueryOperations {
 						.arrayFilters(update.getArrayFilters().stream().map(ArrayFilter::asDocument).collect(Collectors.toList()));
 			}
 
+			if (query != null && query.isSorted()) {
+				options.sort(query.getSortObject());
+			}
+
 			HintFunction.from(getQuery().getHint()).ifPresent(codecRegistryProvider, options::hintString, options::hint);
 			applyCollation(domainType, options::collation);
 
-			if (callback != null) {
-				callback.accept(options);
-			}
-
 			return options;
 		}
 

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

@@ -183,6 +183,7 @@ import com.mongodb.reactivestreams.client.MongoDatabase;
  * @author Roman Puchkovskiy
  * @author Mathieu Ouellet
  * @author Yadhukrishna S Pai
+ * @author Florian Lüdiger
  * @since 2.0
  */
 public class ReactiveMongoTemplate implements ReactiveMongoOperations, ApplicationContextAware {
@@ -1730,12 +1731,6 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
 	protected Mono<UpdateResult> doUpdate(String collectionName, Query query, @Nullable UpdateDefinition update,
 			@Nullable Class<?> entityClass, boolean upsert, boolean multi) {
 
-		if (query.isSorted() && LOGGER.isWarnEnabled()) {
-
-			LOGGER.warn(String.format("%s does not support sort ('%s'); Please use findAndModify() instead",
-					upsert ? "Upsert" : "UpdateFirst", serializeToJsonSafely(query.getSortObject())));
-		}
-
 		MongoPersistentEntity<?> entity = entityClass == null ? null : getPersistentEntity(entityClass);
 
 		UpdateContext updateContext = multi ? queryOperations.updateContext(update, query, upsert)
@@ -1743,7 +1738,7 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
 		updateContext.increaseVersionForUpdateIfNecessary(entity);
 
 		Document queryObj = updateContext.getMappedQuery(entity);
-		UpdateOptions updateOptions = updateContext.getUpdateOptions(entityClass);
+		UpdateOptions updateOptions = updateContext.getUpdateOptions(entityClass, query);
 
 		Flux<UpdateResult> result;
 

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

@@ -126,6 +126,7 @@ import com.mongodb.client.result.UpdateResult;
  * @author Laszlo Csontos
  * @author duozhilin
  * @author Jakub Zurawa
+ * @author Florian Lüdiger
  */
 @ExtendWith(MongoClientExtension.class)
 public class MongoTemplateTests {
@@ -4065,6 +4066,64 @@ public class MongoTemplateTests {
 		assertThat(loaded.mapValue).isEqualTo(sourceMap);
 	}
 
+	@Test // GH-4797
+	public void updateFirstWithSortingAscendingShouldUpdateCorrectEntities() {
+
+		PersonWithIdPropertyOfTypeObjectId youngPerson = new PersonWithIdPropertyOfTypeObjectId();
+		youngPerson.setId(new ObjectId());
+		youngPerson.setAge(27);
+		youngPerson.setFirstName("Dave");
+		template.save(youngPerson);
+
+		PersonWithIdPropertyOfTypeObjectId oldPerson = new PersonWithIdPropertyOfTypeObjectId();
+		oldPerson.setId(new ObjectId());
+		oldPerson.setAge(34);
+		oldPerson.setFirstName("Dave");
+		template.save(oldPerson);
+
+		template.updateFirst(query(where("firstName").is("Dave")).with(Sort.by(Direction.ASC, "age")),
+				update("firstName", "Mike"), PersonWithIdPropertyOfTypeObjectId.class);
+
+		PersonWithIdPropertyOfTypeObjectId oldPersonResult = template.findById(oldPerson.getId(),
+				PersonWithIdPropertyOfTypeObjectId.class);
+		assertThat(oldPersonResult).isNotNull();
+		assertThat(oldPersonResult.getFirstName()).isEqualTo("Dave");
+
+		PersonWithIdPropertyOfTypeObjectId youngPersonResult = template.findById(youngPerson.getId(),
+				PersonWithIdPropertyOfTypeObjectId.class);
+		assertThat(youngPersonResult).isNotNull();
+		assertThat(youngPersonResult.getFirstName()).isEqualTo("Mike");
+	}
+
+	@Test // GH-4797
+	public void updateFirstWithSortingDescendingShouldUpdateCorrectEntities() {
+
+		PersonWithIdPropertyOfTypeObjectId youngPerson = new PersonWithIdPropertyOfTypeObjectId();
+		youngPerson.setId(new ObjectId());
+		youngPerson.setAge(27);
+		youngPerson.setFirstName("Dave");
+		template.save(youngPerson);
+
+		PersonWithIdPropertyOfTypeObjectId oldPerson = new PersonWithIdPropertyOfTypeObjectId();
+		oldPerson.setId(new ObjectId());
+		oldPerson.setAge(34);
+		oldPerson.setFirstName("Dave");
+		template.save(oldPerson);
+
+		template.updateFirst(query(where("firstName").is("Dave")).with(Sort.by(Direction.DESC, "age")),
+				update("firstName", "Mike"), PersonWithIdPropertyOfTypeObjectId.class);
+
+		PersonWithIdPropertyOfTypeObjectId oldPersonResult = template.findById(oldPerson.getId(),
+				PersonWithIdPropertyOfTypeObjectId.class);
+		assertThat(oldPersonResult).isNotNull();
+		assertThat(oldPersonResult.getFirstName()).isEqualTo("Mike");
+
+		PersonWithIdPropertyOfTypeObjectId youngPersonResult = template.findById(youngPerson.getId(),
+				PersonWithIdPropertyOfTypeObjectId.class);
+		assertThat(youngPersonResult).isNotNull();
+		assertThat(youngPersonResult.getFirstName()).isEqualTo("Dave");
+	}
+
 	private AtomicReference<ImmutableVersioned> createAfterSaveReference() {
 
 		AtomicReference<ImmutableVersioned> saved = new AtomicReference<>();

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

@@ -1846,6 +1846,42 @@ public class ReactiveMongoTemplateTests {
 				.verify();
 	}
 
+	@Test // GH-4797
+	public void updateFirstWithSortingAscendingShouldUpdateCorrectEntities() {
+
+		Person youngPerson = new Person("Dave", 27);
+		Person oldPerson = new Person("Dave", 34);
+
+		template.insertAll(List.of(youngPerson, oldPerson))
+				.then(template.updateFirst(new Query(where("firstName").is("Dave")).with(Sort.by(Direction.ASC, "age")),
+						new Update().set("firstName", "Carter"), Person.class))
+				.flatMapMany(p -> template.find(new Query().with(Sort.by(Direction.ASC, "age")), Person.class))
+				.as(StepVerifier::create) //
+				.consumeNextWith(actual -> {
+					assertThat(actual.getFirstName()).isEqualTo("Carter");
+				}).consumeNextWith(actual -> {
+					assertThat(actual.getFirstName()).isEqualTo("Dave");
+				}).verifyComplete();
+	}
+
+	@Test // GH-4797
+	public void updateFirstWithSortingDescendingShouldUpdateCorrectEntities() {
+
+		Person youngPerson = new Person("Dave", 27);
+		Person oldPerson = new Person("Dave", 34);
+
+		template.insertAll(List.of(youngPerson, oldPerson))
+				.then(template.updateFirst(new Query(where("firstName").is("Dave")).with(Sort.by(Direction.DESC, "age")),
+						new Update().set("firstName", "Carter"), Person.class))
+				.flatMapMany(p -> template.find(new Query().with(Sort.by(Direction.ASC, "age")), Person.class))
+				.as(StepVerifier::create) //
+				.consumeNextWith(actual -> {
+					assertThat(actual.getFirstName()).isEqualTo("Dave");
+				}).consumeNextWith(actual -> {
+					assertThat(actual.getFirstName()).isEqualTo("Carter");
+				}).verifyComplete();
+	}
+
 	private PersonWithAList createPersonWithAList(String firstname, int age) {
 
 		PersonWithAList p = new PersonWithAList();