Retain Field Lookup Policy when exposing aggregation fields.

Introduce FieldLookupPolicy and methods to create field-exposing/inheriting AggregationOperationContexts.

Move off RelaxedTypeBasedAggregationOperationContext.

See #4714
Original pull request: #4720

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

@@ -16,15 +16,14 @@
 package org.springframework.data.mongodb.core;
 
 import java.util.List;
-import java.util.Optional;
-import java.util.stream.Collectors;
 
 import org.bson.Document;
+
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.aggregation.Aggregation;
 import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.AggregationOptions.DomainTypeMapping;
-import org.springframework.data.mongodb.core.aggregation.RelaxedTypeBasedAggregationOperationContext;
+import org.springframework.data.mongodb.core.aggregation.FieldLookupPolicy;
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
@@ -52,8 +51,8 @@ class AggregationUtil {
 
 		this.queryMapper = queryMapper;
 		this.mappingContext = mappingContext;
-		this.untypedMappingContext = Lazy
-				.of(() -> new RelaxedTypeBasedAggregationOperationContext(Object.class, mappingContext, queryMapper));
+		this.untypedMappingContext = Lazy.of(() -> new TypeBasedAggregationOperationContext(Object.class, mappingContext,
+				queryMapper, FieldLookupPolicy.relaxed()));
 	}
 
 	AggregationOperationContext createAggregationContext(Aggregation aggregation, @Nullable Class<?> inputType) {
@@ -64,27 +63,18 @@ class AggregationUtil {
 			return Aggregation.DEFAULT_CONTEXT;
 		}
 
-		if (!(aggregation instanceof TypedAggregation)) {
-
-			if(inputType == null) {
-				return untypedMappingContext.get();
-			}
-
-			if (domainTypeMapping == DomainTypeMapping.STRICT
-					&& !aggregation.getPipeline().containsUnionWith()) {
-				return new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper);
-			}
+		FieldLookupPolicy lookupPolicy = domainTypeMapping == DomainTypeMapping.STRICT
+				&& !aggregation.getPipeline().containsUnionWith() ? FieldLookupPolicy.strict() : FieldLookupPolicy.relaxed();
 
-			return new RelaxedTypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper);
+		if (aggregation instanceof TypedAggregation<?> ta) {
+			return new TypeBasedAggregationOperationContext(ta.getInputType(), mappingContext, queryMapper, lookupPolicy);
 		}
 
-		inputType = ((TypedAggregation<?>) aggregation).getInputType();
-		if (domainTypeMapping == DomainTypeMapping.STRICT
-				&& !aggregation.getPipeline().containsUnionWith()) {
-			return new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper);
+		if (inputType == null) {
+			return untypedMappingContext.get();
 		}
 
-		return new RelaxedTypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper);
+		return new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper, lookupPolicy);
 	}
 
 	/**
@@ -109,9 +99,4 @@ class AggregationUtil {
 		return aggregation.toDocument(collection, context);
 	}
 
-	private List<Document> mapAggregationPipeline(List<Document> pipeline) {
-
-		return pipeline.stream().map(val -> queryMapper.getMappedObject(val, Optional.empty()))
-				.collect(Collectors.toList());
-	}
 }

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

@@ -35,6 +35,7 @@ import com.mongodb.MongoClientSettings;
  *
  * @author Oliver Gierke
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @since 1.3
  */
 public interface AggregationOperationContext extends CodecRegistryProvider {
@@ -107,14 +108,46 @@ public interface AggregationOperationContext extends CodecRegistryProvider {
 				.toArray(String[]::new));
 	}
 
