Propagate relaxed field lookup across nested aggregation contexts

pom.xml

@@ -5,7 +5,7 @@
 
  <groupId>org.springframework.data</groupId>
  <artifactId>spring-data-mongodb-parent</artifactId>
- <version>4.4.0-SNAPSHOT</version>
+ <version>4.4.x-GH-4714-SNAPSHOT</version>
  <packaging>pom</packaging>
 
  <name>Spring Data MongoDB</name>

spring-data-mongodb-benchmarks/pom.xml

@@ -7,7 +7,7 @@
  <parent>
  <groupId>org.springframework.data</groupId>
  <artifactId>spring-data-mongodb-parent</artifactId>
- <version>4.4.0-SNAPSHOT</version>
+ <version>4.4.x-GH-4714-SNAPSHOT</version>
  <relativePath>../pom.xml</relativePath>
  </parent>
 

spring-data-mongodb-distribution/pom.xml

@@ -15,7 +15,7 @@
  <parent>
  <groupId>org.springframework.data</groupId>
  <artifactId>spring-data-mongodb-parent</artifactId>
- <version>4.4.0-SNAPSHOT</version>
+ <version>4.4.x-GH-4714-SNAPSHOT</version>
  <relativePath>../pom.xml</relativePath>
  </parent>
 

spring-data-mongodb/pom.xml

@@ -13,7 +13,7 @@
  <parent>
  <groupId>org.springframework.data</groupId>
  <artifactId>spring-data-mongodb-parent</artifactId>
- <version>4.4.0-SNAPSHOT</version>
+ <version>4.4.x-GH-4714-SNAPSHOT</version>
  <relativePath>../pom.xml</relativePath>
  </parent>
 

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

@@ -60,12 +60,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 = contextToUse.inheritAndExpose(fields);
  } else {
  contextToUse = fields.exposesNoFields() ? DEFAULT_CONTEXT
- : new ExposedFieldsAggregationOperationContext(fields, contextToUse);
+ : contextToUse.expose(fields);
  }
  }
+
  }
 
  return operationDocuments;

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);
+ 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);
+ 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

@@ -15,11 +15,14 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import java.util.function.BiFunction;
+
 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;
+import org.springframework.lang.NonNull;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
@@ -37,22 +40,33 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 
  private final ExposedFields exposedFields;
  private final AggregationOperationContext rootContext;
+ private final FieldLookupPolicy lookupPolicy;
+ private final ContextualLookupSupport contextualLookup;
 
  /**
  * 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) {
+ 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.lookupPolicy = lookupPolicy;
+ this.contextualLookup = ContextualLookupSupport.create(lookupPolicy, this::resolveExposedField, (field, name) -> {
+ if (field != null) {
+ return new DirectFieldReference(new ExposedField(field, true));
+ }
+ return new DirectFieldReference(new ExposedField(name, true));
+ });
  }
 
  @Override
@@ -87,25 +101,11 @@ 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");
 
- FieldReference exposedField = resolveExposedField(field, name);
- if (exposedField != null) {
- return exposedField;
- }
-
- if (rootContext instanceof RelaxedTypeBasedAggregationOperationContext) {
-
- if (field != null) {
- return new DirectFieldReference(new ExposedField(field, true));
- }
-
- return new DirectFieldReference(new ExposedField(name, true));
- }
-
- throw new IllegalArgumentException(String.format("Invalid reference '%s'", name));
+ return contextualLookup.get(field, name);
  }
 
  /**
@@ -156,4 +156,90 @@ AggregationOperationContext getRootContext() {
  public CodecRegistry getCodecRegistry() {
  return getRootContext().getCodecRegistry();
  }
+
+ @Override
+ public AggregationOperationContext continueOnMissingFieldReference() {
+ if (!lookupPolicy.isStrict()) {
+ return this;
+ }
+ return new ExposedFieldsAggregationOperationContext(exposedFields, rootContext, FieldLookupPolicy.lenient());
+ }
+
+ @Override
+ public AggregationOperationContext expose(ExposedFields fields) {
+ return new ExposedFieldsAggregationOperationContext(fields, this, lookupPolicy);
+ }
+
+ @Override
+ public AggregationOperationContext inheritAndExpose(ExposedFields fields) {
+ return new InheritingExposedFieldsAggregationOperationContext(fields, this, lookupPolicy);
+ }
+
+ static class ContextualLookupSupport {
+
+ private final BiFunction<Field, String, FieldReference> resolver;
+
+ ContextualLookupSupport(BiFunction<Field, String, FieldReference> resolver) {
+ this.resolver = resolver;
+ }
+
+ public static ContextualLookupSupport create(FieldLookupPolicy lookupPolicy,
+ BiFunction<Field, String, FieldReference> resolver, BiFunction<Field, String, FieldReference> fallback) {
+
+ if (lookupPolicy.isStrict()) {
+ return new StrictContextualLookup(resolver);
+ }
+
+ return new FallbackContextualLookup(resolver, fallback);
+
+ }
+
+ public FieldReference get(@Nullable Field field, String name) {
+ return resolver.apply(field, name);
+ }
+ }
+
+ static class StrictContextualLookup extends ContextualLookupSupport {
+
+ StrictContextualLookup(BiFunction<Field, String, FieldReference> resolver) {
+ super(resolver);
+ }
+
+ @Override
+ @NonNull
+ public FieldReference get(Field field, String name) {
+
+ FieldReference lookup = super.get(field, name);
+
+ if (lookup != null) {
+ return lookup;
+ }
+
+ throw new IllegalArgumentException(String.format("Invalid reference '%s'", name));
+ }
+ }
+
+ static class FallbackContextualLookup extends ContextualLookupSupport {
+
+ private final BiFunction<Field, String, FieldReference> fallback;
+
+ FallbackContextualLookup(BiFunction<Field, String, FieldReference> resolver,
+ BiFunction<Field, String, FieldReference> fallback) {
+ super(resolver);
+ this.fallback = fallback;
+ }
+
+ @Override
+ @NonNull
+ public FieldReference get(@Nullable Field field, String name) {
+
+ FieldReference lookup = super.get(field, name);
+
+ if (lookup != null) {
+ return lookup;
+ }
+
+ return fallback.apply(field, name);
+ }
+ }
 }