Facilitating edge and vertex collection creation

Author: apetenchea

arangoasync/exceptions.py

@@ -263,6 +263,26 @@ class DocumentUpdateError(ArangoServerError):
     """Failed to update document."""
 
 
+class EdgeDefinitionListError(ArangoServerError):
+    """Failed to retrieve edge definitions."""
+
+
+class EdgeDefinitionCreateError(ArangoServerError):
+    """Failed to create edge definition."""
+
+
+class EdgeDefinitionReplaceError(ArangoServerError):
+    """Failed to replace edge definition."""
+
+
+class EdgeDefinitionDeleteError(ArangoServerError):
+    """Failed to delete edge definition."""
+
+
+class EdgeListError(ArangoServerError):
+    """Failed to retrieve edges coming in and out of a vertex."""
+
+
 class GraphCreateError(ArangoServerError):
     """Failed to create the graph."""
 
@@ -389,3 +409,15 @@ class UserReplaceError(ArangoServerError):
 
 class UserUpdateError(ArangoServerError):
     """Failed to update user."""
+
+
+class VertexCollectionCreateError(ArangoServerError):
+    """Failed to create vertex collection."""
+
+
+class VertexCollectionDeleteError(ArangoServerError):
+    """Failed to delete vertex collection."""
+
+
+class VertexCollectionListError(ArangoServerError):
+    """Failed to retrieve vertex collections."""

arangoasync/graph.py

@@ -1,16 +1,26 @@
 __all__ = ["Graph"]
 
 
-from typing import Generic, TypeVar
+from typing import Generic, Optional, Sequence, TypeVar
 
 from arangoasync.collection import EdgeCollection, VertexCollection
-from arangoasync.exceptions import GraphListError
+from arangoasync.exceptions import (
+    EdgeDefinitionCreateError,
+    GraphListError,
+    VertexCollectionCreateError,
+)
 from arangoasync.executor import ApiExecutor
 from arangoasync.request import Method, Request
 from arangoasync.response import Response
 from arangoasync.result import Result
 from arangoasync.serialization import Deserializer, Serializer
-from arangoasync.typings import GraphProperties, Json, Jsons
+from arangoasync.typings import (
+    EdgeDefinitionOptions,
+    GraphProperties,
+    Json,
+    Jsons,
+    VertexCollectionOptions,
+)
 
 T = TypeVar("T")  # Serializer type
 U = TypeVar("U")  # Deserializer loads
@@ -67,7 +77,7 @@ async def properties(self) -> Result[GraphProperties]:
             GraphListError: If the operation fails.
 
         References:
-            - `get-a-graph <https://docs.arangodb.com/3.12/develop/http-api/graphs/named-graphs/#get-a-graph>`__
+            - `get-a-graph <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#get-a-graph>`__
         """  # noqa: E501
         request = Request(method=Method.GET, endpoint=f"/_api/gharial/{self._name}")
 
@@ -96,6 +106,48 @@ def vertex_collection(self, name: str) -> VertexCollection[T, U, V]:
             doc_deserializer=self._doc_deserializer,
         )
 
+    async def create_vertex_collection(
+        self,
+        name: str,
+        options: Optional[VertexCollectionOptions | Json] = None,
+    ) -> Result[VertexCollection[T, U, V]]:
+        """Create a vertex collection in the graph.
+
+        Args:
+            name (str): Vertex collection name.
+            options (dict | VertexCollectionOptions | None): Extra options for
+                creating vertex collections.
+
+        Returns:
+            VertexCollection: Vertex collection API wrapper.
+
+        Raises:
+            VertexCollectionCreateError: If the operation fails.
+
+        References:
+           - `add-a-vertex-collection <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#add-a-vertex-collection>`__
+        """  # noqa: E501
+        data: Json = {"collection": name}
+
+        if options is not None:
+            if isinstance(options, VertexCollectionOptions):
+                data["options"] = options.to_dict()
+            else:
+                data["options"] = options
+
+        request = Request(
+            method=Method.POST,
+            endpoint=f"/_api/gharial/{self._name}/vertex",
+            data=self.serializer.dumps(data),
+        )
+
+        def response_handler(resp: Response) -> VertexCollection[T, U, V]:
+            if not resp.is_success:
+                raise VertexCollectionCreateError(resp, request)
+            return self.vertex_collection(name)
+
+        return await self._executor.execute(request, response_handler)
+
     def edge_collection(self, name: str) -> EdgeCollection[T, U, V]:
         """Returns the edge collection API wrapper.
 
