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.

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

@@ -2136,6 +2136,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 				aggregateIterable = aggregateIterable.batchSize(options.getCursorBatchSize());
 			}
 
+			options.getComment().ifPresent(aggregateIterable::comment);
+
 			MongoIterable<O> iterable = aggregateIterable.map(val -> {
 
 				rawResult.add(val);
@@ -2177,6 +2179,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 				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 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
 			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 class AggregationOptions {
 	 */
 	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 class AggregationOptions {
 		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 class AggregationOptions {
 		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 class AggregationOptions {
 
 		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 class AggregationOptions {
 		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,6 +334,7 @@ public class AggregationOptions {
 		 *
 		 * @param collation can be {@literal null}.
 		 * @return
+		 * @since 2.0
 		 */
 		public Builder collation(@Nullable Collation collation) {
 
@@ -309,13 +342,26 @@ public class AggregationOptions {
 			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,6 +30,7 @@ import org.springframework.data.mongodb.core.aggregation.AggregationOperationCon
 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;
@@ -37,12 +38,14 @@ import org.springframework.data.repository.query.QueryMethodEvaluationContextPro
 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 @@ class AggregationUtils {
 		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 @@ public class ReactiveStringBasedAggregation extends AbstractReactiveMongoQuery {
 
 	private AggregationOptions computeOptions(MongoQueryMethod method, ConvertingParameterAccessor accessor) {
 
-		return AggregationUtils
+		AggregationOptions.Builder builder = Aggregation.newAggregationOptions();
-				.applyCollation(Aggregation.newAggregationOptions(), method.getAnnotatedCollation(), accessor,
+
-						method.getParameters(), expressionParser, evaluationContextProvider) //
+		AggregationUtils.applyCollation(builder, method.getAnnotatedCollation(), accessor, method.getParameters(),
-				.build();
+				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 @@ public class StringBasedAggregation extends AbstractMongoQuery {
 
 	private AggregationOptions computeOptions(MongoQueryMethod method, ConvertingParameterAccessor accessor) {
 
-		return AggregationUtils
+		AggregationOptions.Builder builder = Aggregation.newAggregationOptions();
-				.applyCollation(Aggregation.newAggregationOptions(), method.getAnnotatedCollation(), accessor,
+
-						method.getParameters(), expressionParser, evaluationContextProvider) //
+		AggregationUtils.applyCollation(builder, method.getAnnotatedCollation(), accessor, method.getParameters(),
-				.build();
+				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.domain.Sort;
 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 class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		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.reactivestreams.Publisher;
 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 class ReactiveMongoTemplateUnitTests {
 		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.assertj.core.api.Assertions.*;
-import static org.junit.Assert.*;
 import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
 
 import org.bson.Document;
@@ -40,36 +39,40 @@ public class AggregationOptionsTests {
 		aggregationOptions = newAggregationOptions().explain(true) //
 				.cursorBatchSize(1) //
 				.allowDiskUse(true) //
+				.comment("hola!") //
 				.build();
 	}
 
 	@Test // DATAMONGO-960
 	public void aggregationOptionsBuilderShouldSetOptionsAccordingly() {
 
-		assertThat(aggregationOptions.isAllowDiskUse(), is(true));
+		assertThat(aggregationOptions.isAllowDiskUse()).isTrue();
-		assertThat(aggregationOptions.isExplain(), is(true));
+		assertThat(aggregationOptions.isExplain()).isTrue();
-		assertThat(aggregationOptions.getCursor().get(), is(new Document("batchSize", 1)));
+		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.isAllowDiskUse()).isTrue();
-		assertThat(aggregationOptions.isExplain(), is(true));
+		assertThat(aggregationOptions.isExplain()).isTrue();
-		assertThat(aggregationOptions.getCursor().get(), is(new Document("batchSize", 1)));
+		assertThat(aggregationOptions.getCursor().get()).isEqualTo(new Document("batchSize", 1));
-		assertThat(aggregationOptions.getCursorBatchSize(), is(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!\"}"));
 	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/query/ReactiveStringBasedAggregationUnitTests.java

@@ -39,6 +39,7 @@ import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mongodb.core.ReactiveMongoOperations;
 import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
+import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
 import org.springframework.data.mongodb.core.convert.DbRefResolver;
@@ -48,6 +49,7 @@ import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.core.query.Collation;
 import org.springframework.data.mongodb.repository.Aggregation;
+import org.springframework.data.mongodb.repository.Meta;
 import org.springframework.data.mongodb.repository.Person;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
@@ -62,6 +64,7 @@ import org.springframework.util.ClassUtils;
  * Unit tests for {@link ReactiveStringBasedAggregation}.
  *
  * @author Christoph Strobl
+ * @author Mark Paluch
  */
 @RunWith(MockitoJUnitRunner.class)
 public class ReactiveStringBasedAggregationUnitTests {
@@ -98,6 +101,17 @@ public class ReactiveStringBasedAggregationUnitTests {
 		assertThat(pipelineOf(invocation)).containsExactly(GROUP_BY_LASTNAME, SORT);
 	}
 
+	@Test // DATAMONGO-2153
+	public void plainStringAggregationConsidersMeta() {
+
+		AggregationInvocation invocation = executeAggregation("plainStringAggregation");
+
+		AggregationOptions options = invocation.aggregation.getOptions();
+
+		assertThat(options.getComment()).contains("expensive-aggregation");
+		assertThat(options.getCursorBatchSize()).isEqualTo(42);
+	}
+
 	@Test // DATAMONGO-2153
 	public void plainStringAggregationWithSortParameter() {
 
@@ -107,6 +121,11 @@ public class ReactiveStringBasedAggregationUnitTests {
 		assertThat(inputTypeOf(invocation)).isEqualTo(Person.class);
 		assertThat(targetTypeOf(invocation)).isEqualTo(PersonAggregate.class);
 		assertThat(pipelineOf(invocation)).containsExactly(GROUP_BY_LASTNAME, SORT);
+
+		AggregationOptions options = invocation.aggregation.getOptions();
+
+		assertThat(options.getComment()).isEmpty();
+		assertThat(options.getCursorBatchSize()).isNull();
 	}
 
 	@Test // DATAMONGO-2153
@@ -194,6 +213,7 @@ public class ReactiveStringBasedAggregationUnitTests {
 
 	private interface SampleRepository extends ReactiveCrudRepository<Person, Long> {
 
+		@Meta(cursorBatchSize = 42, comment = "expensive-aggregation")
 		@Aggregation({ RAW_GROUP_BY_LASTNAME_STRING, RAW_SORT_STRING })
 		Mono<PersonAggregate> plainStringAggregation();
 

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/query/StringBasedAggregationUnitTests.java

@@ -25,7 +25,6 @@ import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.function.IntFunction;
 
 import org.bson.Document;
 import org.junit.Before;
@@ -39,6 +38,7 @@ import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
+import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.aggregation.AggregationResults;
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
@@ -49,6 +49,7 @@ import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.core.query.Collation;
 import org.springframework.data.mongodb.repository.Aggregation;
+import org.springframework.data.mongodb.repository.Meta;
 import org.springframework.data.mongodb.repository.Person;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
@@ -63,6 +64,7 @@ import org.springframework.util.ClassUtils;
  * Unit tests for {@link StringBasedAggregation}.
  *
  * @author Christoph Strobl
+ * @author Mark Paluch
  */
 @RunWith(MockitoJUnitRunner.class)
 public class StringBasedAggregationUnitTests {
@@ -100,13 +102,30 @@ public class StringBasedAggregationUnitTests {
 		assertThat(pipelineOf(invocation)).containsExactly(GROUP_BY_LASTNAME, SORT);
 	}
 
+	@Test // DATAMONGO-2153
+	public void plainStringAggregationConsidersMeta() {
+
+		AggregationInvocation invocation = executeAggregation("plainStringAggregation");
+
+		AggregationOptions options = invocation.aggregation.getOptions();
+
+		assertThat(options.getComment()).contains("expensive-aggregation");
+		assertThat(options.getCursorBatchSize()).isEqualTo(42);
+	}
+
 	@Test // DATAMONGO-2153
 	public void returnSingleObject() {
 
 		PersonAggregate expected = new PersonAggregate();
 		when(aggregationResults.getUniqueMappedResult()).thenReturn(Collections.singletonList(expected));
 
-		assertThat(executeAggregation("returnSingleEntity").result).isEqualTo(expected);
+		AggregationInvocation invocation = executeAggregation("returnSingleEntity");
+		assertThat(invocation.result).isEqualTo(expected);
+
+		AggregationOptions options = invocation.aggregation.getOptions();
+
+		assertThat(options.getComment()).isEmpty();
+		assertThat(options.getCursorBatchSize()).isNull();
 	}
 
 	@Test // DATAMONGO-2153
@@ -227,6 +246,7 @@ public class StringBasedAggregationUnitTests {
 
 	private interface SampleRepository extends Repository<Person, Long> {
 
+		@Meta(cursorBatchSize = 42, comment = "expensive-aggregation")
 		@Aggregation({ RAW_GROUP_BY_LASTNAME_STRING, RAW_SORT_STRING })
 		PersonAggregate plainStringAggregation();