DATAMONGO-1824 - Fix aggregation execution for MongoDB 3.6.

We now send aggregation commands along a cursor batch size for compatibility with MongoDB 3.6 that no longer supports aggregations without cursor. We consume the whole cursor before returning and converting results and omit the 16MB aggregation result limit. For MongoDB versions not supporting aggregation cursors we return results directly.

Original pull request: #521.

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

@@ -18,6 +18,7 @@ package org.springframework.data.mongodb.core;
 import static org.springframework.data.mongodb.core.query.Criteria.*;
 import static org.springframework.data.mongodb.core.query.SerializationUtils.*;
 
+import com.mongodb.*;
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
 import lombok.NonNull;
@@ -127,14 +128,6 @@ import org.springframework.util.ObjectUtils;
 import org.springframework.util.ResourceUtils;
 import org.springframework.util.StringUtils;
 
-import com.mongodb.Cursor;
-import com.mongodb.DBCollection;
-import com.mongodb.DBCursor;
-import com.mongodb.Mongo;
-import com.mongodb.MongoClient;
-import com.mongodb.MongoException;
-import com.mongodb.ReadPreference;
-import com.mongodb.WriteConcern;
 import com.mongodb.client.AggregateIterable;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MapReduceIterable;
@@ -1933,16 +1926,10 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 		Assert.notNull(aggregation, "Aggregation pipeline must not be null!");
 		Assert.notNull(outputType, "Output type must not be null!");
 
-		AggregationOperationContext rootContext = context == null ? Aggregation.DEFAULT_CONTEXT : context;
-		Document command = aggregation.toDocument(collectionName, rootContext);
+		Document commandResult = new BatchAggregationLoader(this, readPreference, Integer.MAX_VALUE)
+				.aggregate(collectionName, aggregation, context);
 
-		if (LOGGER.isDebugEnabled()) {
-			LOGGER.debug("Executing aggregation: {}", serializeToJsonSafely(command));
-		}
-
-		Document commandResult = executeCommand(command, this.readPreference);
-
-		return new AggregationResults<O>(returnPotentiallyMappedResults(outputType, commandResult, collectionName),
+		return new AggregationResults<>(returnPotentiallyMappedResults(outputType, commandResult, collectionName),
 				commandResult);
 	}
 
@@ -3074,4 +3061,160 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 	public MongoDbFactory getMongoDbFactory() {
 		return mongoDbFactory;
 	}
+
+	/**
+	 * {@link BatchAggregationLoader} is a little helper that can process cursor results returned by an aggregation
+	 * command execution. On presence of a {@literal nextBatch} indicated by presence of an {@code id} field in the
+	 * {@code cursor} another {@code getMore} command gets executed reading the next batch of documents until everything
+	 * has been loaded.
+	 *
+	 * @author Christoph Strobl
+	 * @since 1.10
+	 */
+	static class BatchAggregationLoader {
+
+		private static final String CURSOR_FIELD = "cursor";
+		private static final String RESULT_FIELD = "result";
+		private static final String BATCH_SIZE_FIELD = "batchSize";
+
+		private final MongoTemplate template;
+		private final ReadPreference readPreference;
+		private final int batchSize;
+
+		BatchAggregationLoader(MongoTemplate template, ReadPreference readPreference, int batchSize) {
+
+			this.template = template;
+			this.readPreference = readPreference;
+			this.batchSize = batchSize;
+		}
+
+		Document aggregate(String collectionName, Aggregation aggregation, AggregationOperationContext context) {
+
+			Document command = AggregationCommandPreparer.INSTANCE.prepareAggregationCommand(collectionName, aggregation,
+					context, batchSize);
+
+			if (LOGGER.isDebugEnabled()) {
+				LOGGER.debug("Executing aggregation: {}", serializeToJsonSafely(command));
+			}
+
+			List<Document> results = aggregateBatched(collectionName, batchSize, command);
+			return mergeArregationCommandResults(results);
+		}
+
+		private Document mergeArregationCommandResults(List<Document> results) {
+
+			Document commandResult = new Document();
+			if (results.size() == 1) {
+				commandResult = results.iterator().next();
+			} else {
+
+				List<Object> allResults = new ArrayList();
+
+				for (Document result : results) {
+					Collection foo = (Collection<?>) result.get(RESULT_FIELD);
+					if (!CollectionUtils.isEmpty(foo)) {
+						allResults.addAll(foo);
+					}
+				}
+
+				// take general info from first batch
+				commandResult.put("serverUsed", results.iterator().next().get("serverUsed"));
+				commandResult.put("ok", results.iterator().next().get("ok"));
+
+				// and append the merged results
+				commandResult.put(RESULT_FIELD, allResults);
+			}
+			return commandResult;
+		}
+
+		private List<Document> aggregateBatched(String collectionName, int batchSize, Document command) {
+
+			List<Document> results = new ArrayList<>();
+
+			Document tmp = template.executeCommand(command, readPreference);
+			results.add(AggregationResultPostProcessor.INSTANCE.process(command, tmp));
+
+			while (hasNext(tmp)) {
+
+				Document getMore = new Document("getMore", getNextBatchId(tmp)) //
+						.append("collection", collectionName) //
+						.append(BATCH_SIZE_FIELD, batchSize); //
+
+				tmp = template.executeCommand(getMore, this.readPreference);
+				results.add(AggregationResultPostProcessor.INSTANCE.process(command, tmp));
+			}
+
+			return results;
+		}
+
+		private boolean hasNext(Document commandResult) {
+
+			if (!commandResult.containsKey(CURSOR_FIELD)) {
+				return false;
+			}
+
+			Object next = getNextBatchId(commandResult);
+			return (next == null || ((Number) next).longValue() == 0L) ? false : true;
+		}
+
+		private Object getNextBatchId(Document commandResult) {
+			return ((Document) commandResult.get(CURSOR_FIELD)).get("id");
+		}
+
+		/**
+		 * Helper to pre process the aggregation command sent to the server by adding {@code cursor} options to match
+		 * execution on different server versions.
+		 *
+		 * @author Christoph Strobl
+		 * @since 1.10
+		 */
+		private static enum AggregationCommandPreparer {
+
+			INSTANCE;
+
+			Document prepareAggregationCommand(String collectionName, Aggregation aggregation,
+					AggregationOperationContext context, int batchSize) {
+
+				AggregationOperationContext rootContext = context == null ? Aggregation.DEFAULT_CONTEXT : context;
+				Document command = aggregation.toDocument(collectionName, rootContext);
+
+				if (!aggregation.getOptions().isExplain()) {
+					command.put(CURSOR_FIELD, new Document(BATCH_SIZE_FIELD, batchSize));
+				}
+
+				return command;
+			}
+		}
+
+		/**
+		 * Helper to post process aggregation command result by copying over required attributes.
+		 *
+		 * @author Christoph Strobl
+		 * @since 1.10
+		 */
+		private static enum AggregationResultPostProcessor {
+
+			INSTANCE;
+
+			Document process(Document command, Document commandResult) {
+
+				if (!commandResult.containsKey(CURSOR_FIELD)) {
+					return commandResult;
+				}
+
+				Document resultObject = new Document("serverUsed", commandResult.get("serverUsed"));
+				resultObject.put("ok", commandResult.get("ok"));
+
+				Document cursor = (Document) commandResult.get(CURSOR_FIELD);
+				if (cursor.containsKey("firstBatch")) {
+					resultObject.put(RESULT_FIELD, cursor.get("firstBatch"));
+				} else {
+					resultObject.put(RESULT_FIELD, cursor.get("nextBatch"));
+				}
+
+				return resultObject;
+			}
+		}
+	}
+
 }

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

