DATAMONGO-753 - Add support for nested field references in aggregation operations.

Aggregation pipelines now correctly handle nested field references in aggregation operations. We introduced FieldsExposingAggregationOperation to mark AggregationOperations that change the set of exposed fields available for processing by later AggregationOperations. Extracted context state out of AggregationOperation to ExposedFieldsAggregationContext for better separation of concerns. Modified toDbObject(…) in Aggregation to only replace the aggregation context when the current AggregationOperation is a FieldExposingAggregationOperation.

Original pull request: #74.

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

@@ -227,8 +227,9 @@ public class Aggregation {
 
 			operationDocuments.add(operation.toDBObject(context));
 
-			if (operation instanceof AggregationOperationContext) {
-				context = (AggregationOperationContext) operation;
+			if (operation instanceof FieldsExposingAggregationOperation) {
+				FieldsExposingAggregationOperation exposedFieldsOperation = (FieldsExposingAggregationOperation) operation;
+				context = new ExposedFieldsAggregationOperationContext(exposedFieldsOperation.getFields());
 			}
 		}
 

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

@@ -17,17 +17,32 @@ package org.springframework.data.mongodb.core.aggregation;
 
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
+import org.springframework.util.Assert;
 
 import com.mongodb.DBObject;
 
 /**
- * Support class to implement {@link AggregationOperation}s that will become an {@link AggregationOperationContext} as
- * well defining {@link ExposedFields}.
+ * {@link AggregationOperationContext} that combines the available field references from a given
+ * {@code AggregationOperationContext} and an {@link FieldsExposingAggregationOperation}.
  * 
+ * @author Thomas Darimont
  * @author Oliver Gierke
- * @since 1.3
+ * @since 1.4
  */
-public abstract class ExposedFieldsAggregationOperationContext implements AggregationOperationContext {
+class ExposedFieldsAggregationOperationContext implements AggregationOperationContext {
+
+	private final ExposedFields exposedFields;
+
+	/**
+	 * Creates a new {@link ExposedFieldsAggregationOperationContext} from the given {@link ExposedFields}.
+	 * 
+	 * @param exposedFields must not be {@literal null}.
+	 */
+	public ExposedFieldsAggregationOperationContext(ExposedFields exposedFields) {
+
+		Assert.notNull(exposedFields, "ExposedFields must not be null!");
+		this.exposedFields = exposedFields;
+	}
 
 	/* 
 	 * (non-Javadoc)
@@ -54,7 +69,7 @@ public abstract class ExposedFieldsAggregationOperationContext implements Aggreg
 	@Override
 	public FieldReference getReference(String name) {
 
-		ExposedField field = getFields().getField(name);
+		ExposedField field = exposedFields.getField(name);
 
 		if (field != null) {
 			return new FieldReference(field);
@@ -62,6 +77,4 @@ public abstract class ExposedFieldsAggregationOperationContext implements Aggreg
 
 		throw new IllegalArgumentException(String.format("Invalid reference '%s'!", name));
 	}
-
-	protected abstract ExposedFields getFields();
 }

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2013 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
+ *
+ *      http://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;
+
+/**
+ * {@link AggregationOperation} that exposes new {@link ExposedFields} that can be used for later aggregation pipeline
+ * {@code AggregationOperation}s.
+ * 
+ * @author Thomas Darimont
+ */
+public interface FieldsExposingAggregationOperation extends AggregationOperation {
+
+	/**
+	 * Returns the fields exposed by the {@link AggregationOperation}.
+	 * 
+	 * @return will never be {@literal null}.
+	 */
+	ExposedFields getFields();
+}

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

@@ -38,7 +38,7 @@ import com.mongodb.DBObject;
  * @author Oliver Gierke
  * @since 1.3
  */
-public class GroupOperation extends ExposedFieldsAggregationOperationContext implements AggregationOperation {
+public class GroupOperation implements FieldsExposingAggregationOperation {
 
 	private final ExposedFields nonSynthecticFields;
 	private final List<Operation> operations;

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

@@ -41,7 +41,7 @@ import com.mongodb.DBObject;
  * @author Oliver Gierke
  * @since 1.3
  */
-public class ProjectionOperation extends ExposedFieldsAggregationOperationContext implements AggregationOperation {
+public class ProjectionOperation implements FieldsExposingAggregationOperation {
 
 	private static final List<Projection> NONE = Collections.emptyList();
 
@@ -149,10 +149,10 @@ public class ProjectionOperation extends ExposedFieldsAggregationOperationContex
 
 	/* 
 	 * (non-Javadoc)
-	 * @see org.springframework.data.mongodb.core.aggregation.ExposedFieldsAggregationOperationContext#getFields()
+	 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperationContextSupport#getFields()
 	 */
 	@Override
-	protected ExposedFields getFields() {
+	public ExposedFields getFields() {
 
 		ExposedFields fields = null;
 

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

@@ -71,9 +71,9 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio
 		return mapper.getMappedObject(dbObject, mappingContext.getPersistentEntity(type));
 	}
 
-	/* 
+	/*
 	 * (non-Javadoc)
-	 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperationContext#getReference(org.springframework.data.mongodb.core.aggregation.ExposedFields.AvailableField)
+	 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperationContext#getReference(org.springframework.data.mongodb.core.aggregation.Field)
 	 */
 	@Override
 	public FieldReference getReference(Field field) {

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

@@ -29,7 +29,7 @@ import com.mongodb.DBObject;
  * @author Oliver Gierke
  * @since 1.3
  */
-public class UnwindOperation extends ExposedFieldsAggregationOperationContext implements AggregationOperation {
+public class UnwindOperation implements AggregationOperation {
 
 	private final ExposedField field;
 
@@ -44,15 +44,6 @@ public class UnwindOperation extends ExposedFieldsAggregationOperationContext im
 		this.field = new ExposedField(field, true);
 	}
 
-	/* 
-	 * (non-Javadoc)
-	 * @see org.springframework.data.mongodb.core.aggregation.ExposedFieldsAggregationOperationContext#getFields()
-	 */
-	@Override
-	protected ExposedFields getFields() {
-		return ExposedFields.from(field);
-	}
-
 	/* 
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperation#toDBObject(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext)

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

@@ -15,7 +15,7 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
-import static org.hamcrest.CoreMatchers.*;
+import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 import static org.springframework.data.domain.Sort.Direction.*;
 import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
@@ -84,6 +84,7 @@ public class AggregationTests {
 		mongoTemplate.dropCollection(INPUT_COLLECTION);
 		mongoTemplate.dropCollection(Product.class);
 		mongoTemplate.dropCollection(UserWithLikes.class);
+		mongoTemplate.dropCollection(DATAMONGO753.class);
 	}
 
 	/**
@@ -452,7 +453,61 @@ public class AggregationTests {
 		assertThat((Double) resultList.get(0).get("netPriceMul2"), is(netPrice * 2));
 		assertThat((Double) resultList.get(0).get("netPriceDiv119"), is(netPrice / 1.19));
 		assertThat((Integer) resultList.get(0).get("spaceUnitsMod2"), is(spaceUnits % 2));
+	}
 
+	/**
+	 * @see DATAMONGO-753
+	 * @see http 
+	 *      ://stackoverflow.com/questions/18653574/spring-data-mongodb-aggregation-framework-invalid-reference-in-group
+	 *      -operati
+	 */
+	@Test
+	public void allowsNestedFieldReferencesAsGroupIdsInGroupExpressions() {
+
+		mongoTemplate.insert(new DATAMONGO753().withPDs(new PD("A", 1), new PD("B", 1), new PD("C", 1)));
+		mongoTemplate.insert(new DATAMONGO753().withPDs(new PD("B", 1), new PD("B", 1), new PD("C", 1)));
+
+		TypedAggregation<DATAMONGO753> agg = newAggregation(DATAMONGO753.class, //
+				unwind("pd"), //
+				group("pd.pDch") // the nested field expression
+						.sum("pd.up").as("uplift"), //
+				project("_id", "uplift"));
+
+		AggregationResults<DBObject> result = mongoTemplate.aggregate(agg, DBObject.class);
+		List<DBObject> stats = result.getMappedResults();
+
+		assertThat(stats.size(), is(3));
+		assertThat(stats.get(0).get("_id").toString(), is("C"));
+		assertThat((Integer) stats.get(0).get("uplift"), is(2));
+		assertThat(stats.get(1).get("_id").toString(), is("B"));
+		assertThat((Integer) stats.get(1).get("uplift"), is(3));
+		assertThat(stats.get(2).get("_id").toString(), is("A"));
+		assertThat((Integer) stats.get(2).get("uplift"), is(1));
+	}
+
+	/**
+	 * @see DATAMONGO-753
+	 * @see http 
+	 *      ://stackoverflow.com/questions/18653574/spring-data-mongodb-aggregation-framework-invalid-reference-in-group
+	 *      -operati
+	 */
+	@Test
+	public void aliasesNestedFieldInProjectionImmediately() {
+
+		mongoTemplate.insert(new DATAMONGO753().withPDs(new PD("A", 1), new PD("B", 1), new PD("C", 1)));
+		mongoTemplate.insert(new DATAMONGO753().withPDs(new PD("B", 1), new PD("B", 1), new PD("C", 1)));
+
+		TypedAggregation<DATAMONGO753> agg = newAggregation(DATAMONGO753.class, //
+				unwind("pd"), //
+				project().and("pd.up").as("up"));
+
+		AggregationResults<DBObject> results = mongoTemplate.aggregate(agg, DBObject.class);
+		List<DBObject> mappedResults = results.getMappedResults();
+
+		assertThat(mappedResults, hasSize(6));
+		for (DBObject element : mappedResults) {
+			assertThat(element.get("up"), is((Object) 1));
+		}
 	}
 
 	private void assertLikeStats(LikeStats like, String id, long count) {
@@ -502,4 +557,22 @@ public class AggregationTests {
 		assertThat(tagCount.getN(), is(n));
 	}
 
+	static class DATAMONGO753 {
+		PD[] pd;
+
+		DATAMONGO753 withPDs(PD... pds) {
+			this.pd = pds;
+			return this;
+		}
+	}
+
+	static class PD {
+		String pDch;
+		@org.springframework.data.mongodb.core.mapping.Field("alias") int up;
+
+		public PD(String pDch, int up) {
+			this.pDch = pDch;
+			this.up = up;
+		}
+	}
 }

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

@@ -15,32 +15,83 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
+import static org.springframework.data.mongodb.core.query.Criteria.*;
+
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
 /**
  * Unit tests for {@link Aggregation}.
  * 
  * @author Oliver Gierke
+ * @author Thomas Darimont
  */
 public class AggregationUnitTests {
 
+	public @Rule ExpectedException exception = ExpectedException.none();
+
 	@Test(expected = IllegalArgumentException.class)
 	public void rejectsNullAggregationOperation() {
-		Aggregation.newAggregation((AggregationOperation[]) null);
+		newAggregation((AggregationOperation[]) null);
 	}
 
 	@Test(expected = IllegalArgumentException.class)
 	public void rejectsNullTypedAggregationOperation() {
-		Aggregation.newAggregation(String.class, (AggregationOperation[]) null);
+		newAggregation(String.class, (AggregationOperation[]) null);
 	}
 
 	@Test(expected = IllegalArgumentException.class)
 	public void rejectsNoAggregationOperation() {
-		Aggregation.newAggregation(new AggregationOperation[0]);
+		newAggregation(new AggregationOperation[0]);
 	}
 
 	@Test(expected = IllegalArgumentException.class)
 	public void rejectsNoTypedAggregationOperation() {
-		Aggregation.newAggregation(String.class, new AggregationOperation[0]);
+		newAggregation(String.class, new AggregationOperation[0]);
+	}
+
+	/**
+	 * @see DATAMONGO-753
+	 */
+	@Test
+	public void checkForCorrectFieldScopeTransfer() {
+
+		exception.expect(IllegalArgumentException.class);
+		exception.expectMessage("Invalid reference");
+		exception.expectMessage("'b'");
+
+		newAggregation( //
+				project("a", "b"), //
+				group("a").count().as("cnt"), // a was introduced to the context by the project operation
+				project("cnt", "b") // b was removed from the context by the group operation
+		).toDbObject("foo", Aggregation.DEFAULT_CONTEXT); // -> triggers IllegalArgumentException
+	}
+
+	/**
+	 * @see DATAMONGO-753
+	 */
+	@Test
+	public void unwindOperationShouldNotChangeAvailableFields() {
+
+		newAggregation( //
+				project("a", "b"), //
+				unwind("a"), //
+				project("a", "b") // b should still be available
+		).toDbObject("foo", Aggregation.DEFAULT_CONTEXT);
+	}
+
+	/**
+	 * @see DATAMONGO-753
+	 */
+	@Test
+	public void matchOperationShouldNotChangeAvailableFields() {
+
+		newAggregation( //
+				project("a", "b"), //
+				match(where("a").gte(1)), //
+				project("a", "b") // b should still be available
+		).toDbObject("foo", Aggregation.DEFAULT_CONTEXT);
 	}
 }