DATAMONGO-2390 - Add support for maxTimeMS to AggregationOptions.

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-mongodb-parent</artifactId>
-	<version>2.3.0.BUILD-SNAPSHOT</version>
+	<version>2.3.0.DATAMONGO-2390-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data MongoDB</name>

spring-data-mongodb-benchmarks/pom.xml

@@ -7,7 +7,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>2.3.0.BUILD-SNAPSHOT</version>
+		<version>2.3.0.DATAMONGO-2390-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb-distribution/pom.xml

@@ -14,7 +14,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>2.3.0.BUILD-SNAPSHOT</version>
+		<version>2.3.0.DATAMONGO-2390-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb/pom.xml

@@ -11,7 +11,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>2.3.0.BUILD-SNAPSHOT</version>
+		<version>2.3.0.DATAMONGO-2390-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

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

@@ -2190,6 +2190,10 @@ protected <O> AggregationResults<O> doAggregate(Aggregation aggregation, String
 
 			options.getComment().ifPresent(aggregateIterable::comment);
 
+			if(options.hasExecutionTimeLimit()) {
+				aggregateIterable = aggregateIterable.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS);
+			}
+
 			MongoIterable<O> iterable = aggregateIterable.map(val -> {
 
 				rawResult.add(val);

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

@@ -1073,6 +1073,10 @@ private <O> Flux<O> aggregateAndMap(MongoCollection<Document> collection, List<D
 				.map(Collation::toMongoCollation) //
 				.ifPresent(cursor::collation);
 
+		if(options.hasExecutionTimeLimit()) {
+			cursor = cursor.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS);
+		}
+
 		return Flux.from(cursor).map(readCallback::doWith);
 	}
 

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

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import java.time.Duration;
 import java.util.Optional;
 
 import org.bson.Document;
@@ -45,12 +46,14 @@ public class AggregationOptions {
 	private static final String ALLOW_DISK_USE = "allowDiskUse";
 	private static final String COLLATION = "collation";
 	private static final String COMMENT = "comment";
+	private static final String MAX_TIME = "maxTimeMS";
 
 	private final boolean allowDiskUse;
 	private final boolean explain;
 	private final Optional<Document> cursor;
 	private final Optional<Collation> collation;
 	private final Optional<String> comment;
+	private Duration maxTime = Duration.ZERO;
 
 	/**
 	 * Creates a new {@link AggregationOptions}.
@@ -129,7 +132,11 @@ public static AggregationOptions fromDocument(Document document) {
 				: null;
 		String comment = document.getString(COMMENT);
 
-		return new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
+		AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
+		if (document.containsKey(MAX_TIME)) {
+			options.maxTime = Duration.ofMillis(document.getLong(MAX_TIME));
+		}
+		return options;
 	}
 
 	/**
@@ -206,6 +213,14 @@ public Optional<String> getComment() {
 		return comment;
 	}
 
+	/**
+	 * @return the time limit for processing. {@link Duration#ZERO} is used for the default unbounded behavior.
+	 * @since 2.3
+	 */
+	public Duration getMaxTime() {
+		return maxTime;
+	}
+
 	/**
 	 * Returns a new potentially adjusted copy for the given {@code aggregationCommandObject} with the configuration
 	 * applied.
@@ -233,6 +248,10 @@ Document applyAndReturnPotentiallyChangedCommand(Document command) {
 			collation.map(Collation::toDocument).ifPresent(val -> result.append(COLLATION, val));
 		}
 
+		if (hasExecutionTimeLimit() && !result.containsKey(MAX_TIME)) {
+			result.append(MAX_TIME, maxTime.toMillis());
+		}
+
 		return result;
 	}
 
@@ -251,9 +270,17 @@ public Document toDocument() {
 		collation.ifPresent(val -> document.append(COLLATION, val.toDocument()));
 		comment.ifPresent(val -> document.append(COMMENT, val));
 
+		if (hasExecutionTimeLimit()) {
+			document.append(MAX_TIME, maxTime.toMillis());
+		}
+
 		return document;
 	}
 
+	public boolean hasExecutionTimeLimit() {
+		return !maxTime.isZero() && !maxTime.isNegative();
+	}
+
 	/* (non-Javadoc)
 	 * @see java.lang.Object#toString()
 	 */
@@ -279,6 +306,7 @@ public static class Builder {
 		private @Nullable Document cursor;
 		private @Nullable Collation collation;
 		private @Nullable String comment;
+		private @Nullable Duration maxTime;
 
 		/**
 		 * Defines whether to off-load intensive sort-operations to disk.
@@ -355,13 +383,33 @@ public Builder comment(@Nullable String comment) {
 			return this;
 		}
 
+		/**
+		 * Set the time limit for processing.
+		 *
+		 * @param maxTime {@link Duration#ZERO} is used for the default unbounded behavior. {@link Duration#isNegative()
+		 *          Negative} values will be ignored.
+		 * @return this.
+		 * @sinve 2.3
+		 */
+		public Builder maxTime(@Nullable Duration maxTime) {
+
+			this.maxTime = maxTime;
+			return this;
+		}
+
 		/**
 		 * Returns a new {@link AggregationOptions} instance with the given configuration.
 		 *
 		 * @return
 		 */
 		public AggregationOptions build() {
-			return new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
+
+			AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
+			if (maxTime != null) {
+				options.maxTime = maxTime;
+			}
+
+			return options;
 		}
 	}
 }

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

@@ -22,11 +22,11 @@
 
 import org.bson.Document;
 import org.bson.conversions.Bson;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.dao.DataAccessException;
 import org.springframework.data.geo.Point;
 import org.springframework.data.mapping.MappingException;
@@ -49,7 +49,7 @@
  * @author Thomas Darimont
  * @author Christoph Strobl
  */
-@RunWith(MockitoJUnitRunner.class)
+@ExtendWith(MockitoExtension.class)
 public abstract class MongoOperationsUnitTests {
 
 	@Mock CollectionCallback<Object> collectionCallback;
@@ -59,7 +59,7 @@ public abstract class MongoOperationsUnitTests {
 	Person person;
 	List<Person> persons;
 
-	@Before
+	@BeforeEach
 	public final void operationsSetUp() {
 
 		person = new Person("Oliver");

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

@@ -60,6 +60,11 @@ public MongoClient mongoClient() {
 		protected String getDatabaseName() {
 			return "collation-tests";
 		}
+
+		@Override
+		protected boolean autoIndexCreation() {
+			return false;
+		}
 	}
 
 	@Autowired MongoTemplate template;

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

@@ -85,6 +85,11 @@ protected String getDatabaseName() {
 			return DB_NAME;
 		}
 
+		@Override
+		protected boolean autoIndexCreation() {
+			return false;
+		}
+
 		@Bean
 		MongoTransactionManager txManager(MongoDbFactory dbFactory) {
 			return new MongoTransactionManager(dbFactory);