From 10d3c510ea7b383463635092b95c99cac7236db8 Mon Sep 17 00:00:00 2001
From: Ido Shraga <ido@cyabra.com>
Date: Sat, 16 Oct 2021 12:30:49 +0300
Subject: [PATCH 1/5] Update with Aggregation Pipeline

---
 mongoengine/queryset/base.py    | 11 +++++++++--
 tests/queryset/test_queryset.py | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index 5dc47e001..f5d868b83 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -524,6 +524,7 @@ def update(
         write_concern=None,
         read_concern=None,
         full_result=False,
+        aggregation_update=False,
         **update,
     ):
         """Perform an atomic update on the fields matched by the query.
@@ -539,6 +540,7 @@ def update(
         :param read_concern: Override the read concern for the operation
         :param full_result: Return the associated ``pymongo.UpdateResult`` rather than just the number
             updated items
+        :param aggregation_update: Update with Aggregation Pipeline https://docs.mongodb.com/manual/reference/method/db.collection.updateMany/#update-with-aggregation-pipeline
         :param update: Django-style update keyword arguments
 
         :returns the number of updated documents (unless ``full_result`` is True)
@@ -551,8 +553,13 @@ def update(
 
         queryset = self.clone()
         query = queryset._query
-        update = transform.update(queryset._document, **update)
-
+        if aggregation_update:
+            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:
diff --git a/tests/queryset/test_queryset.py b/tests/queryset/test_queryset.py
index 9a2f2862c..d7fe10a28 100644
--- a/tests/queryset/test_queryset.py
+++ b/tests/queryset/test_queryset.py
@@ -2217,6 +2217,25 @@ class BlogPost(Document):
         post.reload()
         assert post.tags == ["code", "mongodb"]
 
+        BlogPost.objects(slug="test").update(
+            __raw__=[{"$set": {"slug": {"$concat": ["$slug", " ", "$slug"]}}}],
+            aggregation_update=True,
+        )
+        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
+            ],
+            aggregation_update=True,
+        )
+        post.reload()
+        assert post.slug == "When test test it"
+
     def test_add_to_set_each(self):
         class Item(Document):
             name = StringField(required=True)

From e35dab56bac4e7d7e0c9fb48cbcce0f85eb4a198 Mon Sep 17 00:00:00 2001
From: Ido Shraga <ido@cyabra.com>
Date: Sat, 16 Oct 2021 18:07:56 +0300
Subject: [PATCH 2/5] Update with Aggregation Pipeline doc

---
 docs/guide/querying.rst      | 27 +++++++++++++++++++++++++--
 mongoengine/queryset/base.py |  4 ++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/docs/guide/querying.rst b/docs/guide/querying.rst
index 1cf5ec3c6..1b608608c 100644
--- a/docs/guide/querying.rst
+++ b/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 using the ``aggregation_update=True`` and ``__raw__``
+keyword argument::
+
+
+    # 'tags' field is set to 'coding is fun'
+    Page.objects(tags='coding').update(__raw__=[
+        {"$set": {"tags": {"$concat": ["$tags", "is fun"]}}}
+        ],
+        aggregation_update=True,
+    )
+
 .. versionadded:: 0.4
 
 Sorting/Ordering results
diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index 5261949a5..62f5658e8 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -554,6 +554,10 @@ def update(
         queryset = self.clone()
         query = queryset._query
         if aggregation_update:
+            if "__raw__" not in update:
+                raise OperationError(
+                    "Currently aggregation_update works only with __raw__ value"
+                )
             update = [
                 transform.update(queryset._document, **{"__raw__": u})
                 for u in update["__raw__"]

From b8cf4c202044f7b0d533e1cd0cf8728e67c0c787 Mon Sep 17 00:00:00 2001
From: Ido Shraga <ido@cyabra.com>
Date: Sat, 16 Oct 2021 18:12:35 +0300
Subject: [PATCH 3/5] Update with Aggregation Pipeline doc

---
 docs/guide/querying.rst | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/guide/querying.rst b/docs/guide/querying.rst
index 1b608608c..7471aa158 100644
--- a/docs/guide/querying.rst
+++ b/docs/guide/querying.rst
@@ -237,9 +237,10 @@ 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 using the ``aggregation_update=True`` and ``__raw__``
+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"]}}}

From 04a049f4ce2c8538d21f8a8c754e84d50ba91478 Mon Sep 17 00:00:00 2001
From: Ido Shraga <ido@cyabra.com>
Date: Sun, 17 Oct 2021 13:14:04 +0300
Subject: [PATCH 4/5] fix tester

---
 tests/queryset/test_queryset.py | 14 ++++++++++++++
 tests/utils.py                  |  4 ++++
 2 files changed, 18 insertions(+)

diff --git a/tests/queryset/test_queryset.py b/tests/queryset/test_queryset.py
index d7fe10a28..79f510805 100644
--- a/tests/queryset/test_queryset.py
+++ b/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,19 @@ 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"]}}}],
             aggregation_update=True,
diff --git a/tests/utils.py b/tests/utils.py
index a05b9c14b..0dcdb2dbf 100644
--- a/tests/utils.py
+++ b/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)
 

From 6be253fc6691c4ebca4ce8fd6fefb18ede6c5833 Mon Sep 17 00:00:00 2001
From: Ido Shraga <ido@cyabra.com>
Date: Sat, 6 Nov 2021 17:15:43 +0200
Subject: [PATCH 5/5] following @bagerard point and remove the extra
 aggregation_update parameter

---
 docs/guide/querying.rst         | 3 +--
 mongoengine/queryset/base.py    | 8 +-------
 tests/queryset/test_queryset.py | 2 --
 3 files changed, 2 insertions(+), 11 deletions(-)

diff --git a/docs/guide/querying.rst b/docs/guide/querying.rst
index 7471aa158..1d7617b40 100644
--- a/docs/guide/querying.rst
+++ b/docs/guide/querying.rst
@@ -236,7 +236,7 @@ keyword argument::
 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 using the ``aggregation_update=True`` and ``__raw__``
+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::
@@ -245,7 +245,6 @@ keyword argument::
     Page.objects(tags='coding').update(__raw__=[
         {"$set": {"tags": {"$concat": ["$tags", "is fun"]}}}
         ],
-        aggregation_update=True,
     )
 
 .. versionadded:: 0.4
diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index 62f5658e8..98b4b66e4 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -524,7 +524,6 @@ def update(
         write_concern=None,
         read_concern=None,
         full_result=False,
-        aggregation_update=False,
         **update,
     ):
         """Perform an atomic update on the fields matched by the query.
