Configure relaxed/strict field lookup when creating AggregationOperationContext.

Closes #4714
Original pull request: #4720

Author: christophstrobl

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

@@ -50,6 +50,7 @@ static List<Document> toDocument(List<AggregationOperation> operations, Aggregat
 		List<Document> operationDocuments = new ArrayList<Document>(operations.size());
 
 		AggregationOperationContext contextToUse = rootContext;
+		boolean relaxed = rootContext instanceof RelaxedTypeBasedAggregationOperationContext;
 
 		for (AggregationOperation operation : operations) {
 
@@ -60,12 +61,13 @@ static List<Document> toDocument(List<AggregationOperation> operations, Aggregat
 				ExposedFields fields = exposedFieldsOperation.getFields();
 
 				if (operation instanceof InheritsFieldsAggregationOperation || exposedFieldsOperation.inheritsFields()) {
-					contextToUse = new InheritingExposedFieldsAggregationOperationContext(fields, contextToUse);
+					contextToUse = new InheritingExposedFieldsAggregationOperationContext(fields, contextToUse, relaxed);
 				} else {
 					contextToUse = fields.exposesNoFields() ? DEFAULT_CONTEXT
-							: new ExposedFieldsAggregationOperationContext(fields, contextToUse);
+							: new ExposedFieldsAggregationOperationContext(fields, contextToUse, relaxed);
 				}
 			}
+
 		}
 
 		return operationDocuments;

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

@@ -688,7 +688,7 @@ private Document toFilter(ExposedFields exposedFields, AggregationOperationConte
 
 			Document filterExpression = new Document();
 			InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-					exposedFields, context);
+					exposedFields, context, false);
 
 			filterExpression.putAll(context.getMappedObject(new Document("input", getMappedInput(context))));
 			filterExpression.put("as", as.getTarget());

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

