Vertex collection management

Author: apetenchea

arangoasync/exceptions.py

@@ -295,6 +295,10 @@ class GraphListError(ArangoServerError):
     """Failed to retrieve graphs."""
 
 
+class GraphPropertiesError(ArangoServerError):
+    """Failed to retrieve graph properties."""
+
+
 class IndexCreateError(ArangoServerError):
     """Failed to create collection index."""
 

arangoasync/graph.py

@@ -1,13 +1,15 @@
 __all__ = ["Graph"]
 
 
-from typing import Generic, Optional, Sequence, TypeVar
+from typing import Generic, List, Optional, Sequence, TypeVar
 
 from arangoasync.collection import EdgeCollection, VertexCollection
 from arangoasync.exceptions import (
     EdgeDefinitionCreateError,
-    GraphListError,
+    GraphPropertiesError,
     VertexCollectionCreateError,
+    VertexCollectionDeleteError,
+    VertexCollectionListError,
 )
 from arangoasync.executor import ApiExecutor
 from arangoasync.request import Method, Request
@@ -74,7 +76,7 @@ async def properties(self) -> Result[GraphProperties]:
             GraphProperties: Properties of the graph.
 
         Raises:
-            GraphListError: If the operation fails.
+            GraphProperties: If the operation fails.
 
         References:
             - `get-a-graph <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#get-a-graph>`__
@@ -83,7 +85,7 @@ async def properties(self) -> Result[GraphProperties]:
 
         def response_handler(resp: Response) -> GraphProperties:
             if not resp.is_success:
-                raise GraphListError(resp, request)
+                raise GraphPropertiesError(resp, request)
             body = self.deserializer.loads(resp.raw_body)
             return GraphProperties(body["graph"])
 
@@ -106,6 +108,56 @@ def vertex_collection(self, name: str) -> VertexCollection[T, U, V]:
             doc_deserializer=self._doc_deserializer,
         )
 
