DATAMONGO-2635 - Enforce aggregation pipeline mapping.

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-mongodb-parent</artifactId>
-	<version>3.2.0-SNAPSHOT</version>
+	<version>3.2.0-DATAMONGO-2635-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data MongoDB</name>

spring-data-mongodb-benchmarks/pom.xml

@@ -7,7 +7,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>3.2.0-SNAPSHOT</version>
+		<version>3.2.0-DATAMONGO-2635-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb-distribution/pom.xml

@@ -14,7 +14,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>3.2.0-SNAPSHOT</version>
+		<version>3.2.0-DATAMONGO-2635-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb/pom.xml

@@ -11,7 +11,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>3.2.0-SNAPSHOT</version>
+		<version>3.2.0-DATAMONGO-2635-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

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

@@ -77,7 +77,7 @@ AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
 		}
 
 		if (!(aggregation instanceof TypedAggregation)) {
-			return Aggregation.DEFAULT_CONTEXT;
+			return new RelaxedTypeBasedAggregationOperationContext(Object.class, mappingContext, queryMapper);
 		}
 
 		Class<?> inputType = ((TypedAggregation) aggregation).getInputType();
@@ -98,7 +98,7 @@ AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
 	 */
 	List<Document> createPipeline(Aggregation aggregation, AggregationOperationContext context) {
 
-		if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
+		if (ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
 			return aggregation.toPipeline(context);
 		}
 

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

@@ -707,10 +707,9 @@ Document getMappedShardKey(MongoPersistentEntity<?> entity) {
 		 */
 		List<Document> getUpdatePipeline(@Nullable Class<?> domainType) {
 
-			AggregationOperationContext context = domainType != null
-					? new RelaxedTypeBasedAggregationOperationContext(domainType, mappingContext, queryMapper)
-					: Aggregation.DEFAULT_CONTEXT;
+			Class<?> type = domainType != null ? domainType : Object.class;
 
+			AggregationOperationContext context = new RelaxedTypeBasedAggregationOperationContext(type, mappingContext, queryMapper);
 			return aggregationUtil.createPipeline((AggregationUpdate) update, context);
 		}
 

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

@@ -17,6 +17,7 @@
 
 import static org.springframework.data.mongodb.core.query.SerializationUtils.*;
 
+import org.springframework.data.mongodb.core.aggregation.RelaxedTypeBasedAggregationOperationContext;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.util.function.Tuple2;
@@ -2112,7 +2113,7 @@ List<Document> prepareFilter(ChangeStreamOptions options) {
 			AggregationOperationContext context = agg instanceof TypedAggregation
 					? new TypeBasedAggregationOperationContext(((TypedAggregation<?>) agg).getInputType(),
 							getConverter().getMappingContext(), queryMapper)
-					: Aggregation.DEFAULT_CONTEXT;
+					: new RelaxedTypeBasedAggregationOperationContext(Object.class, mappingContext, queryMapper);
 
 			return agg.toPipeline(new PrefixingDelegatingAggregationOperationContext(context, "fullDocument",
 					Arrays.asList("operationType", "fullDocument", "documentKey", "updateDescription", "ns")));

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

@@ -1450,6 +1450,14 @@ private Object renderFieldValue(AggregationOperationContext context) {
 						return field.getTarget();
 					}
 
+					if(field.getTarget().equals(Fields.UNDERSCORE_ID)) {
+						try {
+							return context.getReference(field).getReferenceValue();
+						} catch (java.lang.IllegalArgumentException e) {
+							return Fields.UNDERSCORE_ID_REF;
+						}
+					}
+
 					// check whether referenced field exists in the context
 					return context.getReference(field).getReferenceValue();
 

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

@@ -21,6 +21,7 @@
 import java.util.List;
 
 import org.bson.Document;
+import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.DirectFieldReference;
@@ -29,6 +30,7 @@
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.data.util.Lazy;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
@@ -46,6 +48,7 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio
 	private final Class<?> type;
 	private final MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
 	private final QueryMapper mapper;
+	private final Lazy<MongoPersistentEntity<?>> entity;
 
 	/**
 	 * Creates a new {@link TypeBasedAggregationOperationContext} for the given type, {@link MappingContext} and
@@ -65,6 +68,7 @@ public TypeBasedAggregationOperationContext(Class<?> type,
 		this.type = type;
 		this.mappingContext = mappingContext;
 		this.mapper = mapper;
+		this.entity = Lazy.of(() -> mappingContext.getPersistentEntity(type));
 	}
 
 	/*
@@ -151,10 +155,14 @@ public AggregationOperationContext continueOnMissingFieldReference(Class<?> type
 
 	protected FieldReference getReferenceFor(Field field) {
 
+		if(entity.getNullable() == null) {
+			return new DirectFieldReference(new ExposedField(field, true));
+		}
+
 		PersistentPropertyPath<MongoPersistentProperty> propertyPath = mappingContext
-				.getPersistentPropertyPath(field.getTarget(), type);
+					.getPersistentPropertyPath(field.getTarget(), type);
 		Field mappedField = field(field.getName(),
-				propertyPath.toDotPath(MongoPersistentProperty.PropertyToFieldNameConverter.INSTANCE));
+					propertyPath.toDotPath(MongoPersistentProperty.PropertyToFieldNameConverter.INSTANCE));
 
 		return new DirectFieldReference(new ExposedField(mappedField, true));
 	}

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

@@ -142,12 +142,8 @@ public Document toDocument(AggregationOperationContext context) {
 
 	private AggregationOperationContext computeContext(AggregationOperationContext source) {
 
-		if (domainType == null) {
-			return Aggregation.DEFAULT_CONTEXT;
-		}
-
 		if (source instanceof TypeBasedAggregationOperationContext) {
-			return ((TypeBasedAggregationOperationContext) source).continueOnMissingFieldReference(domainType);
+			return ((TypeBasedAggregationOperationContext) source).continueOnMissingFieldReference(domainType != null ? domainType : Object.class);
 		}
 
 		if (source instanceof ExposedFieldsAggregationOperationContext) {

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

@@ -1928,6 +1928,22 @@ void skipOutputDoesNotReadBackAggregationResults() {
 		assertThat(results.getRawResults()).isEmpty();
 	}
 
+	@Test // DATAMONGO-2635
+	void mapsEnumsInMatchClauseUsingInCriteriaCorrectly() {
+
+		WithEnum source = new WithEnum();
+		source.enumValue = MyEnum.TWO;
+		source.id = "id-1";
+
+		mongoTemplate.save(source);
+
+		Aggregation agg = newAggregation(match(where("enumValue").in(Collections.singletonList(MyEnum.TWO))));
+
+		AggregationResults<Document> results = mongoTemplate.aggregate(agg, mongoTemplate.getCollectionName(WithEnum.class),
+				Document.class);
+		assertThat(results.getMappedResults()).hasSize(1);
+	}
+
 	private void createUsersWithReferencedPersons() {
 
 		mongoTemplate.dropCollection(User.class);
@@ -2240,4 +2256,15 @@ static class ComplexId {
 		String p1;
 		String p2;
 	}
+
+	static enum MyEnum {
+		ONE, TWO
+	}
+
+	@lombok.Data
+	static class WithEnum {
+
+		@Id String id;
+		MyEnum enumValue;
+	}
 }

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

@@ -29,6 +29,11 @@
 import org.junit.jupiter.api.Test;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mongodb.core.aggregation.ConditionalOperators.Cond;
+import org.springframework.data.mongodb.core.aggregation.ProjectionOperationUnitTests.BookWithFieldAnnotation;
+import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
+import org.springframework.data.mongodb.core.convert.NoOpDbRefResolver;
+import org.springframework.data.mongodb.core.convert.QueryMapper;
+import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.core.query.Criteria;
 
 /**
@@ -582,6 +587,17 @@ void shouldAllowInternalThisAndValueReferences() {
 				"{\"attributeRecordArrays\": {\"$reduce\": {\"input\": \"$attributeRecordArrays\", \"initialValue\": [], \"in\": {\"$concatArrays\": [\"$$value\", \"$$this\"]}}}}"));
 	}
 
+	@Test // DATAMONGO-2644
+	void projectOnIdIsAlwaysValid() {
+
+		MongoMappingContext mappingContext = new MongoMappingContext();
+		Document target = new Aggregation(bucket("start"), project("_id")).toDocument("collection-1",
+				new RelaxedTypeBasedAggregationOperationContext(BookWithFieldAnnotation.class, mappingContext,
+						new QueryMapper(new MappingMongoConverter(NoOpDbRefResolver.INSTANCE, mappingContext))));
+
+		assertThat(extractPipelineElement(target, 1, "$project")).isEqualTo(Document.parse(" { \"_id\" : \"$_id\" }"));
+	}
+
 	private Document extractPipelineElement(Document agg, int index, String operation) {
 
 		List<Document> pipeline = (List<Document>) agg.get("pipeline");