From ec55e28cb3b03f7a310a04bc0867f1fb8802abcd Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@vmware.com>
Date: Fri, 28 Apr 2023 08:23:05 +0200
Subject: [PATCH] Fix missing query options when calling MongoOperations#count.

This commit makes sure to forward maxTimeMsec and comment options from the query to the CountOptions.

Closes: #4374
Original pull request: #4378
---
 .../data/mongodb/core/QueryOperations.java    | 12 +++++++++++
 .../mongodb/core/MongoTemplateUnitTests.java  | 20 ++++++++++++++++++
 .../core/ReactiveMongoTemplateUnitTests.java  | 21 +++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java
index 114a21788..c3acbd994 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java
@@ -21,6 +21,7 @@ import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -66,6 +67,7 @@ import com.mongodb.client.model.CountOptions;
 import com.mongodb.client.model.DeleteOptions;
 import com.mongodb.client.model.ReplaceOptions;
 import com.mongodb.client.model.UpdateOptions;
+import org.springframework.util.StringUtils;
 
 /**
  * {@link QueryOperations} centralizes common operations required before an operation is actually ready to be executed.
@@ -567,6 +569,16 @@ class QueryOperations {
 			if (query.getSkip() > 0) {
 				options.skip((int) query.getSkip());
 			}
+
+			if(query.getMeta().hasValues()) {
+				if(query.getMeta().getMaxTimeMsec() != null && query.getMeta().getMaxTimeMsec() > 0) {
+					options.maxTime(query.getMeta().getMaxTimeMsec(), TimeUnit.MILLISECONDS);
+				}
+				if(StringUtils.hasText(query.getMeta().getComment())) {
+					options.comment(query.getMeta().getComment());
+				}
+			}
+
 			if (StringUtils.hasText(query.getHint())) {
 
 				String hint = query.getHint();
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
index 2944a1d08..24ef90653 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
@@ -2307,6 +2307,26 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		verify(collection).countDocuments(any(Document.class), any());
 	}
 
+	@Test // GH-4374
+	void countConsidersMaxTimeMs() {
+
+		template.count(new BasicQuery("{ 'spring' : 'data-mongodb' }").maxTimeMsec(5000), Human.class);
+
+		ArgumentCaptor<CountOptions> options = ArgumentCaptor.forClass(CountOptions.class);
+		verify(collection).countDocuments(any(Document.class), options.capture());
+		assertThat(options.getValue().getMaxTime(TimeUnit.MILLISECONDS)).isEqualTo(5000);
+	}
+
+	@Test // GH-4374
+	void countPassesOnComment() {
+
+		template.count(new BasicQuery("{ 'spring' : 'data-mongodb' }").comment("rocks!"), Human.class);
+
+		ArgumentCaptor<CountOptions> options = ArgumentCaptor.forClass(CountOptions.class);
+		verify(collection).countDocuments(any(Document.class), options.capture());
+		assertThat(options.getValue().getComment()).isEqualTo(BsonUtils.simpleToBsonValue("rocks!"));
+	}
+
 	@Test // GH-3984
 	void templatePassesOnTimeSeriesOptionsWhenNoTypeGiven() {
 
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java
index 2f9ce1ac8..c2543ade8 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java
@@ -23,6 +23,7 @@ import static org.springframework.data.mongodb.test.util.Assertions.assertThat;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.NoArgsConstructor;
+import org.springframework.data.mongodb.util.BsonUtils;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -1456,6 +1457,26 @@ public class ReactiveMongoTemplateUnitTests {
 		verify(collection).countDocuments(any(Document.class), any());
 	}
 
+	@Test // GH-4374
+	void countConsidersMaxTimeMs() {
+
+		template.count(new BasicQuery("{ 'spring' : 'data-mongodb' }").maxTimeMsec(5000), Person.class).subscribe();
+
+		ArgumentCaptor<CountOptions> options = ArgumentCaptor.forClass(CountOptions.class);
+		verify(collection).countDocuments(any(Document.class), options.capture());
+		assertThat(options.getValue().getMaxTime(TimeUnit.MILLISECONDS)).isEqualTo(5000);
+	}
+
+	@Test // GH-4374
+	void countPassesOnComment() {
+
+		template.count(new BasicQuery("{ 'spring' : 'data-mongodb' }").comment("rocks!"), Person.class).subscribe();
+
+		ArgumentCaptor<CountOptions> options = ArgumentCaptor.forClass(CountOptions.class);
+		verify(collection).countDocuments(any(Document.class), options.capture());
+		assertThat(options.getValue().getComment()).isEqualTo(BsonUtils.simpleToBsonValue("rocks!"));
+	}
+
 	@Test // GH-2911
 	void insertErrorsOnPublisher() {