+    async def vertex_collections(self) -> Result[List[str]]:
+        """Get the names of all vertex collections in the graph.
+
+        Returns:
+            list: List of vertex collection names.
+
+        Raises:
+            VertexCollectionListError: If the operation fails.
+
+        References:
+            - `list-vertex-collections <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#list-vertex-collections>`__
+        """  # noqa: E501
+        request = Request(
+            method=Method.GET,
+            endpoint=f"/_api/gharial/{self._name}/vertex",
+        )
+
+        def response_handler(resp: Response) -> List[str]:
+            if not resp.is_success:
+                raise VertexCollectionListError(resp, request)
+            body = self.deserializer.loads(resp.raw_body)
+            return list(sorted(set(body["collections"])))
+
+        return await self._executor.execute(request, response_handler)
+
+    async def has_vertex_collection(self, name: str) -> Result[bool]:
+        """Check if the graph has the given vertex collection.
+
+        Args:
+            name (str): Vertex collection mame.
+
+        Returns:
+            bool: `True` if the graph has the vertex collection, `False` otherwise.
+
+        Raises:
+            VertexCollectionListError: If the operation fails.
+        """
+        request = Request(
+            method=Method.GET,
+            endpoint=f"/_api/gharial/{self._name}/vertex",
+        )
+
+        def response_handler(resp: Response) -> bool:
+            if not resp.is_success:
+                raise VertexCollectionListError(resp, request)
+            body = self.deserializer.loads(resp.raw_body)
+            return name in body["collections"]
+
+        return await self._executor.execute(request, response_handler)
+
     async def create_vertex_collection(
         self,
         name: str,
@@ -148,6 +200,34 @@ def response_handler(resp: Response) -> VertexCollection[T, U, V]:
 
         return await self._executor.execute(request, response_handler)
 
+    async def delete_vertex_collection(self, name: str, purge: bool = False) -> None:
+        """Remove a vertex collection from the graph.
+
+        Args:
+            name (str): Vertex collection name.
+            purge (bool): If set to `True`, the vertex collection is not just deleted
+                from the graph but also from the database completely. Note that you
+                cannot remove vertex collections that are used in one of the edge
+                definitions of the graph.
+
+        Raises:
+            VertexCollectionDeleteError: If the operation fails.
+
+        References:
+           - `remove-a-vertex-collection <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#remove-a-vertex-collection>`__
+        """  # noqa: E501
+        request = Request(
+            method=Method.DELETE,
+            endpoint=f"/_api/gharial/{self._name}/vertex/{name}",
+            params={"dropCollection": purge},
+        )
+
+        def response_handler(resp: Response) -> None:
+            if not resp.is_success:
+                raise VertexCollectionDeleteError(resp, request)
+
+        await self._executor.execute(request, response_handler)
+
     def edge_collection(self, name: str) -> EdgeCollection[T, U, V]:
         """Returns the edge collection API wrapper.
 

tests/conftest.py

@@ -8,7 +8,12 @@
 from arangoasync.auth import Auth, JwtToken
 from arangoasync.client import ArangoClient
 from arangoasync.typings import UserInfo
-from tests.helpers import generate_col_name, generate_db_name, generate_username
+from tests.helpers import (
+    generate_col_name,
+    generate_db_name,
+    generate_graph_name,
+    generate_username,
+)
 
 
 @dataclass
@@ -19,6 +24,7 @@ class GlobalData:
     secret: str = None
     token: JwtToken = None
     sys_db_name: str = "_system"
+    graph_name: str = "test_graph"
     username: str = generate_username()
     cluster: bool = False
     enterprise: bool = False
@@ -64,6 +70,7 @@ def pytest_configure(config):
     global_data.token = JwtToken.generate_token(global_data.secret)
     global_data.cluster = config.getoption("cluster")
     global_data.enterprise = config.getoption("enterprise")
+    global_data.graph_name = generate_graph_name()
 
     async def get_db_version():
         async with ArangoClient(hosts=global_data.url) as client:
@@ -215,6 +222,11 @@ async def bad_db(arango_client):
     )
 
 
+@pytest_asyncio.fixture
+def bad_graph(bad_db):
+    return bad_db.graph(global_data.graph_name)
+
+
 @pytest_asyncio.fixture
 async def doc_col(db):
     col_name = generate_col_name()

tests/test_graph.py

@@ -1,6 +1,14 @@
 import pytest
 
-from arangoasync.exceptions import GraphCreateError, GraphDeleteError, GraphListError
+from arangoasync.exceptions import (
+    GraphCreateError,
+    GraphDeleteError,
+    GraphListError,
+    GraphPropertiesError,
+    VertexCollectionCreateError,
+    VertexCollectionDeleteError,
+    VertexCollectionListError,
+)
 from arangoasync.typings import GraphOptions
 from tests.helpers import generate_col_name, generate_graph_name
 
@@ -41,13 +49,16 @@ async def test_graph_basic(db, bad_db):
         await bad_db.delete_graph(graph1_name)
 
 
-async def test_graph_properties(db, cluster, enterprise):
+async def test_graph_properties(db, bad_graph, cluster, enterprise):
     # Create a graph
     name = generate_graph_name()
     is_smart = cluster and enterprise
     options = GraphOptions(number_of_shards=3)
     graph = await db.create_graph(name, is_smart=is_smart, options=options)
 
+    with pytest.raises(GraphPropertiesError):
+        await bad_graph.properties()
+
     # Create first vertex collection
     vcol_name = generate_col_name()
     vcol = await graph.create_vertex_collection(vcol_name)
@@ -85,3 +96,33 @@ async def test_graph_properties(db, cluster, enterprise):
     assert properties.edge_definitions[0]["from"][0] == vcol_name
     assert len(properties.edge_definitions[0]["to"]) == 1
     assert properties.edge_definitions[0]["to"][0] == vcol2_name
+
+
+async def test_vertex_collections(db, bad_graph):
+    # Test errors
+    with pytest.raises(VertexCollectionCreateError):
+        await bad_graph.create_vertex_collection("bad_col")
+    with pytest.raises(VertexCollectionListError):
+        await bad_graph.vertex_collections()
+    with pytest.raises(VertexCollectionListError):
+        await bad_graph.has_vertex_collection("bad_col")
+    with pytest.raises(VertexCollectionDeleteError):
+        await bad_graph.delete_vertex_collection("bad_col")
+
+    # Create graph
+    graph = await db.create_graph(generate_graph_name())
+
+    # Create vertex collections
+    names = [generate_col_name() for _ in range(3)]
+    cols = [await graph.create_vertex_collection(name) for name in names]
+
+    # List vertex collection
+    col_list = await graph.vertex_collections()
+    assert len(col_list) == 3
+    for c in cols:
+        assert c.name in col_list
+        assert await graph.has_vertex_collection(c.name)
+
+    # Delete collections
+    await graph.delete_vertex_collection(names[0])
+    assert await graph.has_vertex_collection(names[0]) is False