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 @@ 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.2.7
+ */
+ 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.2.7
+ */
+ 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.2.7, {@link FieldLookupPolicy} should be specified explicitly when creating the
+ * AggregationOperationContext.
  */
+ @Deprecated(since = "4.2.7", 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.2.7", 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.2.7
+ */
+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) {