@@ -540,7 +539,6 @@ def update(
         :param read_concern: Override the read concern for the operation
         :param full_result: Return the associated ``pymongo.UpdateResult`` rather than just the number
             updated items
-        :param aggregation_update: Update with Aggregation Pipeline https://docs.mongodb.com/manual/reference/method/db.collection.updateMany/#update-with-aggregation-pipeline
         :param update: Django-style update keyword arguments
 
         :returns the number of updated documents (unless ``full_result`` is True)
@@ -553,11 +551,7 @@ def update(
 
         queryset = self.clone()
         query = queryset._query
-        if aggregation_update:
-            if "__raw__" not in update:
-                raise OperationError(
-                    "Currently aggregation_update works only with __raw__ value"
-                )
+        if "__raw__" in update and isinstance(update["__raw__"], list):
             update = [
                 transform.update(queryset._document, **{"__raw__": u})
                 for u in update["__raw__"]
diff --git a/tests/queryset/test_queryset.py b/tests/queryset/test_queryset.py
index 79f510805..942db6255 100644
--- a/tests/queryset/test_queryset.py
+++ b/tests/queryset/test_queryset.py
@@ -2233,7 +2233,6 @@ class BlogPost(Document):
 
         BlogPost.objects(slug="test").update(
             __raw__=[{"$set": {"slug": {"$concat": ["$slug", " ", "$slug"]}}}],
-            aggregation_update=True,
         )
         post.reload()
         assert post.slug == "test test"
@@ -2245,7 +2244,6 @@ class BlogPost(Document):
                     "$set": {"slug": {"$concat": ["When", " ", "$slug"]}}
                 },  # When test test it
             ],
-            aggregation_update=True,
         )
         post.reload()
         assert post.slug == "When test test it"