Merge pull request #2578 from idoshr/Update_with_Aggregation_Pipeline

Update with aggregation pipeline

Author: bagerard

docs/guide/querying.rst

@@ -218,12 +218,35 @@ However, this doesn't map well to the syntax so you can also use a capital S ins
 
 Raw queries
 -----------
-It is possible to provide a raw :mod:`PyMongo` query as a query parameter, which will
-be integrated directly into the query. This is done using the ``__raw__``
+It is possible to provide a raw :mod:`PyMongo` query as a query parameter or update as a update parameter , which will
+be integrated directly into the query or update. This is done using the ``__raw__``
 keyword argument::
 
     Page.objects(__raw__={'tags': 'coding'})
 
+    # or for update
+
+    Page.objects(__raw__={'tags': 'coding'}).update(__raw__={'$set': {'tags': 'coding'}})
+
+    Page.objects(tags='coding').update(__raw__={'$set': {'tags': 'coding'}})
+
+.. versionadded:: 0.4
+
+
+Update with Aggregation Pipeline
+-----------
+It is possible to provide a raw :mod:`PyMongo` aggregation update parameter, which will
+be integrated directly into the update. This is done by using ``__raw__`` field and value of array
+pipeline
+`Update with Aggregation Pipeline <https://docs.mongodb.com/manual/reference/method/db.collection.updateMany/#update-with-aggregation->`_
+keyword argument::
+
+    # 'tags' field is set to 'coding is fun'
+    Page.objects(tags='coding').update(__raw__=[
+        {"$set": {"tags": {"$concat": ["$tags", "is fun"]}}}
+        ],
+    )
+
 .. versionadded:: 0.4
 
 Sorting/Ordering results

mongoengine/queryset/base.py

@@ -551,8 +551,13 @@ def update(
 
         queryset = self.clone()
         query = queryset._query
-        update = transform.update(queryset._document, **update)
-
+        if "__raw__" in update and isinstance(update["__raw__"], list):
+            update = [
+                transform.update(queryset._document, **{"__raw__": u})
+                for u in update["__raw__"]
+            ]
+        else:
+            update = transform.update(queryset._document, **update)
         # If doing an atomic upsert on an inheritable class
         # then ensure we add _cls to the update operation
         if upsert and "_cls" in query:

tests/queryset/test_queryset.py

@@ -25,6 +25,7 @@
     queryset_manager,
 )
 from tests.utils import (
+    requires_mongodb_gte_42,
     requires_mongodb_gte_44,
     requires_mongodb_lt_42,
 )
@@ -2217,6 +2218,36 @@ class BlogPost(Document):
         post.reload()
         assert post.tags == ["code", "mongodb"]
 
+    @requires_mongodb_gte_42
+    def test_aggregation_update(self):
+        """Ensure that the 'aggregation_update' update 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"
+
+        BlogPost.objects(slug="test test").update(
+            __raw__=[
+                {"$set": {"slug": {"$concat": ["$slug", " ", "it"]}}},  # test test it
+                {
+                    "$set": {"slug": {"$concat": ["When", " ", "$slug"]}}
+                },  # When test test it
+            ],
+        )
+        post.reload()
+        assert post.slug == "When test test it"
+
     def test_add_to_set_each(self):
         class Item(Document):
             name = StringField(required=True)

tests/utils.py

@@ -37,6 +37,10 @@ def requires_mongodb_lt_42(func):
     return _decorated_with_ver_requirement(func, (4, 2), oper=operator.lt)
 
 
+def requires_mongodb_gte_42(func):
+    return _decorated_with_ver_requirement(func, (4, 2), oper=operator.ge)
+
+
 def requires_mongodb_gte_44(func):
     return _decorated_with_ver_requirement(func, (4, 4), oper=operator.ge)