Adding support for /shards

Author: apetenchea

arangoasync/collection.py

@@ -18,6 +18,7 @@
 from arangoasync.exceptions import (
     CollectionPropertiesError,
     CollectionResponsibleShardError,
+    CollectionShardsError,
     CollectionStatisticsError,
     CollectionTruncateError,
     DocumentCountError,
@@ -627,6 +628,41 @@ def response_handler(resp: Response) -> str:
 
         return await self._executor.execute(request, response_handler)
 
+    async def shards(self, details: Optional[bool] = None) -> Result[Json]:
+        """Return collection shards and properties.
+
+        Available only in a cluster setup.
+
+        Args:
+            details (bool | None): If set to `True`, include responsible
+                servers for these shards.
+
+        Returns:
+            dict: Collection shards and properties.
+
+        Raises:
+            CollectionShardsError: If retrieval fails.
+
+        References:
+            - `get-the-shard-ids-of-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-shard-ids-of-a-collection>`__
+        """  # noqa: E501
+        params: Params = {}
+        if details is not None:
+            params["details"] = details
+
+        request = Request(
+            method=Method.GET,
+            endpoint=f"/_api/collection/{self.name}/shards",
+            params=params,
+        )
+
+        def response_handler(resp: Response) -> Json:
+            if not resp.is_success:
+                raise CollectionShardsError(resp, request)
+            return Response.format_body(self.deserializer.loads(resp.raw_body))
+
+        return await self._executor.execute(request, response_handler)
+
     async def has(
         self,
         document: str | Json,

arangoasync/exceptions.py

@@ -199,6 +199,10 @@ class CollectionResponsibleShardError(ArangoServerError):
     """Failed to retrieve responsible shard."""
 
 
+class CollectionShardsError(ArangoServerError):
+    """Failed to retrieve collection shards."""
+
+
 class CollectionStatisticsError(ArangoServerError):
     """Failed to retrieve collection statistics."""
 

tests/test_collection.py

@@ -6,6 +6,7 @@
 from arangoasync.exceptions import (
     CollectionPropertiesError,
     CollectionResponsibleShardError,
+    CollectionShardsError,
     CollectionStatisticsError,
     CollectionTruncateError,
     DocumentCountError,
@@ -24,7 +25,7 @@ def test_collection_attributes(db, doc_col):
 
 
 @pytest.mark.asyncio
-async def test_collection_misc_methods(doc_col, bad_col, docs):
+async def test_collection_misc_methods(doc_col, bad_col, docs, cluster):
     # Properties
     properties = await doc_col.properties()
     assert properties.name == doc_col.name
@@ -41,11 +42,16 @@ async def test_collection_misc_methods(doc_col, bad_col, docs):
         await bad_col.statistics()
 
     # Shards
-    doc = await doc_col.insert(docs[0])
-    shard = await doc_col.responsible_shard(doc)
-    assert isinstance(shard, str)
-    with pytest.raises(CollectionResponsibleShardError):
-        await bad_col.responsible_shard(doc)
+    if cluster:
+        doc = await doc_col.insert(docs[0])
+        shard = await doc_col.responsible_shard(doc)
+        assert isinstance(shard, str)
+        with pytest.raises(CollectionResponsibleShardError):
+            await bad_col.responsible_shard(doc)
+        shards = await doc_col.shards(details=True)
+        assert isinstance(shards, dict)
+        with pytest.raises(CollectionShardsError):
+            await bad_col.shards()
 
 
 @pytest.mark.asyncio