@@ -50,7 +50,7 @@ protected DocumentEnhancingOperation(Map<Object, Object> source) {
 	public Document toDocument(AggregationOperationContext context) {
 
 		InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-				exposedFields, context);
+				exposedFields, context, false);
 
 		if (valueMap.size() == 1) {
 			return context.getMappedObject(

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

@@ -37,6 +37,7 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 
 	private final ExposedFields exposedFields;
 	private final AggregationOperationContext rootContext;
+	private final boolean relaxedFieldLookup;
 
 	/**
 	 * Creates a new {@link ExposedFieldsAggregationOperationContext} from the given {@link ExposedFields}. Uses the given
@@ -46,13 +47,14 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 	 * @param rootContext must not be {@literal null}.
 	 */
 	public ExposedFieldsAggregationOperationContext(ExposedFields exposedFields,
-			AggregationOperationContext rootContext) {
+			AggregationOperationContext rootContext, boolean relaxedFieldLookup) {
 
 		Assert.notNull(exposedFields, "ExposedFields must not be null");
 		Assert.notNull(rootContext, "RootContext must not be null");
 
 		this.exposedFields = exposedFields;
 		this.rootContext = rootContext;
+		this.relaxedFieldLookup = relaxedFieldLookup;
 	}
 
 	@Override
@@ -87,7 +89,7 @@ public Fields getFields(Class<?> type) {
 	 * @param name must not be {@literal null}.
 	 * @return
 	 */
-	private FieldReference getReference(@Nullable Field field, String name) {
+	protected FieldReference getReference(@Nullable Field field, String name) {
 
 		Assert.notNull(name, "Name must not be null");
 
@@ -96,12 +98,10 @@ private FieldReference getReference(@Nullable Field field, String name) {
 			return exposedField;
 		}
 
-		if (rootContext instanceof RelaxedTypeBasedAggregationOperationContext) {
-
+		if(relaxedFieldLookup) {
 			if (field != null) {
 				return new DirectFieldReference(new ExposedField(field, true));
 			}
-
 			return new DirectFieldReference(new ExposedField(name, true));
 		}
 
@@ -156,4 +156,12 @@ AggregationOperationContext getRootContext() {
 	public CodecRegistry getCodecRegistry() {
 		return getRootContext().getCodecRegistry();
 	}
+
+	@Override
+	public AggregationOperationContext continueOnMissingFieldReference() {
+		if(relaxedFieldLookup) {
+			return this;
+		}
+		return new ExposedFieldsAggregationOperationContext(exposedFields, rootContext, true);
+	}
 }

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

@@ -38,9 +38,9 @@ class InheritingExposedFieldsAggregationOperationContext extends ExposedFieldsAg
 	 * @param previousContext must not be {@literal null}.
 	 */
 	public InheritingExposedFieldsAggregationOperationContext(ExposedFields exposedFields,
-			AggregationOperationContext previousContext) {
+			AggregationOperationContext previousContext, boolean continueOnMissingFieldReference) {
 
-		super(exposedFields, previousContext);
+		super(exposedFields, previousContext, continueOnMissingFieldReference);
 
 		this.previousContext = previousContext;
 	}

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

@@ -171,7 +171,7 @@ private Document toMap(ExposedFields exposedFields, AggregationOperationContext
 
 			Document map = new Document();
 			InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-					exposedFields, context);
+					exposedFields, context, false);
 
 			Document input;
 			if (sourceArray instanceof Field field) {
@@ -308,15 +308,16 @@ private Document toLet(ExposedFields exposedFields, AggregationOperationContext
 
 			Document letExpression = new Document();
 			Document mappedVars = new Document();
-			InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-					exposedFields, context);
 
 			for (ExpressionVariable var : this.vars) {
 				mappedVars.putAll(getMappedVariable(var, context));
 			}
 
 			letExpression.put("vars", mappedVars);
 			if (expression != null) {
+
+				InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
+					exposedFields, context, false);
 				letExpression.put("in", getMappedIn(operationContext));
 			}
 

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

@@ -15,15 +15,26 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
-import static org.assertj.core.api.Assertions.*;
-import static org.mockito.Mockito.*;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.springframework.data.domain.Sort.Direction.DESC;
+import static org.springframework.data.mongodb.core.aggregation.Aggregation.project;
+import static org.springframework.data.mongodb.core.aggregation.Aggregation.sort;
 
 import java.util.List;
 
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
+import org.springframework.data.annotation.Id;
 import org.springframework.data.mongodb.core.aggregation.FieldsExposingAggregationOperation.InheritsFieldsAggregationOperation;
+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.test.util.MongoTestMappingContext;
 
 /**
  * @author Christoph Strobl
@@ -115,4 +126,34 @@ void inheritingFieldsExposingAggregationOperationForcesNewContextForNextStageKee
 				.extracting("previousContext").isSameAs(captor.getAllValues().get(1));
 	}
 
+
+
+	record TestRecord(@Id String field1, String field2, LayerOne layerOne) {
+		record LayerOne(List<LayerTwo> layerTwo) {
+		}
+
+		record LayerTwo(LayerThree layerThree) {
+		}
+
+		record LayerThree(int fieldA, int fieldB)
+		{}
+	}
+
+	@Test
+	void xxx() {
+
+		MongoTestMappingContext ctx = new MongoTestMappingContext(cfg -> {
+			cfg.initialEntitySet(TestRecord.class);
+		});
+
+		MappingMongoConverter mongoConverter = new MappingMongoConverter(NoOpDbRefResolver.INSTANCE, ctx);
+
+		Aggregation agg = Aggregation.newAggregation(
+			Aggregation.unwind("layerOne.layerTwo"),
+			project().and("layerOne.layerTwo.layerThree").as("layerOne.layerThree"),
+			sort(DESC, "layerOne.layerThree.fieldA")
+		);
+
+		AggregationOperationRenderer.toDocument(agg.getPipeline().getOperations(), new RelaxedTypeBasedAggregationOperationContext(TestRecord.class, ctx, new QueryMapper(mongoConverter)));
+	}
 }