Fix mapping custom field names in downstream stages in TypedAggregation pipelines

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-mongodb-parent</artifactId>
-	<version>4.2.0-SNAPSHOT</version>
+	<version>4.2.x-4443-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.2.0-SNAPSHOT</version>
+		<version>4.2.x-4443-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.2.0-SNAPSHOT</version>
+		<version>4.2.x-4443-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.2.0-SNAPSHOT</version>
+		<version>4.2.x-4443-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

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

@@ -63,7 +63,7 @@ static List<Document> toDocument(List<AggregationOperation> operations, Aggregat
 					contextToUse = new InheritingExposedFieldsAggregationOperationContext(fields, contextToUse);
 				} else {
 					contextToUse = fields.exposesNoFields() ? DEFAULT_CONTEXT
-							: new ExposedFieldsAggregationOperationContext(exposedFieldsOperation.getFields(), contextToUse);
+							: new ExposedFieldsAggregationOperationContext(fields, contextToUse);
 				}
 			}
 		}

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

@@ -15,13 +15,15 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import org.bson.Document;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
 
 /**
  * {@link ExposedFieldsAggregationOperationContext} that inherits fields from its parent
  * {@link AggregationOperationContext}.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  * @since 1.9
  */
 class InheritingExposedFieldsAggregationOperationContext extends ExposedFieldsAggregationOperationContext {
@@ -43,6 +45,11 @@ public InheritingExposedFieldsAggregationOperationContext(ExposedFields exposedF
 		this.previousContext = previousContext;
 	}
 
+	@Override
+	public Document getMappedObject(Document document) {
+		return previousContext.getMappedObject(document);
+	}
+
 	@Override
 	protected FieldReference resolveExposedField(Field field, String name) {
 

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

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2023 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;
+
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import java.util.List;
+
+import org.assertj.core.api.InstanceOfAssertFactories;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.springframework.data.mongodb.core.aggregation.FieldsExposingAggregationOperation.InheritsFieldsAggregationOperation;
+
+/**
+ * @author Christoph Strobl
+ */
+public class AggregationOperationRendererUnitTests {
+
+	@Test // GH-4443
+	void nonFieldsExposingAggregationOperationContinuesWithSameContextForNextStage() {
+
+		AggregationOperationContext rootContext = mock(AggregationOperationContext.class);
+		AggregationOperation stage1 = mock(AggregationOperation.class);
+		AggregationOperation stage2 = mock(AggregationOperation.class);
+
+		AggregationOperationRenderer.toDocument(List.of(stage1, stage2), rootContext);
+
+		verify(stage1).toPipelineStages(eq(rootContext));
+		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));
+	}
+
+}

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

@@ -91,6 +91,7 @@
  * @author Sergey Shcherbakov
  * @author Minsu Kim
  * @author Sangyong Choi
