DATAMONGO-2635 - Enforce aggregation pipeline mapping.

Avoid using the Aggregation.DEFAULT_CONTEXT which does not map contained values to the according MongoDB representation. We now use a relaxed aggregation context, preserving given field names, where possible.

Original pull request: #890.

Author: christophstrobl

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/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;
+	}
 }