Fix expression evaluation for documents in a List for projections.

bwaldvogel · Dec 21, 2020 · a6c79a4 · a6c79a4
1 parent 2d81ecc
commit a6c79a4
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/ProjectStage.java b/core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/ProjectStage.java
@@ -2,7 +2,9 @@
 
 import static de.bwaldvogel.mongo.backend.Constants.ID_FIELD;
 
+import java.util.List;
 import java.util.Map.Entry;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import de.bwaldvogel.mongo.backend.Missing;
@@ -54,7 +56,15 @@ Document projectDocument(Document document) {
                 } else {
                     Utils.removeSubdocumentValue(result, field);
                 }
-            } else if (projectionValue == null) {
+            }
+            else if (projectionValue instanceof List) {
+                List<Object> resolvedProjectionValues = ((List<Object>) projectionValue)
+                    .stream()
+                    .map(value -> Expression.evaluateDocument(value, document))
+                    .collect(Collectors.toList());
+                result.put(field, resolvedProjectionValues);
+            }
+            else if (projectionValue == null) {
                 result.put(field, null);
             } else {
                 Object value = Expression.evaluateDocument(projectionValue, document);

diff --git a/core/src/test/java/de/bwaldvogel/mongo/backend/aggregation/stage/ProjectStageTest.java b/core/src/test/java/de/bwaldvogel/mongo/backend/aggregation/stage/ProjectStageTest.java
@@ -4,6 +4,8 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 
+import java.util.Collections;
+
 import org.junit.jupiter.api.Test;
 
 import de.bwaldvogel.mongo.bson.Document;
@@ -41,6 +43,24 @@ void testProject_withNestedInclusion() throws Exception {
             .isEqualTo(json("_id: 1, x: {b: 2, c: 3}"));
     }
 
+    @Test
+    void testProject_withFieldToBeEvaluated() {
+        Document projection = new Document();
+        projection.put("_id", 1);
+        projection.put("x", new Document("count", "$count"));
+        assertThat(project(json("_id: 1, count: 5"), projection))
+            .isEqualTo(json("_id: 1, x: {count: 5}"));
+    }
+
+    @Test
+    void testProject_withFieldWithinArrayToBeEvaluated() {
+        Document projection = new Document();
+        projection.put("_id", 1);
+        projection.put("x", Collections.singletonList(new Document("count", "$count")));
+        assertThat(project(json("_id: 1, count: 5"), projection))
+            .isEqualTo(json("_id: 1, x: [{count: 5}]"));
+    }
+
     private static Document project(Document document, Document projection) {
         return new ProjectStage(projection).projectDocument(document);
     }