+	/**
+	 * Create a nested {@link AggregationOperationContext} from this context that exposes {@link ExposedFields fields}.
+	 * <p>
+	 * Implementations of {@link AggregationOperationContext} retain their {@link FieldLookupPolicy}. If no policy is
+	 * specified, then lookup defaults to {@link FieldLookupPolicy#strict()}.
+	 *
+	 * @param fields the fields to expose, must not be {@literal null}.
+	 * @return the new {@link AggregationOperationContext} exposing {@code fields}.
+	 * @since 4.3.1
+	 */
+	default AggregationOperationContext expose(ExposedFields fields) {
+		return new ExposedFieldsAggregationOperationContext(fields, this, FieldLookupPolicy.strict());
+	}
+
+	/**
+	 * Create a nested {@link AggregationOperationContext} from this context that inherits exposed fields from this
+	 * context and exposes {@link ExposedFields fields}.
+	 * <p>
+	 * Implementations of {@link AggregationOperationContext} retain their {@link FieldLookupPolicy}. If no policy is
+	 * specified, then lookup defaults to {@link FieldLookupPolicy#strict()}.
+	 *
+	 * @param fields the fields to expose, must not be {@literal null}.
+	 * @return the new {@link AggregationOperationContext} exposing {@code fields}.
+	 * @since 4.3.1
+	 */
+	default AggregationOperationContext inheritAndExpose(ExposedFields fields) {
+		return new InheritingExposedFieldsAggregationOperationContext(fields, this, FieldLookupPolicy.strict());
+	}
+
 	/**
 	 * This toggle allows the {@link AggregationOperationContext context} to use any given field name without checking for
-	 * its existence. Typically the {@link AggregationOperationContext} fails when referencing unknown fields, those that
+	 * its existence. Typically, the {@link AggregationOperationContext} fails when referencing unknown fields, those that
 	 * are not present in one of the previous stages or the input source, throughout the pipeline.
 	 *
 	 * @return a more relaxed {@link AggregationOperationContext}.
 	 * @since 3.0
+	 * @deprecated since 4.3.1, {@link FieldLookupPolicy} should be specified explicitly when creating the
+	 *             AggregationOperationContext.
 	 */
+	@Deprecated(since = "4.3.1", forRemoval = true)
 	default AggregationOperationContext continueOnMissingFieldReference() {
 		return this;
 	}
@@ -123,4 +156,5 @@ public interface AggregationOperationContext extends CodecRegistryProvider {
 	default CodecRegistry getCodecRegistry() {
 		return MongoClientSettings.getDefaultCodecRegistry();
 	}
+
 }

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

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

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

