DATAMONGO-954 - Add support for system variables in aggregation operations.

We now support referring to system variables like for instance $$ROOT or $$CURRENT from within aggregation framework pipeline projection and group expressions.

Original pull request: spring-projects#190.

Author: Thomas Darimont

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

@@ -44,11 +44,22 @@
  */
 public class Aggregation {
 
+	/**
+	 * References the root document, i.e. the top-level document, currently being processed in the aggregation pipeline
+	 * stage.
+	 */
+	public static final String ROOT = SystemVariable.ROOT.toString();
+
+	/**
+	 * References the start of the field path being processed in the aggregation pipeline stage. Unless documented
+	 * otherwise, all stages start with CURRENT the same as ROOT.
+	 */
+	public static final String CURRENT = SystemVariable.CURRENT.toString();
+
 	public static final AggregationOperationContext DEFAULT_CONTEXT = new NoOpAggregationOperationContext();
 	public static final AggregationOptions DEFAULT_OPTIONS = newAggregationOptions().build();
 
 	protected final List<AggregationOperation> operations;
-
 	private final AggregationOptions options;
 
 	/**
@@ -363,4 +374,51 @@ public FieldReference getReference(String name) {
 			return new FieldReference(new ExposedField(new AggregationField(name), true));
 		}
 	}
+
+	/**
+	 * Describes the system variables available in MongoDB aggregation framework pipeline expressions.
+	 * 
+	 * @author Thomas Darimont
+	 * @see http://docs.mongodb.org/manual/reference/aggregation-variables
+	 */
+	enum SystemVariable {
+
+		ROOT, CURRENT;
+
+		private static final String PREFIX = "$$";
+
+		/**
+		 * Return {@literal true} if the given {@code fieldRef} denotes a well-known system variable, {@literal false}
+		 * otherwise.
+		 * 
+		 * @param fieldRef may be {@literal null}.
+		 * @return
+		 */
+		public static boolean isReferingToSystemVariable(String fieldRef) {
+
+			if (fieldRef == null || !fieldRef.startsWith(PREFIX) || fieldRef.length() <= 2) {
+				return false;
+			}
+
+			int indexOfFirstDot = fieldRef.indexOf('.');
+			String candidate = fieldRef.substring(2, indexOfFirstDot == -1 ? fieldRef.length() : indexOfFirstDot);
+
+			for (SystemVariable value : values()) {
+				if (value.name().equals(candidate)) {
+					return true;
+				}
+			}
+
+			return false;
+		}
+
+		/* 
+		 * (non-Javadoc)
+		 * @see java.lang.Enum#toString()
+		 */
+		@Override
+		public String toString() {
+			return PREFIX.concat(name());
+		}
+	}
 }

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

@@ -30,6 +30,7 @@
  * Value object to capture a list of {@link Field} instances.
  * 
  * @author Oliver Gierke
