DATAMONGO-2153 - Apply Meta comment and cursor size to AggregationOptions.

We now apply propagate the cursor size to aggregation options. We introduced AggregationOptions.comment to propagate the meta comment to aggregation execution.

Original pull request: #743.

Author: mp911de

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

@@ -2136,6 +2136,8 @@ protected <O> AggregationResults<O> doAggregate(Aggregation aggregation, String
 				aggregateIterable = aggregateIterable.batchSize(options.getCursorBatchSize());
 			}
 
+			options.getComment().ifPresent(aggregateIterable::comment);
+
 			MongoIterable<O> iterable = aggregateIterable.map(val -> {
 
 				rawResult.add(val);
@@ -2177,6 +2179,8 @@ protected <O> CloseableIterator<O> aggregateStream(Aggregation aggregation, Stri
 				cursor = cursor.batchSize(options.getCursorBatchSize());
 			}
 
+			options.getComment().ifPresent(cursor::comment);
+
 			Class<?> domainType = aggregation instanceof TypedAggregation ? ((TypedAggregation) aggregation).getInputType()
 					: null;
 

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

@@ -1036,6 +1036,8 @@ private <O> Flux<O> aggregateAndMap(MongoCollection<Document> collection, List<D
 			cursor = cursor.batchSize(options.getCursorBatchSize());
 		}
 
+		options.getComment().ifPresent(cursor::comment);
+
 		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

@@ -44,11 +44,13 @@ public class AggregationOptions {
 	private static final String EXPLAIN = "explain";
 	private static final String ALLOW_DISK_USE = "allowDiskUse";
 	private static final String COLLATION = "collation";
+	private static final String COMMENT = "comment";
 
 	private final boolean allowDiskUse;
 	private final boolean explain;
 	private final Optional<Document> cursor;
 	private final Optional<Collation> collation;
+	private final Optional<String> comment;
 
 	/**
 	 * Creates a new {@link AggregationOptions}.
@@ -73,11 +75,28 @@ 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);
+	}
+
+	/**
+	 * Creates a new {@link AggregationOptions}.
+	 *
+	 * @param allowDiskUse whether to off-load intensive sort-operations to disk.
+	 * @param explain whether to get the execution plan for the aggregation instead of the actual results.
+	 * @param cursor can be {@literal null}, used to pass additional options (such as {@code batchSize}) to the
+	 *          aggregation.
+	 * @param collation collation for string comparison. Can be {@literal null}.
+	 * @param comment execution comment. Can be {@literal null}.
+	 * @since 2.2
+	 */
+	public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Document cursor,
+			@Nullable Collation collation, @Nullable String comment) {
 
 		this.allowDiskUse = allowDiskUse;
 		this.explain = explain;
 		this.cursor = Optional.ofNullable(cursor);
 		this.collation = Optional.ofNullable(collation);
+		this.comment = Optional.ofNullable(comment);
 	}
 
 	/**
@@ -108,8 +127,9 @@ public static AggregationOptions fromDocument(Document document) {
 		Document cursor = document.get(CURSOR, Document.class);
 		Collation collation = document.containsKey(COLLATION) ? Collation.from(document.get(COLLATION, Document.class))
 				: null;
+		String comment = document.getString(COMMENT);
 
-		return new AggregationOptions(allowDiskUse, explain, cursor, collation);
+		return new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
 	}
 
 	/**
@@ -176,6 +196,16 @@ public Optional<Collation> getCollation() {
 		return collation;
 	}
 
+	/**
+	 * Get the comment for the aggregation.
+	 *
+	 * @return
+	 * @since 2.2
+	 */
+	public Optional<String> getComment() {
+		return comment;
+	}
+
 	/**
 	 * Returns a new potentially adjusted copy for the given {@code aggregationCommandObject} with the configuration
 	 * applied.
@@ -219,6 +249,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));
 
 		return document;
 	}
@@ -247,6 +278,7 @@ public static class Builder {
 		private boolean explain;
 		private @Nullable Document cursor;
 		private @Nullable Collation collation;
+		private @Nullable String comment;
 
 		/**
 		 * Defines whether to off-load intensive sort-operations to disk.
@@ -302,20 +334,34 @@ public Builder cursorBatchSize(int batchSize) {
 		 *
 		 * @param collation can be {@literal null}.
 		 * @return
+		 * @since 2.0
 		 */
 		public Builder collation(@Nullable Collation collation) {
 
 			this.collation = collation;
 			return this;
 		}
 
+		/**
+		 * Define a comment to describe the execution.
+		 *
+		 * @param comment can be {@literal null}.
+		 * @return
+		 * @since 2.2
+		 */
+		public Builder comment(@Nullable String comment) {
+
+			this.comment = comment;
+			return this;
+		}
+
 		/**
 		 * Returns a new {@link AggregationOptions} instance with the given configuration.
 		 *
 		 * @return
 		 */
 		public AggregationOptions build() {
-			return new AggregationOptions(allowDiskUse, explain, cursor, collation);
+			return new AggregationOptions(allowDiskUse, explain, cursor, collation, comment);
 		}
 	}
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/AggregationUtils.java

@@ -30,19 +30,22 @@
 import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
 import org.springframework.data.mongodb.core.query.Collation;
+import org.springframework.data.mongodb.core.query.Meta;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.util.json.ParameterBindingContext;
 import org.springframework.data.mongodb.util.json.ParameterBindingDocumentCodec;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.lang.Nullable;
 import org.springframework.util.ClassUtils;
+import org.springframework.util.StringUtils;
 
 /**
  * Internal utility class to help avoid duplicate code required in both the reactive and the sync {@link Aggregation}
  * support offered by repositories.
  *
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @since 2.2
  */
 @UtilityClass
@@ -72,6 +75,27 @@ static AggregationOptions.Builder applyCollation(AggregationOptions.Builder buil
 		return collation == null ? builder : builder.collation(collation);
 	}
 
+	/**
+	 * Apply {@link Meta#getComment()} and {@link Meta#getCursorBatchSize()}.
+	 *
+	 * @param builder must not be {@literal null}.
+	 * @param queryMethod must not be {@literal null}.
+	 */
+	static AggregationOptions.Builder applyMeta(AggregationOptions.Builder builder, MongoQueryMethod queryMethod) {
+
+		Meta meta = queryMethod.getQueryMetaAttributes();
+
+		if (StringUtils.hasText(meta.getComment())) {
+			builder.comment(meta.getComment());
+		}
+
+		if (meta.getCursorBatchSize() != null) {
+			builder.cursorBatchSize(meta.getCursorBatchSize());
+		}
+
+		return builder;
+	}
+
 	/**
 	 * Compute the {@link AggregationOperation aggregation} pipeline for the given {@link MongoQueryMethod}. The raw
 	 * {@link org.springframework.data.mongodb.repository.Aggregation#pipeline()} is parsed with a

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/ReactiveStringBasedAggregation.java

@@ -115,10 +115,13 @@ List<AggregationOperation> computePipeline(ConvertingParameterAccessor accessor)
 
 	private AggregationOptions computeOptions(MongoQueryMethod method, ConvertingParameterAccessor accessor) {
 
-		return AggregationUtils
-				.applyCollation(Aggregation.newAggregationOptions(), method.getAnnotatedCollation(), accessor,
-						method.getParameters(), expressionParser, evaluationContextProvider) //
-				.build();
+		AggregationOptions.Builder builder = Aggregation.newAggregationOptions();
+
+		AggregationUtils.applyCollation(builder, method.getAnnotatedCollation(), accessor, method.getParameters(),
+				expressionParser, evaluationContextProvider);
+		AggregationUtils.applyMeta(builder, method);
+
+		return builder.build();
 	}
 
 	/*

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/StringBasedAggregation.java

@@ -125,10 +125,13 @@ List<AggregationOperation> computePipeline(MongoQueryMethod method, ConvertingPa
 
 	private AggregationOptions computeOptions(MongoQueryMethod method, ConvertingParameterAccessor accessor) {
 
-		return AggregationUtils
-				.applyCollation(Aggregation.newAggregationOptions(), method.getAnnotatedCollation(), accessor,
-						method.getParameters(), expressionParser, evaluationContextProvider) //
-				.build();
+		AggregationOptions.Builder builder = Aggregation.newAggregationOptions();
+
+		AggregationUtils.applyCollation(builder, method.getAnnotatedCollation(), accessor, method.getParameters(),
+				expressionParser, evaluationContextProvider);
+		AggregationUtils.applyMeta(builder, method);
+
+		return builder.build();
 	}
 
 	/*

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

@@ -55,6 +55,7 @@
 import org.springframework.data.geo.Point;
 import org.springframework.data.mongodb.MongoDbFactory;
 import org.springframework.data.mongodb.core.aggregation.Aggregation;
+import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver;
 import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
 import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
@@ -418,6 +419,16 @@ public void aggregateShouldIgnoreReadPreferenceWhenNotSet() {
 		verify(collection, never()).withReadPreference(any());
 	}
 
+	@Test // DATAMONGO-2153
+	public void aggregateShouldHonorOptionsComment() {
+
+		AggregationOptions options = AggregationOptions.builder().comment("expensive").build();
+
+		template.aggregate(newAggregation(Aggregation.unwind("foo")).withOptions(options), "collection-1", Wrapper.class);
+
+		verify(aggregateIterable).comment("expensive");
+	}
+
 	@Test // DATAMONGO-1166, DATAMONGO-2264
 	public void geoNearShouldHonorReadPreferenceWhenSet() {
 

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

@@ -43,6 +43,7 @@
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.mongodb.core.MongoTemplateUnitTests.AutogenerateableId;
+import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
 import org.springframework.data.mongodb.core.convert.NoOpDbRefResolver;
 import org.springframework.data.mongodb.core.mapping.Field;
@@ -539,6 +540,17 @@ public void aggreateShouldUseCollationFromOptionsEvenIfDefaultCollationIsPresent
 		verify(aggregatePublisher).collation(eq(com.mongodb.client.model.Collation.builder().locale("fr").build()));
 	}
 
+	@Test // DATAMONGO-2153
+	public void aggregateShouldHonorOptionsComment() {
+
+		AggregationOptions options = AggregationOptions.builder().comment("expensive").build();
+
+		template.aggregate(newAggregation(Sith.class, project("id")).withOptions(options), AutogenerateableId.class,
+				Document.class).subscribe();
+
+		verify(aggregatePublisher).comment("expensive");
+	}
+
 	@Test // DATAMONGO-18545
 	public void findAndReplaceShouldUseCollationWhenPresent() {
 

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

@@ -15,8 +15,7 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
-import static org.hamcrest.CoreMatchers.*;
-import static org.junit.Assert.*;
+import static org.assertj.core.api.Assertions.*;
 import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
 
 import org.bson.Document;
@@ -40,36 +39,40 @@ public void setup() {
 		aggregationOptions = newAggregationOptions().explain(true) //
 				.cursorBatchSize(1) //
 				.allowDiskUse(true) //
+				.comment("hola!") //
 				.build();
 	}
 
 	@Test // DATAMONGO-960
 	public void aggregationOptionsBuilderShouldSetOptionsAccordingly() {
 
-		assertThat(aggregationOptions.isAllowDiskUse(), is(true));
-		assertThat(aggregationOptions.isExplain(), is(true));
-		assertThat(aggregationOptions.getCursor().get(), is(new Document("batchSize", 1)));
+		assertThat(aggregationOptions.isAllowDiskUse()).isTrue();
+		assertThat(aggregationOptions.isExplain()).isTrue();
+		assertThat(aggregationOptions.getCursor().get()).isEqualTo(new Document("batchSize", 1));
 	}
 
-	@Test // DATAMONGO-1637
+	@Test // DATAMONGO-1637, DATAMONGO-2153
 	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!");
 
 		aggregationOptions = AggregationOptions.fromDocument(document);
 
-		assertThat(aggregationOptions.isAllowDiskUse(), is(true));
-		assertThat(aggregationOptions.isExplain(), is(true));
-		assertThat(aggregationOptions.getCursor().get(), is(new Document("batchSize", 1)));
-		assertThat(aggregationOptions.getCursorBatchSize(), is(1));
+		assertThat(aggregationOptions.isAllowDiskUse()).isTrue();
+		assertThat(aggregationOptions.isExplain()).isTrue();
+		assertThat(aggregationOptions.getCursor().get()).isEqualTo(new Document("batchSize", 1));
+		assertThat(aggregationOptions.getCursorBatchSize()).isEqualTo(1);
+		assertThat(aggregationOptions.getComment().get()).isEqualTo("hola!");
 	}
 
-	@Test // DATAMONGO-960
+	@Test // DATAMONGO-960, DATAMONGO-2153
 	public void aggregationOptionsToString() {
-		assertThat(aggregationOptions.toDocument(),
-				is(Document.parse("{ \"allowDiskUse\" : true , \"explain\" : true , \"cursor\" : { \"batchSize\" : 1}}")));
+
+		assertThat(aggregationOptions.toDocument()).isEqualTo(Document.parse(
+				"{ \"allowDiskUse\" : true , \"explain\" : true , \"cursor\" : { \"batchSize\" : 1}, \"comment\": \"hola!\"}"));
 	}
 }