Adding truncate and count operations

Author: apetenchea

arangoasync/collection.py

@@ -11,6 +11,8 @@
 )
 from arangoasync.exceptions import (
     CollectionPropertiesError,
+    CollectionTruncateError,
+    DocumentCountError,
     DocumentDeleteError,
     DocumentGetError,
     DocumentInsertError,
@@ -392,6 +394,66 @@ def response_handler(resp: Response) -> CollectionProperties:
 
         return await self._executor.execute(request, response_handler)
 
+    async def truncate(
+        self,
+        wait_for_sync: Optional[bool] = None,
+        compact: Optional[bool] = None,
+    ) -> None:
+        """Removes all documents, but leaves indexes intact.
+
+        Args:
+            wait_for_sync (bool | None): If set to `True`, the data is synchronized
+                to disk before returning from the truncate operation.
+            compact (bool | None): If set to `True`, the storage engine is told to
+                start a compaction in order to free up disk space. This can be
+                resource intensive. If the only intention is to start over with an
+                empty collection, specify `False`.
+
+        Raises:
+            CollectionTruncateError: If truncation fails.
+
+        References:
+            - `truncate-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#truncate-a-collection>`__
+        """  # noqa: E501
+        params: Params = {}
+        if wait_for_sync is not None:
+            params["waitForSync"] = wait_for_sync
+        if compact is not None:
+            params["compact"] = compact
+
+        request = Request(
+            method=Method.PUT,
+            endpoint=f"/_api/collection/{self.name}/truncate",
+            params=params,
+        )
+
+        def response_handler(resp: Response) -> None:
+            if not resp.is_success:
+                raise CollectionTruncateError(resp, request)
+
+        await self._executor.execute(request, response_handler)
+
+    async def count(self) -> Result[int]:
+        """Return the total document count.
+
+        Returns:
+            int: Total document count.
+
+        Raises:
+            DocumentCountError: If retrieval fails.
+        """
+        request = Request(
+            method=Method.GET, endpoint=f"/_api/collection/{self.name}/count"
+        )
+
+        def response_handler(resp: Response) -> int:
+            if resp.is_success:
+                result: int = self.deserializer.loads(resp.raw_body)["count"]
+                return result
+            raise DocumentCountError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
     async def get(
         self,
         document: str | Json,

arangoasync/exceptions.py

@@ -187,6 +187,10 @@ class ClientConnectionError(ArangoClientError):
     """The request was unable to reach the server."""
 
 
+class CollectionTruncateError(ArangoServerError):
+    """Failed to truncate collection."""
+
+
 class CursorCloseError(ArangoServerError):
     """Failed to delete the cursor result from server."""
 
@@ -227,6 +231,10 @@ class DeserializationError(ArangoClientError):
     """Failed to deserialize the server response."""
 
 
+class DocumentCountError(ArangoServerError):
+    """Failed to retrieve document count."""
+
+
 class DocumentDeleteError(ArangoServerError):
     """Failed to delete document."""
 

tests/test_collection.py

@@ -5,6 +5,8 @@
 from arangoasync.errno import DATA_SOURCE_NOT_FOUND, INDEX_NOT_FOUND
 from arangoasync.exceptions import (
     CollectionPropertiesError,
+    CollectionTruncateError,
+    DocumentCountError,
     IndexCreateError,
     IndexDeleteError,
     IndexGetError,
@@ -157,3 +159,26 @@ async def test_collection_index(doc_col, bad_col, cluster):
         await doc_col.delete_index(idx1.id)
     assert err.value.error_code == INDEX_NOT_FOUND
     assert await doc_col.delete_index(idx2.id, ignore_missing=True) is False
+
+
+@pytest.mark.asyncio
+async def test_collection_truncate_count(docs, doc_col, bad_col):
+    # Test errors
+    with pytest.raises(CollectionTruncateError):
+        await bad_col.truncate()
+    with pytest.raises(DocumentCountError):
+        await bad_col.count()
+
+    # Test regular operations
+    await asyncio.gather(*[doc_col.insert(doc) for doc in docs])
+    cnt = await doc_col.count()
+    assert cnt == len(docs)
+
+    await doc_col.truncate()
+    cnt = await doc_col.count()
+    assert cnt == 0
+
+    await asyncio.gather(*[doc_col.insert(doc) for doc in docs])
+    await doc_col.truncate(wait_for_sync=True, compact=True)
+    cnt = await doc_col.count()
+    assert cnt == 0