DATAMONGO-1986 - Always provide a typed AggregationOperationContext for TypedAggregation.

We now initialize a TypeBasedAggregationOperationContext for TypedAggregations if no context is provided. This makes sure that potential Criteria objects are run trough the QueryMapper.
In case the default context is used we now also make sure to at least run the aggregation pipeline through the QueryMapper to avoid passing on non MongoDB simple types to the driver.

Original pull request: #564.

Author: christophstrobl

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

@@ -27,6 +27,7 @@
 import java.util.*;
 import java.util.Map.Entry;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import org.bson.Document;
 import org.bson.conversions.Bson;
@@ -204,7 +205,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 	 * @param databaseName must not be {@literal null} or empty.
 	 */
 	public MongoTemplate(MongoClient mongoClient, String databaseName) {
-		this(new SimpleMongoDbFactory(mongoClient, databaseName), null);
+		this(new SimpleMongoDbFactory(mongoClient, databaseName),  null);
 	}
 
 	/**
@@ -2576,6 +2577,73 @@ private Document addFieldsForProjection(Document fields, Class<?> domainType, Cl
 		return fields;
 	}
 
+	/**
+	 * Prepare the {@link AggregationOperationContext} for a given aggregation by either returning the context itself it
+	 * is not {@literal null}, create a {@link TypeBasedAggregationOperationContext} if the aggregation contains type
+	 * information (is a {@link TypedAggregation}) or use the {@link Aggregation#DEFAULT_CONTEXT}.
+	 *
+	 * @param aggregation must not be {@literal null}.
+	 * @param context can be {@literal null}.
+	 * @return the root {@link AggregationOperationContext} to use.
+	 */
+	private AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
+			@Nullable AggregationOperationContext context) {
+
+		if (context != null) {
+			return context;
+		}
+
+		if (aggregation instanceof TypedAggregation) {
+			return new TypeBasedAggregationOperationContext(((TypedAggregation) aggregation).getInputType(), mappingContext,
+					queryMapper);
+		}
+
+		return Aggregation.DEFAULT_CONTEXT;
+	}
+
+	/**
+	 * Extract and map the aggregation pipeline.
+	 *
+	 * @param aggregation
+	 * @param context
+	 * @return
+	 */
+	private Document aggregationToPipeline(String inputCollectionName, Aggregation aggregation, AggregationOperationContext context) {
+
+		if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
+			return aggregation.toDocument(inputCollectionName, context);
+		}
+
+		return queryMapper.getMappedObject(aggregation.toDocument(inputCollectionName, context), Optional.empty());
+	}
+
+	/**
+	 * Extract the command and map the aggregation pipeline.
+	 *
+	 * @param aggregation
+	 * @param context
+	 * @return
+	 */
+	private Document aggregationToCommand(String collection, Aggregation aggregation,
+			AggregationOperationContext context) {
+
+		Document command = aggregation.toDocument(collection, context);
+
+		if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
+			return command;
+		}
+
+		command.put("pipeline", mapAggregationPipeline(command.get("pipeline", List.class)));
+
+		return command;
+	}
+
+	private List<Document> mapAggregationPipeline(List<Document> pipeline) {
+
+		return pipeline.stream().map(val -> queryMapper.getMappedObject(val, Optional.empty()))
+				.collect(Collectors.toList());
+	}
+
 	/**
 	 * Tries to convert the given {@link RuntimeException} into a {@link DataAccessException} but returns the original
 	 * exception if the conversation failed. Thus allows safe re-throwing of the return value.

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

@@ -733,8 +733,8 @@ protected <O> Flux<O> aggregate(Aggregation aggregation, String collectionName,
 		Assert.hasText(collectionName, "Collection name must not be null or empty!");
 		Assert.notNull(outputType, "Output type must not be null!");
 
-		AggregationOperationContext rootContext = context == null ? Aggregation.DEFAULT_CONTEXT : context;
-		Document command = aggregation.toDocument(collectionName, rootContext);
+		AggregationOperationContext rootContext = prepareAggregationContext(aggregation, context);
+		Document command = aggregationToPipeline(collectionName, aggregation, rootContext);
 		AggregationOptions options = AggregationOptions.fromDocument(command);
 
 		Assert.isTrue(!options.isExplain(), "Cannot use explain option with streaming!");
@@ -752,8 +752,8 @@ protected <O> Flux<O> aggregate(Aggregation aggregation, String collectionName,
 	private <O> Flux<O> aggregateAndMap(MongoCollection<Document> collection, List<Document> pipeline,
 			AggregationOptions options, ReadDocumentCallback<O> readCallback) {
 
-		AggregatePublisher<Document> cursor = collection.aggregate(pipeline).allowDiskUse(options.isAllowDiskUse())
-				.useCursor(true);
+		AggregatePublisher<Document> cursor = collection.aggregate(pipeline)
+				.allowDiskUse(options.isAllowDiskUse());
 
 		if (options.getCollation().isPresent()) {
 			cursor = cursor.collation(options.getCollation().map(Collation::toMongoCollation).get());
@@ -2197,6 +2197,46 @@ private Function<Throwable, Throwable> translateException() {
 		};
 	}
 
+	/**
+	 * Prepare the {@link AggregationOperationContext} for a given aggregation by either returning the context itself it
+	 * is not {@literal null}, create a {@link TypeBasedAggregationOperationContext} if the aggregation contains type
+	 * information (is a {@link TypedAggregation}) or use the {@link Aggregation#DEFAULT_CONTEXT}.
+	 *
+	 * @param aggregation must not be {@literal null}.
+	 * @param context can be {@literal null}.
+	 * @return the root {@link AggregationOperationContext} to use.
+	 */
+	private AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
+			@Nullable AggregationOperationContext context) {
+
+		if (context != null) {
+			return context;
+		}
+
+		if (aggregation instanceof TypedAggregation) {
+			return new TypeBasedAggregationOperationContext(((TypedAggregation) aggregation).getInputType(), mappingContext,
+					queryMapper);
+		}
+
+		return Aggregation.DEFAULT_CONTEXT;
+	}
+
+	/**
+	 * Extract and map the aggregation pipeline.
+	 *
+	 * @param aggregation
+	 * @param context
+	 * @return
+	 */
+	private Document aggregationToPipeline(String inputCollectionName, Aggregation aggregation, AggregationOperationContext context) {
+
+		if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
+			return aggregation.toDocument(inputCollectionName, context);
+		}
+
+		return queryMapper.getMappedObject(aggregation.toDocument(inputCollectionName, context), Optional.empty());
+	}
+
 	/**
 	 * Tries to convert the given {@link RuntimeException} into a {@link DataAccessException} but returns the original
 	 * exception if the conversation failed. Thus allows safe re-throwing of the return value.

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

@@ -41,7 +41,7 @@ class AggregationOperationRenderer {
 	 * {@link Document} representation.
 	 *
 	 * @param operations must not be {@literal null}.
-	 * @param context must not be {@literal null}.
+	 * @param rootContext must not be {@literal null}.
 	 * @return the {@link List} of {@link Document}.
 	 */
 	static List<Document> toDocument(List<AggregationOperation> operations, AggregationOperationContext rootContext) {

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

@@ -660,10 +660,10 @@ public void readsPlainDbObjectById() {
 	@Test // DATAMONGO-1444
 	public void geoNear() {
 
-		List<Venue> venues = Arrays.asList(new Venue("Penn Station", -73.99408, 40.75057), //
-				new Venue("10gen Office", -73.99171, 40.738868), //
-				new Venue("Flatiron Building", -73.988135, 40.741404), //
-				new Venue("Maplewood, NJ", -74.2713, 40.73137));
+		List<Venue> venues = Arrays.asList(TestEntities.geolocation().pennStation(), //
+				TestEntities.geolocation().tenGenOffice(), //
+				TestEntities.geolocation().flatironBuilding(), //
+				TestEntities.geolocation().maplewoodNJ());
 
 		StepVerifier.create(template.insertAll(venues)).expectNextCount(4).verifyComplete();
 

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

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2018 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 java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A simple collection of grouped test entities used throughout the test suite.
+ *
+ * @author Christoph Strobl
+ */
+public class TestEntities {
+
+	private static final GeoEntities GEO = new GeoEntities();
+
+	public static GeoEntities geolocation() {
+		return GEO;
+	}
+
+	public static class GeoEntities {
+
+		/**
+		 * <pre>
+		 * X: -73.99408
+		 * Y: 40.75057
+		 * </pre>
+		 *
+		 * @return new {@link Venue}
+		 */
+		public Venue pennStation() {
+			return new Venue("Penn Station", -73.99408, 40.75057);
+		}
+
+		/**
+		 * <pre>
+		 * X: -73.99171
+		 * Y: 40.738868
+		 * </pre>
+		 *
+		 * @return new {@link Venue}
+		 */
+
+		public Venue tenGenOffice() {
+			return new Venue("10gen Office", -73.99171, 40.738868);
+		}
+
+		/**
+		 * <pre>
+		 * X: -73.988135
+		 * Y: 40.741404
+		 * </pre>
+		 *
+		 * @return new {@link Venue}
+		 */
+		public Venue flatironBuilding() {
+			return new Venue("Flatiron Building", -73.988135, 40.741404);
+		}
+
+		/**
+		 * <pre>
+		 * X: -74.2713
+		 * Y: 40.73137
+		 * </pre>
+		 *
+		 * @return new {@link Venue}
+		 */
+		public Venue maplewoodNJ() {
+			return new Venue("Maplewood, NJ", -74.2713, 40.73137);
+		}
+
+		public List<Venue> newYork() {
+
+			List<Venue> venues = new ArrayList<>();
+
+			venues.add(pennStation());
+			venues.add(tenGenOffice());
+			venues.add(flatironBuilding());
+			venues.add(new Venue("Players Club", -73.997812, 40.739128));
+			venues.add(new Venue("City Bakery ", -73.992491, 40.738673));
+			venues.add(new Venue("Splash Bar", -73.992491, 40.738673));
+			venues.add(new Venue("Momofuku Milk Bar", -73.985839, 40.731698));
+			venues.add(new Venue("Shake Shack", -73.98820, 40.74164));
+			venues.add(new Venue("Penn Station", -73.99408, 40.75057));
+			venues.add(new Venue("Empire State Building", -73.98602, 40.74894));
+			venues.add(new Venue("Ulaanbaatar, Mongolia", 106.9154, 47.9245));
+			venues.add(maplewoodNJ());
+
+			return venues;
+		}
+	}
+}

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

@@ -53,10 +53,13 @@
 import org.springframework.dao.DataAccessException;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.domain.Sort.Direction;
+import org.springframework.data.geo.Box;
 import org.springframework.data.geo.Metrics;
+import org.springframework.data.geo.Point;
 import org.springframework.data.mapping.MappingException;
 import org.springframework.data.mongodb.core.CollectionCallback;
 import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.TestEntities;
 import org.springframework.data.mongodb.core.Venue;
 import org.springframework.data.mongodb.core.aggregation.AggregationTests.CarDescriptor.Entry;
 import org.springframework.data.mongodb.core.aggregation.BucketAutoOperation.Granularities;
@@ -143,6 +146,8 @@ private void cleanDb() {
 		mongoTemplate.dropCollection(Sales2.class);
 		mongoTemplate.dropCollection(Employee.class);
 		mongoTemplate.dropCollection(Art.class);
+		mongoTemplate.dropCollection("personQueryTemp");
+		mongoTemplate.dropCollection(Venue.class);
 	}
 
 	/**
@@ -1484,9 +1489,8 @@ public void shouldRetrieveDateTimeFragementsCorrectly() throws Exception {
 	@Test // DATAMONGO-1127
 	public void shouldSupportGeoNearQueriesForAggregationWithDistanceField() {
 
-		mongoTemplate.insert(new Venue("Penn Station", -73.99408, 40.75057));
-		mongoTemplate.insert(new Venue("10gen Office", -73.99171, 40.738868));
-		mongoTemplate.insert(new Venue("Flatiron Building", -73.988135, 40.741404));
+		mongoTemplate.insertAll(Arrays.asList(TestEntities.geolocation().pennStation(),
+				TestEntities.geolocation().tenGenOffice(), TestEntities.geolocation().flatironBuilding()));
 
 		mongoTemplate.indexOps(Venue.class).ensureIndex(new GeospatialIndex("location"));
 
@@ -1898,6 +1902,36 @@ public void facetShouldCreateFacets() {
 		assertThat(categorizeByYear, hasSize(3));
 	}
 
+	@Test // DATAMONGO-1986
+	public void runMatchOperationCriteriaThroughQueryMapperForTypedAggregation() {
+
+		mongoTemplate.insertAll(TestEntities.geolocation().newYork());
+
+		Aggregation aggregation = newAggregation(Venue.class,
+				match(Criteria.where("location")
+						.within(new Box(new Point(-73.99756, 40.73083), new Point(-73.988135, 40.741404)))),
+				project("id", "location", "name"));
+
+		AggregationResults<Document> groupResults = mongoTemplate.aggregate(aggregation, "newyork", Document.class);
+
+		assertThat(groupResults.getMappedResults().size(), is(4));
+	}
+
+	@Test // DATAMONGO-1986
+	public void runMatchOperationCriteriaThroughQueryMapperForUntypedAggregation() {
+
+		mongoTemplate.insertAll(TestEntities.geolocation().newYork());
+
+		Aggregation aggregation = newAggregation(
+				match(Criteria.where("location")
+						.within(new Box(new Point(-73.99756, 40.73083), new Point(-73.988135, 40.741404)))),
+				project("id", "location", "name"));
+
+		AggregationResults<Document> groupResults = mongoTemplate.aggregate(aggregation, "newyork", Document.class);
+
+		assertThat(groupResults.getMappedResults().size(), is(4));
+	}
+
 	private void createUsersWithReferencedPersons() {
 
 		mongoTemplate.dropCollection(User.class);