DATAMONGO-1836 - Add support to hint in aggregation options.

Original pull request: #878.

Author: thekaleidoscope

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

@@ -155,6 +155,7 @@
  * @author Cimon Lucas
  * @author Michael J. Simons
  * @author Roman Puchkovskiy
+ * @author Yadhukrishna S Pai
  */
 public class MongoTemplate implements MongoOperations, ApplicationContextAware, IndexOperationsProvider {
 
@@ -2149,6 +2150,8 @@ protected <O> AggregationResults<O> doAggregate(Aggregation aggregation, String
 
 			options.getComment().ifPresent(aggregateIterable::comment);
 
+			options.getHint().ifPresent(aggregateIterable::hint);
+
 			if (options.hasExecutionTimeLimit()) {
 				aggregateIterable = aggregateIterable.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS);
 			}
@@ -2207,6 +2210,8 @@ protected <O> CloseableIterator<O> aggregateStream(Aggregation aggregation, Stri
 
 			options.getComment().ifPresent(cursor::comment);
 
+			options.getHint().ifPresent(cursor::hint);
+
 			Class<?> domainType = aggregation instanceof TypedAggregation ? ((TypedAggregation) aggregation).getInputType()
 					: null;
 

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

@@ -146,6 +146,7 @@
  * @author Christoph Strobl
  * @author Roman Puchkovskiy
  * @author Mathieu Ouellet
+ * @author Yadhukrishna S Pai
  * @since 2.0
  */
 public class ReactiveMongoTemplate implements ReactiveMongoOperations, ApplicationContextAware {
@@ -1024,6 +1025,8 @@ private <O> Flux<O> aggregateAndMap(MongoCollection<Document> collection, List<D
 
 		options.getComment().ifPresent(cursor::comment);
 
+		options.getHint().ifPresent(cursor::hint);
+
 		Optionals.firstNonEmpty(options::getCollation, () -> operations.forType(inputType).getCollation()) //
 				.map(Collation::toMongoCollation) //
 				.ifPresent(cursor::collation);

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

@@ -32,6 +32,7 @@
  * @author Oliver Gierke
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Yadhukrishna S Pai
  * @see Aggregation#withOptions(AggregationOptions)
  * @see TypedAggregation#withOptions(AggregationOptions)
  * @since 1.6
@@ -45,12 +46,14 @@ public class AggregationOptions {
 	private static final String COLLATION = "collation";
 	private static final String COMMENT = "comment";
 	private static final String MAX_TIME = "maxTimeMS";
+	private static final String HINT = "hint";
 
 	private final boolean allowDiskUse;
 	private final boolean explain;
 	private final Optional<Document> cursor;
 	private final Optional<Collation> collation;
 	private final Optional<String> comment;
+	private final Optional<Document> hint;
 	private Duration maxTime = Duration.ZERO;
 	private ResultOptions resultOptions = ResultOptions.READ;
 
@@ -77,7 +80,7 @@ public AggregationOptions(boolean allowDiskUse, boolean explain, Document cursor
 	 */
 	public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Document cursor,
 			@Nullable Collation collation) {
-		this(allowDiskUse, explain, cursor, collation, null);
+		this(allowDiskUse, explain, cursor, collation, null, null);
 	}
 
 	/**
@@ -89,16 +92,18 @@ public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Docum
 	 * aggregation.
 	 * @param collation collation for string comparison. Can be {@literal null}.
 	 * @param comment execution comment. Can be {@literal null}.
+	 * @param hint can be {@literal null}, used to provide an index that would be forcibly used by query optimizer.
 	 * @since 2.2
 	 */
 	public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Document cursor,
-			@Nullable Collation collation, @Nullable String comment) {
+			@Nullable Collation collation, @Nullable String comment, @Nullable Document hint) {
 
 		this.allowDiskUse = allowDiskUse;
 		this.explain = explain;
 		this.cursor = Optional.ofNullable(cursor);
 		this.collation = Optional.ofNullable(collation);
 		this.comment = Optional.ofNullable(comment);
+		this.hint = Optional.ofNullable(hint);
 	}
 
 	/**
@@ -130,8 +135,9 @@ public static AggregationOptions fromDocument(Document document) {
 		Collation collation = document.containsKey(COLLATION) ? Collation.from(document.get(COLLATION, Document.class))
 				: null;
 		String comment = document.getString(COMMENT);
+		Document hint = document.get(HINT, Document.class);
 
-		AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
+		AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment, hint);
 		if (document.containsKey(MAX_TIME)) {
 			options.maxTime = Duration.ofMillis(document.getLong(MAX_TIME));
 		}
@@ -212,6 +218,16 @@ public Optional<String> getComment() {
 		return comment;
 	}
 
+	/**
+	 * Get the hint used to to fulfill the aggregation.
+	 *
+	 * @return never {@literal null}.
+	 */
+	public Optional<Document> getHint() {
+		return hint;
+	}
+
+
 	/**
 	 * @return the time limit for processing. {@link Duration#ZERO} is used for the default unbounded behavior.
 	 * @since 3.0
@@ -248,6 +264,10 @@ Document applyAndReturnPotentiallyChangedCommand(Document command) {
 			result.put(EXPLAIN, explain);
 		}
 
+		if (result.containsKey(HINT)) {
+			hint.ifPresent(val -> result.append(HINT, val));
+		}
+
 		if (!result.containsKey(CURSOR)) {
 			cursor.ifPresent(val -> result.put(CURSOR, val));
 		}
@@ -277,6 +297,7 @@ public Document toDocument() {
 		cursor.ifPresent(val -> document.put(CURSOR, val));
 		collation.ifPresent(val -> document.append(COLLATION, val.toDocument()));
 		comment.ifPresent(val -> document.append(COMMENT, val));
+		hint.ifPresent(val -> document.append(HINT, val));
 
 		if (hasExecutionTimeLimit()) {
 			document.append(MAX_TIME, maxTime.toMillis());
@@ -318,6 +339,7 @@ public static class Builder {
 		private @Nullable Document cursor;
 		private @Nullable Collation collation;
 		private @Nullable String comment;
+		private @Nullable Document hint;
 		private @Nullable Duration maxTime;
 		private @Nullable ResultOptions resultOptions;
 
@@ -396,6 +418,19 @@ public Builder comment(@Nullable String comment) {
 			return this;
 		}
 
+		/**
+		 * Define a hint is used forcibly by query optimizer to to fulfill the aggregation.
+		 *
+		 * @param hint can be {@literal null}.
+		 * @return this.
+		 * @since 2.2
+		 */
+		public Builder hint(@Nullable Document hint) {
+
+			this.hint = hint;
+			return this;
+		}
+
 		/**
 		 * Set the time limit for processing.
 		 *
@@ -431,7 +466,13 @@ public Builder skipOutput() {
 		 */
 		public AggregationOptions build() {
 
-			AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
+			AggregationOptions options = new AggregationOptions(
+					allowDiskUse,
+					explain,
+					cursor,
+					collation,
+					comment,
+					hint);
 			if (maxTime != null) {
 				options.maxTime = maxTime;
 			}

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

@@ -139,6 +139,7 @@
  * @author Mark Paluch
  * @author Michael J. Simons
  * @author Roman Puchkovskiy
+ * @author Yadhukrishna S Pai
  */
 @MockitoSettings(strictness = Strictness.LENIENT)
 public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
@@ -468,6 +469,17 @@ void aggregateShouldHonorOptionsComment() {
 		verify(aggregateIterable).comment("expensive");
 	}
 
+	@Test // DATAMONGO-1836
+	void aggregateShouldHonorOptionsHint() {
+
+		Document hint = new Document("dummyField", 1);
+		AggregationOptions options = AggregationOptions.builder().hint(hint).build();
+
+		template.aggregate(newAggregation(Aggregation.unwind("foo")).withOptions(options), "collection-1", Wrapper.class);
+
+		verify(aggregateIterable).hint(hint);
+	}
+
 	@Test // DATAMONGO-1166, DATAMONGO-2264
 	void geoNearShouldHonorReadPreferenceWhenSet() {
 

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

@@ -117,6 +117,7 @@
  * @author Christoph Strobl
  * @author Roman Puchkovskiy
  * @author Mathieu Ouellet
+ * @author Yadhukrishna S Pai
  */
 @ExtendWith(MockitoExtension.class)
 @MockitoSettings(strictness = Strictness.LENIENT)
@@ -628,6 +629,17 @@ void aggregateShouldHonorOptionsComment() {
 		verify(aggregatePublisher).comment("expensive");
 	}
 
+	@Test // DATAMONGO-1836
+	void aggregateShouldHonorOptionsHint() {
+		Document hint = new Document("dummyHint", 1);
+		AggregationOptions options = AggregationOptions.builder().hint(hint).build();
+
+		template.aggregate(newAggregation(Sith.class, project("id")).withOptions(options), AutogenerateableId.class,
+				Document.class).subscribe();
+
+		verify(aggregatePublisher).hint(hint);
+	}
+
 	@Test // DATAMONGO-2390
 	void aggregateShouldNoApplyZeroOrNegativeMaxTime() {
 

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

@@ -28,10 +28,12 @@
  * @author Thomas Darimont
  * @author Mark Paluch
  * @author Christoph Strobl
+ * @author Yadhukrishna S Pai
  * @since 1.6
  */
 public class AggregationOptionsTests {
 
+	private final Document dummyHint = new Document("dummyField", 1);
 	AggregationOptions aggregationOptions;
 
 	@BeforeEach
@@ -40,25 +42,28 @@ public void setup() {
 				.cursorBatchSize(1) //
 				.allowDiskUse(true) //
 				.comment("hola!") //
+				.hint(dummyHint) //
 				.build();
 	}
 
-	@Test // DATAMONGO-960
+	@Test // DATAMONGO-960, DATAMONGO-1836
 	public void aggregationOptionsBuilderShouldSetOptionsAccordingly() {
 
 		assertThat(aggregationOptions.isAllowDiskUse()).isTrue();
 		assertThat(aggregationOptions.isExplain()).isTrue();
 		assertThat(aggregationOptions.getCursor().get()).isEqualTo(new Document("batchSize", 1));
+		assertThat(aggregationOptions.getHint().get()).isEqualTo(dummyHint);
 	}
 
-	@Test // DATAMONGO-1637, DATAMONGO-2153
+	@Test // DATAMONGO-1637, DATAMONGO-2153, DATAMONGO-1836
 	public void shouldInitializeFromDocument() {
 
 		Document document = new Document();
 		document.put("cursor", new Document("batchSize", 1));
 		document.put("explain", true);
 		document.put("allowDiskUse", true);
 		document.put("comment", "hola!");
+		document.put("hint", dummyHint);
 
 		aggregationOptions = AggregationOptions.fromDocument(document);
 
@@ -67,12 +72,19 @@ public void shouldInitializeFromDocument() {
 		assertThat(aggregationOptions.getCursor().get()).isEqualTo(new Document("batchSize", 1));
 		assertThat(aggregationOptions.getCursorBatchSize()).isEqualTo(1);
 		assertThat(aggregationOptions.getComment().get()).isEqualTo("hola!");
+		assertThat(aggregationOptions.getHint().get()).isEqualTo(dummyHint);
 	}
 
-	@Test // DATAMONGO-960, DATAMONGO-2153
+	@Test // DATAMONGO-960, DATAMONGO-2153, DATAMONGO-1836
 	public void aggregationOptionsToString() {
 
 		assertThat(aggregationOptions.toDocument()).isEqualTo(Document.parse(
-				"{ \"allowDiskUse\" : true , \"explain\" : true , \"cursor\" : { \"batchSize\" : 1}, \"comment\": \"hola!\"}"));
+				"{ " +
+						"\"allowDiskUse\" : true , " +
+						"\"explain\" : true , " +
+						"\"cursor\" : { \"batchSize\" : 1}, " +
+						"\"comment\": \"hola!\", " +
+						"\"hint\" : { \"dummyField\" : 1}" +
+						"}"));
 	}
 }