@@ -687,8 +687,7 @@ public class ArrayOperators {
 		private Document toFilter(ExposedFields exposedFields, AggregationOperationContext context) {
 
 			Document filterExpression = new Document();
-			InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-					exposedFields, context, false);
+			AggregationOperationContext operationContext = context.inheritAndExpose(exposedFields);
 
 			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

@@ -49,8 +49,7 @@ abstract class DocumentEnhancingOperation implements InheritsFieldsAggregationOp
 	@Override
 	public Document toDocument(AggregationOperationContext context) {
 
-		InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-				exposedFields, context, false);
+		AggregationOperationContext operationContext = context.inheritAndExpose(exposedFields);
 
 		if (valueMap.size() == 1) {
 			return context.getMappedObject(

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

@@ -17,6 +17,7 @@ package org.springframework.data.mongodb.core.aggregation;
 
 import org.bson.Document;
 import org.bson.codecs.configuration.CodecRegistry;
+
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.DirectFieldReference;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
@@ -37,7 +38,7 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 
 	private final ExposedFields exposedFields;
 	private final AggregationOperationContext rootContext;
-	private final boolean relaxedFieldLookup;
+	private final FieldLookupPolicy lookupPolicy;
 
 	/**
 	 * Creates a new {@link ExposedFieldsAggregationOperationContext} from the given {@link ExposedFields}. Uses the given
@@ -45,16 +46,18 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 	 *
 	 * @param exposedFields must not be {@literal null}.
 	 * @param rootContext must not be {@literal null}.
+	 * @param lookupPolicy must not be {@literal null}.
 	 */
-	public ExposedFieldsAggregationOperationContext(ExposedFields exposedFields,
-			AggregationOperationContext rootContext, boolean relaxedFieldLookup) {
+	public ExposedFieldsAggregationOperationContext(ExposedFields exposedFields, AggregationOperationContext rootContext,
+			FieldLookupPolicy lookupPolicy) {
 
 		Assert.notNull(exposedFields, "ExposedFields must not be null");
 		Assert.notNull(rootContext, "RootContext must not be null");
+		Assert.notNull(lookupPolicy, "FieldLookupPolicy must not be null");
 
 		this.exposedFields = exposedFields;
 		this.rootContext = rootContext;
-		this.relaxedFieldLookup = relaxedFieldLookup;
+		this.lookupPolicy = lookupPolicy;
 	}
 
 	@Override
@@ -89,7 +92,7 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 	 * @param name must not be {@literal null}.
 	 * @return
 	 */
-	protected FieldReference getReference(@Nullable Field field, String name) {
+	private FieldReference getReference(@Nullable Field field, String name) {
 
 		Assert.notNull(name, "Name must not be null");
 
@@ -98,14 +101,15 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 			return exposedField;
 		}
 
-		if(relaxedFieldLookup) {
-			if (field != null) {
-				return new DirectFieldReference(new ExposedField(field, true));
-			}
-			return new DirectFieldReference(new ExposedField(name, true));
+		if (lookupPolicy.isStrict()) {
+			throw new IllegalArgumentException(String.format("Invalid reference '%s'", name));
 		}
 
-		throw new IllegalArgumentException(String.format("Invalid reference '%s'", name));
+		if (field != null) {
+			return new DirectFieldReference(new ExposedField(field, true));
+		}
+
+		return new DirectFieldReference(new ExposedField(name, true));
 	}
 
 	/**
@@ -158,10 +162,22 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 	}
 
 	@Override
+	@Deprecated(since = "4.3.1", forRemoval = true)
 	public AggregationOperationContext continueOnMissingFieldReference() {
-		if(relaxedFieldLookup) {
+		if (!lookupPolicy.isStrict()) {
 			return this;
 		}
-		return new ExposedFieldsAggregationOperationContext(exposedFields, rootContext, true);
+		return new ExposedFieldsAggregationOperationContext(exposedFields, rootContext, FieldLookupPolicy.relaxed());
 	}
+
+	@Override
+	public AggregationOperationContext expose(ExposedFields fields) {
+		return new ExposedFieldsAggregationOperationContext(fields, this, lookupPolicy);
+	}
+
+	@Override
+	public AggregationOperationContext inheritAndExpose(ExposedFields fields) {
+		return new InheritingExposedFieldsAggregationOperationContext(fields, this, lookupPolicy);
+	}
+
 }

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2024 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
+ *
+ *      https://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.aggregation;
+
+/**
+ * Lookup policy for aggregation fields. Allows strict lookups that fail if the field is absent or relaxed ones that
+ * pass-thru the requested field even if we have to assume that the field isn't present because of the limited scope of
+ * our input.
+ *
+ * @author Mark Paluch
+ * @since 4.3.1
+ */
+public abstract class FieldLookupPolicy {
+
+	private static final FieldLookupPolicy STRICT = new FieldLookupPolicy() {
+		@Override
+		boolean isStrict() {
+			return true;
+		}
+	};
+
+	private static final FieldLookupPolicy RELAXED = new FieldLookupPolicy() {
+		@Override
+		boolean isStrict() {
+			return false;
+		}
+	};
+
+	private FieldLookupPolicy() {}
+
+	/**
+	 * @return a relaxed lookup policy.
+	 */
+	public static FieldLookupPolicy relaxed() {
+		return RELAXED;
+	}
+
+	/**
+	 * @return a strict lookup policy.
+	 */
+	public static FieldLookupPolicy strict() {
+		return STRICT;
+	}
+
+	/**
+	 * @return {@code true} if the policy uses a strict lookup; {@code false} to allow references to fields that cannot be
+	 *         determined to be exactly present.
+	 */
+	abstract boolean isStrict();
+
+}

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

@@ -17,6 +17,7 @@ package org.springframework.data.mongodb.core.aggregation;
 
 import org.bson.Document;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
+import org.springframework.lang.Nullable;
 
 /**
  * {@link ExposedFieldsAggregationOperationContext} that inherits fields from its parent
@@ -36,11 +37,12 @@ class InheritingExposedFieldsAggregationOperationContext extends ExposedFieldsAg
 	 *
 	 * @param exposedFields must not be {@literal null}.
 	 * @param previousContext must not be {@literal null}.
+	 * @param lookupPolicy must not be {@literal null}.
 	 */
 	public InheritingExposedFieldsAggregationOperationContext(ExposedFields exposedFields,
-			AggregationOperationContext previousContext, boolean continueOnMissingFieldReference) {
+			AggregationOperationContext previousContext, FieldLookupPolicy lookupPolicy) {
 
-		super(exposedFields, previousContext, continueOnMissingFieldReference);
+		super(exposedFields, previousContext, lookupPolicy);
 
 		this.previousContext = previousContext;
 	}
@@ -51,7 +53,7 @@ class InheritingExposedFieldsAggregationOperationContext extends ExposedFieldsAg
 	}
 
 	@Override
-	protected FieldReference resolveExposedField(Field field, String name) {
+	protected FieldReference resolveExposedField(@Nullable Field field, String name) {
 
 		FieldReference fieldReference = super.resolveExposedField(field, name);
 		if (fieldReference != null) {

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

@@ -15,12 +15,8 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
-import org.springframework.data.mapping.MappingException;
 import org.springframework.data.mapping.context.InvalidPersistentPropertyPath;
 import org.springframework.data.mapping.context.MappingContext;
-import org.springframework.data.mongodb.core.aggregation.ExposedFields.DirectFieldReference;
-import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField;
-import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
@@ -31,7 +27,9 @@ import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
  *
  * @author Christoph Strobl
  * @since 3.0
+ * @deprecated since 4.3.1
  */
+@Deprecated(since = "4.3.1")
 public class RelaxedTypeBasedAggregationOperationContext extends TypeBasedAggregationOperationContext {
 
 	/**
@@ -44,16 +42,6 @@ public class RelaxedTypeBasedAggregationOperationContext extends TypeBasedAggreg
 	 */
 	public RelaxedTypeBasedAggregationOperationContext(Class<?> type,
 			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext, QueryMapper mapper) {
-		super(type, mappingContext, mapper);
-	}
-
-	@Override
-	protected FieldReference getReferenceFor(Field field) {
-
-		try {
-			return super.getReferenceFor(field);
-		} catch (MappingException e) {
-			return new DirectFieldReference(new ExposedField(field, true));
-		}
+		super(type, mappingContext, mapper, FieldLookupPolicy.relaxed());
 	}
 }

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

@@ -21,8 +21,9 @@ import java.util.ArrayList;
 import java.util.List;
 
 import org.bson.Document;
-
 import org.bson.codecs.configuration.CodecRegistry;
+
+import org.springframework.data.mapping.MappingException;
 import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.DirectFieldReference;
@@ -50,6 +51,7 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio
 	private final MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
 	private final QueryMapper mapper;
 	private final Lazy<MongoPersistentEntity<?>> entity;
+	private final FieldLookupPolicy lookupPolicy;
 
 	/**
 	 * Creates a new {@link TypeBasedAggregationOperationContext} for the given type, {@link MappingContext} and
@@ -61,15 +63,33 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio
 	 */
 	public TypeBasedAggregationOperationContext(Class<?> type,
 			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext, QueryMapper mapper) {
+		this(type, mappingContext, mapper, FieldLookupPolicy.strict());
+	}
+
+	/**
+	 * Creates a new {@link TypeBasedAggregationOperationContext} for the given type, {@link MappingContext} and
+	 * {@link QueryMapper}.
+	 *
+	 * @param type must not be {@literal null}.
+	 * @param mappingContext must not be {@literal null}.
+	 * @param mapper must not be {@literal null}.
+	 * @param lookupPolicy must not be {@literal null}.
+	 * @since 4.3.1
+	 */
+	public TypeBasedAggregationOperationContext(Class<?> type,
+			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext, QueryMapper mapper,
+			FieldLookupPolicy lookupPolicy) {
 
 		Assert.notNull(type, "Type must not be null");
 		Assert.notNull(mappingContext, "MappingContext must not be null");
 		Assert.notNull(mapper, "QueryMapper must not be null");
+		Assert.notNull(lookupPolicy, "FieldLookupPolicy must not be null");
 
 		this.type = type;
 		this.mappingContext = mappingContext;
 		this.mapper = mapper;
 		this.entity = Lazy.of(() -> mappingContext.getPersistentEntity(type));
+		this.lookupPolicy = lookupPolicy;
 	}
 
 	@Override
@@ -113,6 +133,7 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio
 	}
 
 	@Override
+	@Deprecated(since = "4.3.1", forRemoval = true)
 	public AggregationOperationContext continueOnMissingFieldReference() {
 		return continueOnMissingFieldReference(type);
 	}
@@ -128,19 +149,43 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio
 	 * @see RelaxedTypeBasedAggregationOperationContext
 	 */
 	public AggregationOperationContext continueOnMissingFieldReference(Class<?> type) {
-		return new RelaxedTypeBasedAggregationOperationContext(type, mappingContext, mapper);
+		return new TypeBasedAggregationOperationContext(type, mappingContext, mapper, FieldLookupPolicy.relaxed());
+	}
+
+	@Override
+	public AggregationOperationContext expose(ExposedFields fields) {
+		return new ExposedFieldsAggregationOperationContext(fields, this, lookupPolicy);
+	}
+
+	@Override
+	public AggregationOperationContext inheritAndExpose(ExposedFields fields) {
+		return new InheritingExposedFieldsAggregationOperationContext(fields, this, lookupPolicy);
 	}
 
 	protected FieldReference getReferenceFor(Field field) {
 
-		if(entity.getNullable() == null || AggregationVariable.isVariable(field)) {
+		try {
+			return doGetFieldReference(field);
+		} catch (MappingException e) {
+
+			if (lookupPolicy.isStrict()) {
+				throw e;
+			}
+
+			return new DirectFieldReference(new ExposedField(field, true));
+		}
+	}
+
+	private DirectFieldReference doGetFieldReference(Field field) {
+
+		if (entity.getNullable() == null || AggregationVariable.isVariable(field)) {
 			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/VariableOperators.java

@@ -170,8 +170,7 @@ public class VariableOperators {
 		private Document toMap(ExposedFields exposedFields, AggregationOperationContext context) {
 
 			Document map = new Document();
-			InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-					exposedFields, context, false);
+			AggregationOperationContext operationContext = context.inheritAndExpose(exposedFields);
 
 			Document input;
 			if (sourceArray instanceof Field field) {
@@ -316,8 +315,7 @@ public class VariableOperators {
 			letExpression.put("vars", mappedVars);
 			if (expression != null) {
 
-				InheritingExposedFieldsAggregationOperationContext operationContext = new InheritingExposedFieldsAggregationOperationContext(
-					exposedFields, context, false);
+				AggregationOperationContext operationContext = context.inheritAndExpose(exposedFields);
 				letExpression.put("in", getMappedIn(operationContext));
 			}
 

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

@@ -558,7 +558,7 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 			protected <O> AggregationResults<O> doAggregate(Aggregation aggregation, String collectionName,
 					Class<O> outputType, AggregationOperationContext context) {
 
-				assertThat(context).isInstanceOf(RelaxedTypeBasedAggregationOperationContext.class);
+				assertThat(ReflectionTestUtils.getField(context, "lookupPolicy")).isEqualTo(FieldLookupPolicy.relaxed());
 				return super.doAggregate(aggregation, collectionName, outputType, context);
 			}
 		};

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

@@ -25,12 +25,13 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.MongoDatabaseFactory;
 import org.springframework.data.mongodb.core.QueryOperations.AggregationDefinition;
 import org.springframework.data.mongodb.core.aggregation.Aggregation;
 import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
-import org.springframework.data.mongodb.core.aggregation.RelaxedTypeBasedAggregationOperationContext;
+import org.springframework.data.mongodb.core.aggregation.FieldLookupPolicy;
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.convert.UpdateMapper;
@@ -38,6 +39,7 @@ import org.springframework.data.mongodb.core.mapping.MongoId;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.test.util.ReflectionTestUtils;
 
 /**
  * Unit tests for {@link QueryOperations}.
@@ -72,27 +74,33 @@ class QueryOperationsUnitTests {
 	void createAggregationContextUsesRelaxedOneForUntypedAggregationsWhenNoInputTypeProvided() {
 
 		Aggregation aggregation = Aggregation.newAggregation(Aggregation.project("name"));
-		AggregationDefinition ctx = queryOperations.createAggregation(aggregation, (Class<?>) null);
+		AggregationDefinition def = queryOperations.createAggregation(aggregation, (Class<?>) null);
+		TypeBasedAggregationOperationContext ctx = (TypeBasedAggregationOperationContext) def
+				.getAggregationOperationContext();
 
-		assertThat(ctx.getAggregationOperationContext()).isInstanceOf(RelaxedTypeBasedAggregationOperationContext.class);
+		assertThat(ReflectionTestUtils.getField(ctx, "lookupPolicy")).isEqualTo(FieldLookupPolicy.relaxed());
 	}
 
 	@Test // GH-3542
 	void createAggregationContextUsesRelaxedOneForTypedAggregationsWhenNoInputTypeProvided() {
 
 		Aggregation aggregation = Aggregation.newAggregation(Person.class, Aggregation.project("name"));
-		AggregationDefinition ctx = queryOperations.createAggregation(aggregation, (Class<?>) null);
+		AggregationDefinition def = queryOperations.createAggregation(aggregation, Person.class);
+		TypeBasedAggregationOperationContext ctx = (TypeBasedAggregationOperationContext) def
+				.getAggregationOperationContext();
 
-		assertThat(ctx.getAggregationOperationContext()).isInstanceOf(RelaxedTypeBasedAggregationOperationContext.class);
+		assertThat(ReflectionTestUtils.getField(ctx, "lookupPolicy")).isEqualTo(FieldLookupPolicy.relaxed());
 	}
 
 	@Test // GH-3542
 	void createAggregationContextUsesRelaxedOneForUntypedAggregationsWhenInputTypeProvided() {
 
 		Aggregation aggregation = Aggregation.newAggregation(Aggregation.project("name"));
-		AggregationDefinition ctx = queryOperations.createAggregation(aggregation, Person.class);
+		AggregationDefinition def = queryOperations.createAggregation(aggregation, Person.class);
+		TypeBasedAggregationOperationContext ctx = (TypeBasedAggregationOperationContext) def
+				.getAggregationOperationContext();
 
-		assertThat(ctx.getAggregationOperationContext()).isInstanceOf(RelaxedTypeBasedAggregationOperationContext.class);
+		assertThat(ReflectionTestUtils.getField(ctx, "lookupPolicy")).isEqualTo(FieldLookupPolicy.relaxed());
 	}
 
 	@Test // GH-3542

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

@@ -15,22 +15,15 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
-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 static org.mockito.Mockito.*;
+import static org.springframework.data.domain.Sort.Direction.*;
+import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
 
 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;
@@ -54,80 +47,6 @@ public class AggregationOperationRendererUnitTests {
 		verify(stage2).toPipelineStages(eq(rootContext));
 	}
 
-	@Test // GH-4443
-	void fieldsExposingAggregationOperationNotExposingFieldsForcesUseOfDefaultContextForNextStage() {
-
-		AggregationOperationContext rootContext = mock(AggregationOperationContext.class);
-		FieldsExposingAggregationOperation stage1 = mock(FieldsExposingAggregationOperation.class);
-		ExposedFields stage1fields = mock(ExposedFields.class);
-		AggregationOperation stage2 = mock(AggregationOperation.class);
-
-		when(stage1.getFields()).thenReturn(stage1fields);
-		when(stage1fields.exposesNoFields()).thenReturn(true);
-
-		AggregationOperationRenderer.toDocument(List.of(stage1, stage2), rootContext);
-
-		verify(stage1).toPipelineStages(eq(rootContext));
-		verify(stage2).toPipelineStages(eq(AggregationOperationRenderer.DEFAULT_CONTEXT));
-	}
-
-	@Test // GH-4443
-	void fieldsExposingAggregationOperationForcesNewContextForNextStage() {
-
-		AggregationOperationContext rootContext = mock(AggregationOperationContext.class);
-		FieldsExposingAggregationOperation stage1 = mock(FieldsExposingAggregationOperation.class);
-		ExposedFields stage1fields = mock(ExposedFields.class);
-		AggregationOperation stage2 = mock(AggregationOperation.class);
-
-		when(stage1.getFields()).thenReturn(stage1fields);
-		when(stage1fields.exposesNoFields()).thenReturn(false);
-
-		ArgumentCaptor<AggregationOperationContext> captor = ArgumentCaptor.forClass(AggregationOperationContext.class);
-
-		AggregationOperationRenderer.toDocument(List.of(stage1, stage2), rootContext);
-
-		verify(stage1).toPipelineStages(eq(rootContext));
-		verify(stage2).toPipelineStages(captor.capture());
-
-		assertThat(captor.getValue()).isInstanceOf(ExposedFieldsAggregationOperationContext.class)
-				.isNotInstanceOf(InheritingExposedFieldsAggregationOperationContext.class);
-	}
-
-	@Test // GH-4443
-	void inheritingFieldsExposingAggregationOperationForcesNewContextForNextStageKeepingReferenceToPreviousContext() {
-
-		AggregationOperationContext rootContext = mock(AggregationOperationContext.class);
-		InheritsFieldsAggregationOperation stage1 = mock(InheritsFieldsAggregationOperation.class);
-		InheritsFieldsAggregationOperation stage2 = mock(InheritsFieldsAggregationOperation.class);
-		InheritsFieldsAggregationOperation stage3 = mock(InheritsFieldsAggregationOperation.class);
-
-		ExposedFields exposedFields = mock(ExposedFields.class);
-		when(exposedFields.exposesNoFields()).thenReturn(false);
-		when(stage1.getFields()).thenReturn(exposedFields);
-		when(stage2.getFields()).thenReturn(exposedFields);
-		when(stage3.getFields()).thenReturn(exposedFields);
-
-		ArgumentCaptor<AggregationOperationContext> captor = ArgumentCaptor.forClass(AggregationOperationContext.class);
-
-		AggregationOperationRenderer.toDocument(List.of(stage1, stage2, stage3), rootContext);
-
-		verify(stage1).toPipelineStages(captor.capture());
-		verify(stage2).toPipelineStages(captor.capture());
-		verify(stage3).toPipelineStages(captor.capture());
-
-		assertThat(captor.getAllValues().get(0)).isEqualTo(rootContext);
-
-		assertThat(captor.getAllValues().get(1))
-				.asInstanceOf(InstanceOfAssertFactories.type(InheritingExposedFieldsAggregationOperationContext.class))
-				.extracting("previousContext").isSameAs(captor.getAllValues().get(0));
-
-		assertThat(captor.getAllValues().get(2))
-				.asInstanceOf(InstanceOfAssertFactories.type(InheritingExposedFieldsAggregationOperationContext.class))
-				.extracting("previousContext").isSameAs(captor.getAllValues().get(1));
-	}
-
-
-
 	record TestRecord(@Id String field1, String field2, LayerOne layerOne) {
 		record LayerOne(List<LayerTwo> layerTwo) {
 		}