From cadcbf61067efe06d1a010b27f1870a6ab0cdd03 Mon Sep 17 00:00:00 2001 From: Thomas Darimont Date: Fri, 13 Jun 2014 14:44:40 +0200 Subject: [PATCH] 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: #190. --- .../mongodb/core/aggregation/Aggregation.java | 60 +++++++++++++++- .../data/mongodb/core/aggregation/Fields.java | 7 +- .../core/aggregation/GroupOperation.java | 11 ++- .../core/aggregation/ProjectionOperation.java | 4 ++ .../core/aggregation/AggregationTests.java | 68 +++++++++++++++++++ .../aggregation/AggregationUnitTests.java | 27 ++++++++ 6 files changed, 174 insertions(+), 3 deletions(-) diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java index 0341d7398..a9cf63535 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java @@ -44,11 +44,22 @@ import com.mongodb.DBObject; */ 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 operations; - private final AggregationOptions options; /** @@ -363,4 +374,51 @@ public class Aggregation { 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()); + } + } } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Fields.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Fields.java index dc7ee6615..3c7160a77 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Fields.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Fields.java @@ -30,6 +30,7 @@ import org.springframework.util.StringUtils; * 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 { @@ -186,7 +187,7 @@ public final class Fields implements Iterable { 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 @@ public final class Fields implements Iterable { return source; } + if (Aggregation.SystemVariable.isReferingToSystemVariable(source)) { + return source; + } + int dollarIndex = source.lastIndexOf('$'); return dollarIndex == -1 ? source : source.substring(dollarIndex + 1); } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/GroupOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/GroupOperation.java index dd1009e4d..11d285ec7 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/GroupOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/GroupOperation.java @@ -364,7 +364,16 @@ public class GroupOperation implements FieldsExposingAggregationOperation { } 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 diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ProjectionOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ProjectionOperation.java index bbc89bc0f..45139b23d 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ProjectionOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ProjectionOperation.java @@ -627,6 +627,10 @@ public class ProjectionOperation implements FieldsExposingAggregationOperation { // 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(); diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java index 018c0c5fc..3f6b002f5 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java @@ -44,11 +44,13 @@ import org.springframework.beans.factory.annotation.Autowired; 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 @@ public class AggregationTests { 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 class AggregationTests { 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 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 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 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 class AggregationTests { } } } + + 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; + } + } } diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationUnitTests.java index 2cb7cf6c4..3116560f9 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationUnitTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationUnitTests.java @@ -27,6 +27,7 @@ import java.util.List; 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 class AggregationUnitTests { )); } + /** + * @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 pipeline = (List) agg.get("pipeline");