Support update with aggregation pipeline in modify method

Author: gowthamvbhat

mongoengine/queryset/base.py

@@ -733,7 +733,15 @@ def modify(
             query["$where"] = where_clause
 
         if not remove:
-            update = transform.update(queryset._document, **update)
+            if "__raw__" in update and isinstance(
+                update["__raw__"], list
+            ):  # Case of Update with Aggregation Pipeline
+                update = [
+                    transform.update(queryset._document, **{"__raw__": u})
+                    for u in update["__raw__"]
+                ]
+            else:
+                update = transform.update(queryset._document, **update)
         sort = queryset._ordering
 
         try:

tests/queryset/test_modify.py

@@ -130,6 +130,35 @@ class BlogPost(Document):
         )
         assert blog.tags == ["python", "go", "rust", "code", "java"]
 
+    def test_modify_with_aggregation_update(self):
+        """Ensure that the 'aggregation_update' modify works correctly."""
+
+        class BlogPost(Document):
+            slug = StringField()
+            tags = ListField(StringField())
+
+        BlogPost.drop_collection()
+
+        post = BlogPost(slug="test")
+        post.save()
+
+        BlogPost.objects(slug="test").update(
+            __raw__=[{"$set": {"slug": {"$concat": ["$slug", " ", "$slug"]}}}],
+        )
+        post.reload()
+        assert post.slug == "test test"
+
+        post = BlogPost.objects(slug="test test").modify(
+            __raw__=[
+                {"$set": {"slug": {"$concat": ["$slug", " ", "it"]}}},  # test test it
+                {
+                    "$set": {"slug": {"$concat": ["When", " ", "$slug"]}}
+                },  # When test test it
+            ],
+            new=True,
+        )
+        assert post.slug == "When test test it"
+
 
 if __name__ == "__main__":
     unittest.main()