From 954f52ca7993cc4806b6e82adc50f8d5e3ae9fb9 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mark.paluch@broadcom.com>
Date: Mon, 19 Feb 2024 15:25:25 +0100
Subject: [PATCH] Apply `AggregationOptions.maxTime` to
 `MongoTemplate.aggregateStream()`.

Closes #4644
---
 .../data/mongodb/core/MongoTemplate.java      |  4 +++
 .../mongodb/core/MongoTemplateUnitTests.java  | 31 ++++++++++++++++---
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java
index 1125b516c..5c8c2702b 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java
@@ -2268,6 +2268,10 @@ public class MongoTemplate
 				cursor = hintFunction.apply(mongoDbFactory, cursor::hintString, cursor::hint);
 			}
 
+			if (options.hasExecutionTimeLimit()) {
+				cursor = cursor.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS);
+			}
+
 			Class<?> domainType = aggregation instanceof TypedAggregation typedAggregation ? typedAggregation.getInputType()
 					: null;
 
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 0d3c9545d..ec609db00 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
@@ -48,6 +48,7 @@ import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
+
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationListener;
@@ -118,7 +119,15 @@ import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoCursor;
 import com.mongodb.client.MongoDatabase;
-import com.mongodb.client.model.*;
+import com.mongodb.client.model.CountOptions;
+import com.mongodb.client.model.CreateCollectionOptions;
+import com.mongodb.client.model.DeleteOptions;
+import com.mongodb.client.model.FindOneAndDeleteOptions;
+import com.mongodb.client.model.FindOneAndReplaceOptions;
+import com.mongodb.client.model.FindOneAndUpdateOptions;
+import com.mongodb.client.model.MapReduceAction;
+import com.mongodb.client.model.TimeSeriesGranularity;
+import com.mongodb.client.model.UpdateOptions;
 import com.mongodb.client.result.DeleteResult;
 import com.mongodb.client.result.UpdateResult;
 
@@ -498,6 +507,17 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		verify(collection).withReadPreference(ReadPreference.secondary());
 	}
 
+	@Test // GH-4644
+	void aggregateStreamShouldHonorMaxTimeIfSet() {
+
+		AggregationOptions options = AggregationOptions.builder().maxTime(Duration.ofSeconds(20)).build();
+
+		template.aggregateStream(newAggregation(Aggregation.unwind("foo")).withOptions(options), "collection-1",
+				Wrapper.class);
+
+		verify(aggregateIterable).maxTime(20000, TimeUnit.MILLISECONDS);
+	}
+
 	@Test // DATAMONGO-2153
 	void aggregateShouldHonorOptionsComment() {
 
@@ -2489,7 +2509,8 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 	@Test // GH-4462
 	void replaceShouldUseHintIfPresent() {
 
-		template.replace(new BasicQuery("{}").withHint("index-to-use"), new Sith(), ReplaceOptions.replaceOptions().upsert());
+		template.replace(new BasicQuery("{}").withHint("index-to-use"), new Sith(),
+				ReplaceOptions.replaceOptions().upsert());
 
 		ArgumentCaptor<com.mongodb.client.model.ReplaceOptions> options = ArgumentCaptor
 				.forClass(com.mongodb.client.model.ReplaceOptions.class);
@@ -2509,7 +2530,8 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 			}
 		});
 
-		template.replace(new BasicQuery("{}").withHint("index-to-use"), new Sith(), ReplaceOptions.replaceOptions().upsert());
+		template.replace(new BasicQuery("{}").withHint("index-to-use"), new Sith(),
+				ReplaceOptions.replaceOptions().upsert());
 
 		verify(collection).withWriteConcern(eq(WriteConcern.UNACKNOWLEDGED));
 	}
@@ -2551,8 +2573,7 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		@Id String id;
 		String firstname;
 
-		public Person() {
-		}
+		public Person() {}
 
 		public Person(String id, String firstname) {
 			this.id = id;