+ * @author Julia Lee
  */
 @ExtendWith(MongoTemplateExtension.class)
 public class AggregationTests {
@@ -119,7 +120,7 @@ private void cleanDb() {
 
 		mongoTemplate.flush(Product.class, UserWithLikes.class, DATAMONGO753.class, Data.class, DATAMONGO788.class,
 				User.class, Person.class, Reservation.class, Venue.class, MeterData.class, LineItem.class, InventoryItem.class,
-				Sales.class, Sales2.class, Employee.class, Art.class, Venue.class);
+				Sales.class, Sales2.class, Employee.class, Art.class, Venue.class, Item.class);
 
 		mongoTemplate.dropCollection(INPUT_COLLECTION);
 		mongoTemplate.dropCollection("personQueryTemp");
@@ -1992,6 +1993,42 @@ void considersMongoIdWithinTypedCollections() {
 		assertThat(aggregate.getMappedResults()).contains(widget);
 	}
 
+	@Test // GH-4443
+	void shouldHonorFieldAliasesForFieldReferencesUsingFieldExposingOperation() {
+
+		Item item1 = Item.builder().itemId("1").tags(Arrays.asList("a", "b")).build();
+		Item item2 = Item.builder().itemId("1").tags(Arrays.asList("a", "c")).build();
+		mongoTemplate.insert(Arrays.asList(item1, item2), Item.class);
+
+		TypedAggregation<Item> aggregation = newAggregation(Item.class,
+				match(where("itemId").is("1")),
+				unwind("tags"),
+				match(where("itemId").is("1").and("tags").is("c")));
+		AggregationResults<Document> results = mongoTemplate.aggregate(aggregation, Document.class);
+		List<Document> mappedResults = results.getMappedResults();
+		assertThat(mappedResults).hasSize(1);
+		assertThat(mappedResults.get(0)).containsEntry("item_id", "1");
+	}
+
+	@Test // GH-4443
+	void projectShouldResetContextToAvoidMappingFieldsAgainstANoLongerExistingTarget() {
+
+		Item item1 = Item.builder().itemId("1").tags(Arrays.asList("a", "b")).build();
+		Item item2 = Item.builder().itemId("1").tags(Arrays.asList("a", "c")).build();
+		mongoTemplate.insert(Arrays.asList(item1, item2), Item.class);
+
+		TypedAggregation<Item> aggregation = newAggregation(Item.class,
+				match(where("itemId").is("1")),
+				unwind("tags"),
+				project().and("itemId").as("itemId").and("tags").as("tags"),
+				match(where("itemId").is("1").and("tags").is("c")));
+
+		AggregationResults<Document> results = mongoTemplate.aggregate(aggregation, Document.class);
+		List<Document> mappedResults = results.getMappedResults();
+		assertThat(mappedResults).hasSize(1);
+		assertThat(mappedResults.get(0)).containsEntry("itemId", "1");
+	}
+
 	private void createUsersWithReferencedPersons() {
 
 		mongoTemplate.dropCollection(User.class);
@@ -2314,19 +2351,21 @@ public String toString() {
 		}
 	}
 
-	// DATAMONGO-1491
+	// DATAMONGO-1491, GH-4443
 	static class Item {
 
 		@org.springframework.data.mongodb.core.mapping.Field("item_id") //
 		String itemId;
 		Integer quantity;
 		Long price;
+		List<String> tags = new ArrayList<>();
 
-		Item(String itemId, Integer quantity, Long price) {
+		Item(String itemId, Integer quantity, Long price, List<String> tags) {
 
 			this.itemId = itemId;
 			this.quantity = quantity;
 			this.price = price;
+			this.tags = tags;
 		}
 
 		public static ItemBuilder builder() {
@@ -2385,6 +2424,7 @@ public static class ItemBuilder {
 			private String itemId;
 			private Integer quantity;
 			private Long price;
+			private List<String> tags;
 
 			ItemBuilder() {}
 
@@ -2403,8 +2443,13 @@ public ItemBuilder price(Long price) {
 				return this;
 			}
 
+			public ItemBuilder tags(List<String> tags) {
+				this.tags = tags;
+				return this;
+			}
+
 			public Item build() {
-				return new Item(itemId, quantity, price);
+				return new Item(itemId, quantity, price, tags);
 			}
 
 			public String toString() {

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

@@ -49,6 +49,7 @@
  * @author Thomas Darimont
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Julia Lee
  */
 public class AggregationUnitTests {
 
@@ -612,7 +613,7 @@ void shouldNotConvertIncludeExcludeValuesForProjectOperation() {
 				WithRetypedIdField.class, mappingContext,
 				new QueryMapper(new MappingMongoConverter(NoOpDbRefResolver.INSTANCE, mappingContext)));
 		Document document = project(WithRetypedIdField.class).toDocument(context);
-		assertThat(document).isEqualTo(new Document("$project", new Document("_id", 1).append("renamed-field", 1)));
+		assertThat(document).isEqualTo(new Document("$project", new Document("_id", 1).append("renamed-field", 1).append("entries", 1)));
 	}
 
 	@Test // GH-4038
@@ -653,6 +654,22 @@ void inheritedFieldsExposingContextShouldNotFailOnUnknownFieldReferenceForRelaxe
 		assertThat(documents.get(2)).isEqualTo("{ $sort : { 'serial_number' : -1, 'label_name' : -1 } }");
 	}
 
+	@Test // GH-4443
+	void fieldsExposingContextShouldUseCustomFieldNameFromRelaxedRootContext() {
+
+		MongoMappingContext mappingContext = new MongoMappingContext();
+		RelaxedTypeBasedAggregationOperationContext context = new RelaxedTypeBasedAggregationOperationContext(
+				WithRetypedIdField.class, mappingContext,
+				new QueryMapper(new MappingMongoConverter(NoOpDbRefResolver.INSTANCE, mappingContext)));
+
+		TypedAggregation<WithRetypedIdField> agg = newAggregation(WithRetypedIdField.class,
+				unwind("entries"), match(where("foo").is("value 2")));
+		List<Document> pipeline = agg.toPipeline(context);
+
+		Document fields = getAsDocument(pipeline.get(1), "$match");
+		assertThat(fields.get("renamed-field")).isEqualTo("value 2");
+	}
+
 	private Document extractPipelineElement(Document agg, int index, String operation) {
 
 		List<Document> pipeline = (List<Document>) agg.get("pipeline");
@@ -672,5 +689,7 @@ public class WithRetypedIdField {
 
 		@org.springframework.data.mongodb.core.mapping.Field("renamed-field") private String foo;
 
+		private List<String> entries = new ArrayList<>();
+
 	}
 }