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

Author: mp911de

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 @@ AggregationOperationContext createAggregationContext(Aggregation aggregation, @N
 			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 @@ Document createCommand(String collection, Aggregation aggregation, AggregationOp
 		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 @@
  *
  * @author Oliver Gierke
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @since 1.3
  */
 public interface AggregationOperationContext extends CodecRegistryProvider {
@@ -107,14 +108,46 @@ default Fields getFields(Class<?> type) {
 				.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 @@ default AggregationOperationContext continueOnMissingFieldReference() {
 	default CodecRegistry getCodecRegistry() {
 		return MongoClientSettings.getDefaultCodecRegistry();
 	}
+
 }

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

@@ -50,7 +50,6 @@ 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) {
 
@@ -61,10 +60,10 @@ static List<Document> toDocument(List<AggregationOperation> operations, Aggregat
 				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 Document toDocument(final AggregationOperationContext context) {
 		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 @@ protected DocumentEnhancingOperation(Map<Object, Object> source) {
 	@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 @@
 
 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,24 +38,26 @@ 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
 	 * {@link AggregationOperationContext} to perform a mapping to mongo types if necessary.
 	 *
 	 * @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 @@ public Fields getFields(Class<?> type) {
 	 * @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 @@ protected FieldReference getReference(@Nullable Field field, String name) {
 			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 @@ public CodecRegistry getCodecRegistry() {
 	}
 
 	@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 @@
 
 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 @@ public Document getMappedObject(Document document) {
 	}
 
 	@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) {