Adding support for /responsibleShard

Author: apetenchea

arangoasync/collection.py

@@ -17,6 +17,7 @@
 )
 from arangoasync.exceptions import (
     CollectionPropertiesError,
+    CollectionResponsibleShardError,
     CollectionStatisticsError,
     CollectionTruncateError,
     DocumentCountError,
@@ -594,6 +595,38 @@ def response_handler(resp: Response) -> CollectionStatistics:
 
         return await self._executor.execute(request, response_handler)
 
+    async def responsible_shard(self, document: Json) -> Result[str]:
+        """Return the ID of the shard responsible for given document.
+
+        If the document does not exist, return the shard that would be
+        responsible.
+
+        Args:
+            document (dict): Document body with "_key" field.
+
+        Returns:
+            str: Shard ID.
+
+        Raises:
+            CollectionResponsibleShardError: If retrieval fails.
+
+        References:
+            - `get-the-responsible-shard-for-a-document <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-responsible-shard-for-a-document>`__
+        """  # noqa: E501
+        request = Request(
+            method=Method.PUT,
+            endpoint=f"/_api/collection/{self.name}/responsibleShard",
+            data=self.serializer.dumps(document),
+        )
+
+        def response_handler(resp: Response) -> str:
+            if resp.is_success:
+                body = self.deserializer.loads(resp.raw_body)
+                return cast(str, body["shardId"])
+            raise CollectionResponsibleShardError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
     async def has(
         self,
         document: str | Json,

arangoasync/exceptions.py

@@ -195,6 +195,10 @@ class CollectionPropertiesError(ArangoServerError):
     """Failed to retrieve collection properties."""
 
 
+class CollectionResponsibleShardError(ArangoServerError):
+    """Failed to retrieve responsible shard."""
+
+
 class CollectionStatisticsError(ArangoServerError):
     """Failed to retrieve collection statistics."""
 

tests/test_collection.py

@@ -5,6 +5,7 @@
 from arangoasync.errno import DATA_SOURCE_NOT_FOUND, INDEX_NOT_FOUND
 from arangoasync.exceptions import (
     CollectionPropertiesError,
+    CollectionResponsibleShardError,
     CollectionStatisticsError,
     CollectionTruncateError,
     DocumentCountError,
@@ -23,20 +24,29 @@ def test_collection_attributes(db, doc_col):
 
 
 @pytest.mark.asyncio
-async def test_collection_misc_methods(doc_col, bad_col):
+async def test_collection_misc_methods(doc_col, bad_col, docs):
     # Properties
     properties = await doc_col.properties()
     assert properties.name == doc_col.name
     assert properties.is_system is False
     assert len(properties.format()) > 1
     with pytest.raises(CollectionPropertiesError):
         await bad_col.properties()
+
+    # Statistics
     statistics = await doc_col.statistics()
     assert statistics.name == doc_col.name
     assert "figures" in statistics
     with pytest.raises(CollectionStatisticsError):
         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)
+
 
 @pytest.mark.asyncio
 async def test_collection_index(doc_col, bad_col, cluster):