+ * @author Thomas Darimont
  * @since 1.3
  */
 public final class Fields implements Iterable<Field> {
@@ -186,7 +187,7 @@ static class AggregationField implements Field {
 		private final String target;
 
 		/**
-		 * Creates an aggregation fieldwith the given name. As no target is set explicitly, the name will be used as target
+		 * Creates an aggregation field with the given name. As no target is set explicitly, the name will be used as target
 		 * as well.
 		 * 
 		 * @param key
@@ -217,6 +218,10 @@ private static final String cleanUp(String source) {
 				return source;
 			}
 
+			if (Aggregation.SystemVariable.isReferingToSystemVariable(source)) {
+				return source;
+			}
+
 			int dollarIndex = source.lastIndexOf('$');
 			return dollarIndex == -1 ? source : source.substring(dollarIndex + 1);
 		}

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

@@ -364,7 +364,16 @@ public DBObject toDBObject(AggregationOperationContext context) {
 		}
 
 		public Object getValue(AggregationOperationContext context) {
-			return reference == null ? value : context.getReference(reference).toString();
+
+			if (reference == null) {
+				return value;
+			}
+
+			if (Aggregation.SystemVariable.isReferingToSystemVariable(reference)) {
+				return reference;
+			}
+
+			return context.getReference(reference).toString();
 		}
 
 		@Override

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

@@ -627,6 +627,10 @@ private Object renderFieldValue(AggregationOperationContext context) {
 				// implicit reference or explicit include?
 				if (value == null || Boolean.TRUE.equals(value)) {
 
+					if (Aggregation.SystemVariable.isReferingToSystemVariable(field.getTarget())) {
+						return field.getTarget();
+					}
+
 					// check whether referenced field exists in the context
 					return context.getReference(field).getReferenceValue();
 

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

@@ -44,11 +44,13 @@
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.dao.DataAccessException;
 import org.springframework.data.annotation.Id;
+import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mapping.model.MappingException;
 import org.springframework.data.mongodb.core.CollectionCallback;
 import org.springframework.data.mongodb.core.MongoTemplate;
 import org.springframework.data.mongodb.core.aggregation.AggregationTests.CarDescriptor.Entry;
 import org.springframework.data.mongodb.core.query.Query;
+import org.springframework.data.mongodb.repository.Person;
 import org.springframework.data.util.Version;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -113,6 +115,8 @@ private void cleanDb() {
 		mongoTemplate.dropCollection(Data.class);
 		mongoTemplate.dropCollection(DATAMONGO788.class);
 		mongoTemplate.dropCollection(User.class);
+		mongoTemplate.dropCollection(Person.class);
+		mongoTemplate.dropCollection(Reservation.class);
 	}
 
 	/**
@@ -903,6 +907,55 @@ public void returnFiveMostCommonLikesShouldReturnStageExecutionInformationWithEx
 		assertThat(rawResult.containsField("stages"), is(true));
 	}
 
+	/**
+	 * @see DATAMONGO-954
+	 */
+	@Test
+	public void shouldSupportReturningCurrentAggregationRoot() {
+
+		mongoTemplate.save(new Person("p1_first", "p1_last", 25));
+		mongoTemplate.save(new Person("p2_first", "p2_last", 32));
+		mongoTemplate.save(new Person("p3_first", "p3_last", 25));
+		mongoTemplate.save(new Person("p4_first", "p4_last", 15));
+
+		List<DBObject> personsWithAge25 = mongoTemplate.find(Query.query(where("age").is(25)), DBObject.class,
+				mongoTemplate.getCollectionName(Person.class));
+
+		Aggregation agg = newAggregation(group("age").push(Aggregation.ROOT).as("users"));
+		AggregationResults<DBObject> result = mongoTemplate.aggregate(agg, Person.class, DBObject.class);
+
+		assertThat(result.getMappedResults(), hasSize(3));
+		DBObject o = (DBObject) result.getMappedResults().get(2);
+
+		assertThat(o.get("_id"), is((Object) 25));
+		assertThat((List<?>) o.get("users"), hasSize(2));
+		assertThat((List<?>) o.get("users"), is(contains(personsWithAge25.toArray())));
+	}
+
+	/**
+	 * @see DATAMONGO-954
+	 * @see http
+	 *      ://stackoverflow.com/questions/24185987/using-root-inside-spring-data-mongodb-for-retrieving-whole-document
+	 */
+	@Test
+	public void shouldSupportReturningCurrentAggregationRootInReference() {
+
+		mongoTemplate.save(new Reservation("0123", "42", 100));
+		mongoTemplate.save(new Reservation("0360", "43", 200));
+		mongoTemplate.save(new Reservation("0360", "44", 300));
+
+		Aggregation agg = newAggregation( //
+				match(where("hotelCode").is("0360")), //
+				sort(Direction.DESC, "confirmationNumber", "timestamp"), //
+				group("confirmationNumber") //
+						.first("timestamp").as("timestamp") //
+						.first(Aggregation.ROOT).as("reservationImage") //
+		);
+		AggregationResults<DBObject> result = mongoTemplate.aggregate(agg, Reservation.class, DBObject.class);
+
+		assertThat(result.getMappedResults(), hasSize(2));
+	}
+
 	private void assertLikeStats(LikeStats like, String id, long count) {
 
 		assertThat(like, is(notNullValue()));
@@ -1067,4 +1120,19 @@ public Entry(String make, String model, int year) {
 			}
 		}
 	}
+
+	static class Reservation {
+
+		String hotelCode;
+		String confirmationNumber;
+		int timestamp;
+
+		public Reservation() {}
+
+		public Reservation(String hotelCode, String confirmationNumber, int timestamp) {
+			this.hotelCode = hotelCode;
+			this.confirmationNumber = confirmationNumber;
+			this.timestamp = timestamp;
+		}
+	}
 }

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

@@ -27,6 +27,7 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
+import org.springframework.data.domain.Sort.Direction;
 
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
@@ -256,6 +257,32 @@ public void shouldRenderAggregationWithCustomOptionsCorrectly() {
 		));
 	}
 
+	/**
+	 * @see DATAMONGO-954
+	 */
+	@Test
+	public void shouldSupportReferencingSystemVariables() {
+
+		DBObject agg = newAggregation( //
+				project("someKey") //
+						.and("a").as("a1") //
+						.and(Aggregation.CURRENT + ".a").as("a2") //
+				, sort(Direction.DESC, "a") //
+				, group("someKey").first(Aggregation.ROOT).as("doc") //
+		).toDbObject("foo", Aggregation.DEFAULT_CONTEXT);
+
+		DBObject projection0 = extractPipelineElement(agg, 0, "$project");
+		assertThat(projection0, is((DBObject) new BasicDBObject("someKey", 1).append("a1", "$a")
+				.append("a2", "$$CURRENT.a")));
+
+		DBObject sort = extractPipelineElement(agg, 1, "$sort");
+		assertThat(sort, is((DBObject) new BasicDBObject("a", -1)));
+
+		DBObject group = extractPipelineElement(agg, 2, "$group");
+		assertThat(group,
+				is((DBObject) new BasicDBObject("_id", "$someKey").append("doc", new BasicDBObject("$first", "$$ROOT"))));
+	}
+
 	private DBObject extractPipelineElement(DBObject agg, int index, String operation) {
 
 		List<DBObject> pipeline = (List<DBObject>) agg.get("pipeline");