From 83923e0e2a330688c55f05dd73d7df9d8a2e8369 Mon Sep 17 00:00:00 2001 From: sangyongchoi Date: Sun, 15 Jan 2023 18:24:40 +0900 Subject: [PATCH] Add support for 'let' and 'pipeline' in $lookup This commit introduces let and pipline to the Lookup aggregation stage. Closes: #3322 Original Pull Request: #4272 --- .../mongodb/core/aggregation/Aggregation.java | 13 ++++ .../core/aggregation/LookupOperation.java | 66 +++++++++++++++++++ .../core/aggregation/AggregationTests.java | 46 ++++++++++++- 3 files changed, 123 insertions(+), 2 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 2d69c799e..cc09f54ec 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 @@ -50,6 +50,7 @@ import org.springframework.util.Assert; * @author Nikolay Bogdanov * @author Gustavo de Geus * @author Jérôme Guyon + * @author Sangyong Choi * @since 1.3 */ public class Aggregation { @@ -664,6 +665,18 @@ public class Aggregation { return new LookupOperation(from, localField, foreignField, as); } + public static LookupOperation lookup(String from, String localField, String foreignField, String as, List aggregationOperations) { + return lookup(field(from), field(localField), field(foreignField), field(as), null, new AggregationPipeline(aggregationOperations)); + } + + public static LookupOperation lookup(String from, String localField, String foreignField, String as, List letExpressionVars, List aggregationOperations) { + return lookup(field(from), field(localField), field(foreignField), field(as), new LookupOperation.Let(letExpressionVars), new AggregationPipeline(aggregationOperations)); + } + + public static LookupOperation lookup(Field from, Field localField, Field foreignField, Field as, LookupOperation.Let let, AggregationPipeline pipeline) { + return new LookupOperation(from, localField, foreignField, as, let, pipeline); + } + /** * Creates a new {@link CountOperationBuilder}. * diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java index 0439876db..ff7999e43 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java @@ -15,6 +15,8 @@ */ package org.springframework.data.mongodb.core.aggregation; +import java.util.List; + import org.bson.Document; import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField; import org.springframework.data.mongodb.core.aggregation.FieldsExposingAggregationOperation.InheritsFieldsAggregationOperation; @@ -28,6 +30,7 @@ import org.springframework.util.Assert; * @author Alessio Fachechi * @author Christoph Strobl * @author Mark Paluch + * @author Sangyong Choi * @since 1.9 * @see MongoDB Aggregation Framework: * $lookup @@ -39,6 +42,11 @@ public class LookupOperation implements FieldsExposingAggregationOperation, Inhe private final Field foreignField; private final ExposedField as; + @Nullable + private final Let let; + @Nullable + private final AggregationPipeline pipeline; + /** * Creates a new {@link LookupOperation} for the given {@link Field}s. * @@ -48,7 +56,10 @@ public class LookupOperation implements FieldsExposingAggregationOperation, Inhe * @param as must not be {@literal null}. */ public LookupOperation(Field from, Field localField, Field foreignField, Field as) { + this(from, localField, foreignField, as, null, null); + } + public LookupOperation(Field from, Field localField, Field foreignField, Field as, @Nullable Let let, @Nullable AggregationPipeline pipeline) { Assert.notNull(from, "From must not be null"); Assert.notNull(localField, "LocalField must not be null"); Assert.notNull(foreignField, "ForeignField must not be null"); @@ -58,6 +69,8 @@ public class LookupOperation implements FieldsExposingAggregationOperation, Inhe this.localField = localField; this.foreignField = foreignField; this.as = new ExposedField(as, true); + this.let = let; + this.pipeline = pipeline; } @Override @@ -75,6 +88,14 @@ public class LookupOperation implements FieldsExposingAggregationOperation, Inhe lookupObject.append("foreignField", foreignField.getTarget()); lookupObject.append("as", as.getTarget()); + if (let != null) { + lookupObject.append("let", let.toDocument(context)); + } + + if (pipeline != null) { + lookupObject.append("pipeline", pipeline.toDocuments(context)); + } + return new Document(getOperator(), lookupObject); } @@ -184,4 +205,49 @@ public class LookupOperation implements FieldsExposingAggregationOperation, Inhe return this; } } + + public static class Let implements AggregationExpression{ + + private final List vars; + + public Let(List vars) { + Assert.notEmpty(vars, "'let' must not be null or empty"); + this.vars = vars; + } + + @Override + public Document toDocument(AggregationOperationContext context) { + return toLet(); + } + + private Document toLet() { + Document mappedVars = new Document(); + + for (ExpressionVariable var : this.vars) { + mappedVars.putAll(getMappedVariable(var)); + } + + return mappedVars; + } + + private Document getMappedVariable(ExpressionVariable var) { + return new Document(var.variableName, prefixDollarSign(var.expression)); + } + + private String prefixDollarSign(String expression) { + return "$" + expression; + } + + public static class ExpressionVariable { + + private final String variableName; + + private final String expression; + + public ExpressionVariable(String variableName, String expression) { + this.variableName = variableName; + this.expression = expression; + } + } + } } 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 c98b0c0b5..77f06796c 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 @@ -43,7 +43,6 @@ import java.util.stream.Stream; import org.assertj.core.data.Offset; import org.bson.Document; -import org.bson.types.ObjectId; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -90,6 +89,7 @@ import com.mongodb.client.MongoCollection; * @author Maninder Singh * @author Sergey Shcherbakov * @author Minsu Kim + * @author Sangyong Choi */ @ExtendWith(MongoTemplateExtension.class) public class AggregationTests { @@ -499,7 +499,7 @@ public class AggregationTests { /* //complex mongodb aggregation framework example from https://docs.mongodb.org/manual/tutorial/aggregation-examples/#largest-and-smallest-cities-by-state - + db.zipcodes.aggregate( { $group: { @@ -1518,6 +1518,48 @@ public class AggregationTests { assertThat(firstItem).containsEntry("linkedPerson.[0].firstname", "u1"); } + @Test + void shouldLookupPeopleCorrectlyWithPipeline() { + createUsersWithReferencedPersons(); + + TypedAggregation agg = newAggregation(User.class, // + lookup("person", "_id", "firstname", "linkedPerson", List.of(match(where("firstname").is("u1")))), // + sort(ASC, "id")); + + AggregationResults results = mongoTemplate.aggregate(agg, User.class, Document.class); + + List mappedResults = results.getMappedResults(); + + Document firstItem = mappedResults.get(0); + + assertThat(firstItem).containsEntry("_id", "u1"); + assertThat(firstItem).containsEntry("linkedPerson.[0].firstname", "u1"); + } + + @Test + void shouldLookupPeopleCorrectlyWithPipelineAndLet() { + createUsersWithReferencedPersons(); + + TypedAggregation agg = newAggregation(User.class, // + lookup( + "person", + "_id", + "firstname", + "linkedPerson", + List.of(new LookupOperation.Let.ExpressionVariable("personFirstname", "firstname")), + List.of(match(where("firstname").is("u1")))), + sort(ASC, "id")); + + AggregationResults results = mongoTemplate.aggregate(agg, User.class, Document.class); + + List mappedResults = results.getMappedResults(); + + Document firstItem = mappedResults.get(0); + + assertThat(firstItem).containsEntry("_id", "u1"); + assertThat(firstItem).containsEntry("linkedPerson.[0].firstname", "u1"); + } + @Test // DATAMONGO-1326 void shouldGroupByAndLookupPeopleCorectly() {