@@ -112,3 +164,66 @@ def edge_collection(self, name: str) -> EdgeCollection[T, U, V]:
             doc_serializer=self._doc_serializer,
             doc_deserializer=self._doc_deserializer,
         )
+
+    async def create_edge_definition(
+        self,
+        edge_collection: str,
+        from_vertex_collections: Sequence[str],
+        to_vertex_collections: Sequence[str],
+        options: Optional[EdgeDefinitionOptions | Json] = None,
+    ) -> Result[EdgeCollection[T, U, V]]:
+        """Create an edge definition in the graph.
+
+        This edge definition has to contain a collection and an array of each from
+        and to vertex collections.
+
+        .. code-block:: python
+
+            {
+                "edge_collection": "edge_collection_name",
+                "from_vertex_collections": ["from_vertex_collection_name"],
+                "to_vertex_collections": ["to_vertex_collection_name"]
+            }
+
+        Args:
+            edge_collection (str): Edge collection name.
+            from_vertex_collections (list): List of vertex collections
+                that can be used as the "from" vertex in edges.
+            to_vertex_collections (list): List of vertex collections
+                that can be used as the "to" vertex in edges.
+            options (dict | EdgeDefinitionOptions | None): Extra options for
+                creating edge definitions.
+
+        Returns:
+            EdgeCollection: Edge collection API wrapper.
+
+        Raises:
+            EdgeDefinitionCreateError: If the operation fails.
+
+        References:
+            - `add-an-edge-definition <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#add-an-edge-definition>`__
+        """  # noqa: E501
+        data: Json = {
+            "collection": edge_collection,
+            "from": from_vertex_collections,
+            "to": to_vertex_collections,
+        }
+
+        if options is not None:
+            if isinstance(options, VertexCollectionOptions):
+                data["options"] = options.to_dict()
+            else:
+                data["options"] = options
+
+        request = Request(
+            method=Method.POST,
+            endpoint=f"/_api/gharial/{self._name}/edge",
+            data=self.serializer.dumps(data),
+        )
+
+        def response_handler(resp: Response) -> EdgeCollection[T, U, V]:
+            if not resp.is_success:
+                raise EdgeDefinitionCreateError(resp, request)
+            return self.edge_collection(edge_collection)
+
+        return await self._executor.execute(request, response_handler)

tests/test_graph.py

@@ -1,7 +1,8 @@
 import pytest
 
 from arangoasync.exceptions import GraphCreateError, GraphDeleteError, GraphListError
-from tests.helpers import generate_graph_name
+from arangoasync.typings import GraphOptions
+from tests.helpers import generate_col_name, generate_graph_name
 
 
 @pytest.mark.asyncio
@@ -40,11 +41,47 @@ async def test_graph_basic(db, bad_db):
         await bad_db.delete_graph(graph1_name)
 
 
-async def test_graph_properties(db):
+async def test_graph_properties(db, cluster, enterprise):
     # Create a graph
     name = generate_graph_name()
-    graph = await db.create_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)
+
+    # Create first vertex collection
+    vcol_name = generate_col_name()
+    vcol = await graph.create_vertex_collection(vcol_name)
+    assert vcol.name == vcol_name
 
     # Get the properties of the graph
     properties = await graph.properties()
     assert properties.name == name
+    assert properties.is_smart == is_smart
+    assert properties.number_of_shards == options.number_of_shards
+    assert properties.orphan_collections == [vcol_name]
+
+    # Create second vertex collection
+    vcol2_name = generate_col_name()
+    vcol2 = await graph.create_vertex_collection(vcol2_name)
+    assert vcol2.name == vcol2_name
+    properties = await graph.properties()
+    assert len(properties.orphan_collections) == 2
+
+    # Create an edge definition
+    edge_name = generate_col_name()
+    edge_col = await graph.create_edge_definition(
+        edge_name,
+        from_vertex_collections=[vcol_name],
+        to_vertex_collections=[vcol2_name],
+    )
+    assert edge_col.name == edge_name
+
+    # There should be no more orphan collections
+    properties = await graph.properties()
+    assert len(properties.orphan_collections) == 0
+    assert len(properties.edge_definitions) == 1
+    assert properties.edge_definitions[0]["collection"] == edge_name
+    assert len(properties.edge_definitions[0]["from"]) == 1
+    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