Adding replace_match

Author: apetenchea

arangoasync/collection.py

@@ -1179,7 +1179,7 @@ async def update_match(
         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.
+            limit (int | str | None): Maximum number of documents to update.
             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.
@@ -1198,8 +1198,6 @@ async def update_match(
         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
@@ -1237,6 +1235,69 @@ def response_handler(resp: Response) -> int:
 
         return await self._executor.execute(request, response_handler)
 
+    async def replace_match(
+        self,
+        filters: Json,
+        body: T,
+        limit: Optional[int | str] = None,
+        wait_for_sync: Optional[bool] = None,
+        allow_dirty_read: Optional[bool] = None,
+    ) -> Result[int]:
+        """Replace matching documents.
+
+        Args:
+            filters (dict | None): Query filters.
+            body (dict): New document body.
+            limit (int | str | None): Maximum number of documents to replace.
+            wait_for_sync (bool | None): Wait until operation has been synced to disk.
+            allow_dirty_read (bool | None): Allow reads from followers in a cluster.
+
+        Returns:
+            int: Number of documents that got replaced.
+
+        Raises:
+           DocumentReplaceError: If replace 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")
+
+        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 ""}
+                REPLACE doc WITH @body IN @@collection
+                {f"OPTIONS {{ {sync} }}" if sync else ""}
+        """  # noqa: E201 E202
+        bind_vars = {
+            "@collection": self.name,
+            "body": body,
+        }
+        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 DocumentReplaceError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
     async def insert_many(
         self,
         documents: Sequence[T],

tests/test_document.py

@@ -485,3 +485,41 @@ async def test_document_update_match(doc_col, bad_col, docs):
     async for doc in await doc_col.find():
         assert doc["val"] != -1
     assert count == 0
+
+
+@pytest.mark.asyncio
+async def test_document_replace_match(doc_col, bad_col, docs):
+    # Check errors first
+    with pytest.raises(DocumentReplaceError):
+        await bad_col.replace_match({}, {})
+    with pytest.raises(ValueError):
+        await doc_col.replace_match({}, {}, limit=-1)
+    with pytest.raises(ValueError):
+        await doc_col.replace_match("abcd", {})
+
+    # Replace all documents
+    await doc_col.insert_many(docs)
+    count = await doc_col.replace_match({}, {"replacement": 42})
+    async for doc in await doc_col.find():
+        assert "replacement" in doc
+        assert "val" not in doc
+    assert count == len(docs)
+    await doc_col.truncate()
+
+    # Replace documents partially
+    await doc_col.insert_many(docs)
+    count = await doc_col.replace_match({"text": "foo"}, {"replacement": 24})
+    async for doc in await doc_col.find():
+        if doc.get("text") == "bar":
+            assert "replacement" not in doc
+        else:
+            assert "replacement" in doc
+    assert count == sum([1 for doc in docs if doc["text"] == "foo"])
+    await doc_col.truncate()
+
+    # No matching documents
+    await doc_col.insert_many(docs)
+    count = await doc_col.replace_match({"text": "no_matching"}, {"val": -1})
+    async for doc in await doc_col.find():
+        assert doc["val"] != -1
+    assert count == 0