@@ -601,6 +601,16 @@ public class Aggregation {
 		return SerializationUtils.serializeToJsonSafely(toDocument("__collection__", DEFAULT_CONTEXT));
 	}
 
+	/**
+	 * Get {@link AggregationOptions} to apply.
+	 *
+	 * @return never {@literal null}
+	 * @since 1.10
+	 */
+	public AggregationOptions getOptions() {
+		return options;
+	}
+
 	/**
 	 * Describes the system variables available in MongoDB aggregation framework pipeline expressions.
 	 *

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

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2017 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core;
+
+import static java.util.Collections.*;
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
+
+import java.util.List;
+
+import org.bson.Document;
+import org.hamcrest.core.IsCollectionContaining;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.springframework.data.mongodb.core.MongoTemplate.BatchAggregationLoader;
+import org.springframework.data.mongodb.core.aggregation.Aggregation;
+import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
+
+import com.mongodb.ReadPreference;
+
+/**
+ * @author Christoph Strobl
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class BatchAggregationLoaderUnitTests {
+
+	static final TypedAggregation<Person> AGGREGATION = newAggregation(Person.class,
+			project().and("firstName").as("name"));
+
+	@Mock MongoTemplate template;
+	@Mock Document aggregationResult;
+	@Mock Document getMoreResult;
+
+	BatchAggregationLoader loader;
+
+	Document cursorWithoutMore = new Document("firstBatch", singletonList(new Document("name", "luke")));
+	Document cursorWithMore = new Document("id", 123).append("firstBatch", singletonList(new Document("name", "luke")));
+	Document cursorWithNoMore = new Document("id", 0).append("nextBatch", singletonList(new Document("name", "han")));
+
+	@Before
+	public void setUp() {
+		loader = new BatchAggregationLoader(template, ReadPreference.primary(), 10);
+	}
+
+	@Test // DATAMONGO-1824
+	public void shouldLoadJustOneBatchWhenAlreayDoneWithFirst() {
+
+		when(template.executeCommand(any(Document.class), any(ReadPreference.class))).thenReturn(aggregationResult);
+		when(aggregationResult.containsKey("cursor")).thenReturn(true);
+		when(aggregationResult.get("cursor")).thenReturn(cursorWithoutMore);
+
+		Document result = loader.aggregate("person", AGGREGATION, Aggregation.DEFAULT_CONTEXT);
+		
+		assertThat((List<Object>) result.get("result"),
+				IsCollectionContaining.<Object> hasItem(new Document("name", "luke")));
+
+		verify(template).executeCommand(any(Document.class), any(ReadPreference.class));
+		verifyNoMoreInteractions(template);
+	}
+
+	@Test // DATAMONGO-1824
+	public void shouldBatchLoadWhenRequired() {
+
+		when(template.executeCommand(any(Document.class), any(ReadPreference.class))).thenReturn(aggregationResult)
+				.thenReturn(getMoreResult);
+		when(aggregationResult.containsKey("cursor")).thenReturn(true);
+		when(aggregationResult.get("cursor")).thenReturn(cursorWithMore);
+		when(getMoreResult.containsKey("cursor")).thenReturn(true);
+		when(getMoreResult.get("cursor")).thenReturn(cursorWithNoMore);
+
+		Document result = loader.aggregate("person", AGGREGATION, Aggregation.DEFAULT_CONTEXT);
+		assertThat((List<Object>) result.get("result"),
+				IsCollectionContaining.<Object> hasItems(new Document("name", "luke"), new Document("name", "han")));
+
+		verify(template, times(2)).executeCommand(any(Document.class), any(ReadPreference.class));
+		verifyNoMoreInteractions(template);
+	}
+}