fix(server/mongo): evolution step filters on all index columns

Signed-off-by: Luka Peschke <[email protected]>
ToucanToco · Feb 6, 2024 · 30b5c28 · 30b5c28
1 parent bdafb61
commit 30b5c28
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 1 deletion.
diff --git a/server/CHANGELOG.md b/server/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+- Mongo: Fixed the `evolution` step when several index columns are used. The generated pipeline now filters on all of them
+  rather than jsut the last one.
+
 ## [0.26.9] - 2023-10-05
 
 - Pypika: Fixed the `todate` step by ensuring the right formatting is used for every backend

diff --git a/server/src/weaverbird/backends/mongo_translator/steps/evolution.py b/server/src/weaverbird/backends/mongo_translator/steps/evolution.py
@@ -82,7 +82,10 @@ def translate_evolution(step: EvolutionStep) -> list[MongoStep]:
                         "cond": {
                             "$and": [
                                 {"$eq": ["$_VQB_DATE_PREV", f"$$item.{step.date_col}"]},
-                                {"$eq": [f"${col}", f"$$item.{col}"] for col in step.index_columns},
+                                *(
+                                    {"$eq": [f"${col}", f"$$item.{col}"]}
+                                    for col in step.index_columns
+                                ),
                             ],
                         },
                     },

diff --git a/server/tests/backends/mongo_translator/steps/test_evolution.py b/server/tests/backends/mongo_translator/steps/test_evolution.py
@@ -0,0 +1,85 @@
+from weaverbird.backends.mongo_translator.steps.evolution import translate_evolution
+from weaverbird.pipeline.steps.evolution import EvolutionStep
+
+
+def test_translate_evolution() -> None:
+    step = EvolutionStep(
+        dateCol="foo",
+        valueCol="bar",
+        evolutionFormat="abs",
+        evolutionType="vsLastDay",
+        indexColumns=["a", "b", "c"],
+    )
+
+    assert translate_evolution(step) == [
+        {
+            "$addFields": {
+                "_VQB_DATE_PREV": {
+                    "$dateSubtract": {"amount": 1, "startDate": "$foo", "unit": "week"}
+                }
+            }
+        },
+        {
+            "$facet": {
+                "_VQB_COPIES_ARRAY": [
+                    {"$group": {"_VQB_ALL_DOCS": {"$push": "$$ROOT"}, "_id": None}}
+                ],
+                "_VQB_ORIGINALS": [{"$project": {"_id": 0}}],
+            }
+        },
+        {"$unwind": "$_VQB_ORIGINALS"},
+        {
+            "$project": {
+                "_VQB_ORIGINALS": {
+                    "$mergeObjects": [
+                        "$_VQB_ORIGINALS",
+                        {"$arrayElemAt": ["$_VQB_COPIES_ARRAY", 0]},
+                    ]
+                }
+            }
+        },
+        {"$replaceRoot": {"newRoot": "$_VQB_ORIGINALS"}},
+        {
+            "$addFields": {
+                "_VQB_ALL_DOCS": {
+                    "$filter": {
+                        "as": "item",
+                        "cond": {
+                            "$and": [
+                                {"$eq": ["$_VQB_DATE_PREV", "$$item.foo"]},
+                                {"$eq": ["$a", "$$item.a"]},
+                                {"$eq": ["$b", "$$item.b"]},
+                                {"$eq": ["$c", "$$item.c"]},
+                            ]
+                        },
+                        "input": "$_VQB_ALL_DOCS",
+                    }
+                }
+            }
+        },
+        {
+            "$addFields": {
+                "_VQB_VALUE_PREV": {
+                    "$cond": [
+                        {"$gt": [{"$size": "$_VQB_ALL_DOCS.bar"}, 1]},
+                        "Error",
+                        {"$arrayElemAt": ["$_VQB_ALL_DOCS.bar", 0]},
+                    ]
+                }
+            }
+        },
+        {
+            "$addFields": {
+                "bar_EVOL_ABS": {
+                    "$cond": [
+                        {"$eq": ["$_VQB_VALUE_PREV", "Error"]},
+                        "Error: More than one previous "
+                        "date found for the specified "
+                        "index columns",
+                        {"$subtract": ["$bar", "$_VQB_VALUE_PREV"]},
+                    ]
+                }
+            }
+        },
+        {"$project": {"_VQB_ALL_DOCS": 0, "_VQB_DATE_PREV": 0, "_VQB_VALUE_PREV": 0}},
+    ]