Adding update_match

Author: apetenchea

arangoasync/collection.py

@@ -1164,6 +1164,79 @@ def response_handler(resp: Response) -> Cursor:
 
         return await self._executor.execute(request, response_handler)
 
+    async def update_match(
+        self,
+        filters: Json,
+        body: T,
+        limit: Optional[int | str] = None,
+        keep_none: Optional[bool] = None,
+        wait_for_sync: Optional[bool] = None,
+        merge_objects: Optional[bool] = None,
+        allow_dirty_read: Optional[bool] = None,
+    ) -> Result[int]:
+        """Update matching documents.
+
+        Args:
+            filters (dict | None): Query filters.
+            body (dict): Full or partial document body with the updates.
+            limit (int | str | None): Maximum number of documents to return.
+            keep_none (bool | None): If set to `True`, fields with value `None` are
+                retained in the document. Otherwise, they are removed completely.
+            wait_for_sync (bool | None): Wait until operation has been synced to disk.
+            merge_objects (bool | None): If set to `True`, sub-dictionaries are merged
+                instead of the new one overwriting the old one.
+            allow_dirty_read (bool | None): Allow reads from followers in a cluster.
+
+        Returns:
+            int: Number of documents updated.
+
+        Raises:
+           DocumentUpdateError: If update fails.
+        """
+        if not self._is_none_or_dict(filters):
+            raise ValueError("filters parameter must be a dict")
+        if not (self._is_none_or_int(limit) or limit == "null"):
+            raise ValueError("limit parameter must be a non-negative int")
+
+        # If the waitForSync parameter is not specified or set to false,
+        # then the collection’s default waitForSync behavior is applied.
+        sync = f", waitForSync: {wait_for_sync}" if wait_for_sync is not None else ""
+        query = f"""
+            FOR doc IN @@collection
+                {self._build_filter_conditions(filters)}
+                {f"LIMIT {limit}" if limit is not None else ""}
+                UPDATE doc WITH @body IN @@collection
+                OPTIONS {{ keepNull: @keep_none, mergeObjects: @merge {sync} }}
+        """  # noqa: E201 E202
+        bind_vars = {
+            "@collection": self.name,
+            "body": body,
+            "keep_none": keep_none,
+            "merge": merge_objects,
+        }
+        data = {"query": query, "bindVars": bind_vars}
+        headers: RequestHeaders = {}
+        if allow_dirty_read is not None:
+            if allow_dirty_read is True:
+                headers["x-arango-allow-dirty-read"] = "true"
+            else:
+                headers["x-arango-allow-dirty-read"] = "false"
+
+        request = Request(
+            method=Method.POST,
+            endpoint="/_api/cursor",
+            data=self.serializer.dumps(data),
+            headers=headers,
+        )
+
+        def response_handler(resp: Response) -> int:
+            if resp.is_success:
+                result = self.deserializer.loads(resp.raw_body)
+                return cast(int, result["extra"]["stats"]["writesExecuted"])
+            raise DocumentUpdateError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
     async def insert_many(
         self,
         documents: Sequence[T],

tests/test_document.py

@@ -452,3 +452,36 @@ async def test_document_delete_many(doc_col, bad_col, docs):
     assert len(result) == len(docs)
     for res in result:
         assert "error" in res
+
+
+@pytest.mark.asyncio
+async def test_document_update_match(doc_col, bad_col, docs):
+    # Check errors first
+    with pytest.raises(DocumentUpdateError):
+        await bad_col.update_match({}, {})
+    with pytest.raises(ValueError):
+        await doc_col.update_match({}, {}, limit=-1)
+    with pytest.raises(ValueError):
+        await doc_col.update_match("abcd", {})
+
+    # Insert all documents
+    await doc_col.insert_many(docs)
+
+    # Update value for all documents
+    count = await doc_col.update_match({}, {"val": 42})
+    async for doc in await doc_col.find():
+        assert doc["val"] == 42
+    assert count == len(docs)
+
+    # Update documents partially
+    count = await doc_col.update_match({"text": "foo"}, {"val": 24})
+    async for doc in await doc_col.find():
+        if doc["text"] == "foo":
+            assert doc["val"] == 24
+    assert count == sum([1 for doc in docs if doc["text"] == "foo"])
+
+    # No matching documents
+    count = await doc_col.update_match({"text": "no_matching"}, {"val": -1})
+    async for doc in await doc_col.find():
+        assert doc["val"] != -1
+    assert count == 0