Fix graph creation flags and consolidate option (#107)

* adding failing test

* improve GraphProperties

* Fix graph creation flags and add GraphOptions support while deprecating separate arguments

Author: apetenchea

arangoasync/database.py

@@ -1240,11 +1240,18 @@ async def create_graph(
                 edge collection), "from" (list of vertex collection names) and "to"
                 (list of vertex collection names).
             is_disjoint (bool | None): Whether to create a Disjoint SmartGraph
-                instead of a regular SmartGraph (Enterprise Edition only).
-            is_smart (bool | None): Define if the created graph should be smart
-                (Enterprise Edition only).
+                instead of a regular SmartGraph. Only has an effect with
+                `smart_graph_attribute` or `is_smart` (Enterprise Edition only).
+                Deprecated: use `GraphOptions.is_disjoint` instead.
+            is_smart (bool | None): If `True`, create a SmartGraph when
+                `options.smart_graph_attribute` is set, or an EnterpriseGraph
+                otherwise. Ignored for SatelliteGraphs (Enterprise Edition only).
+                Deprecated: use `GraphOptions.is_smart` instead.
             options (GraphOptions | dict | None): Options for creating collections
-                within this graph.
+                within this graph. Setting `smart_graph_attribute` creates a
+                SmartGraph even if `is_smart` is omitted or `False`. Explicit
+                `is_smart` and `is_disjoint` arguments override the corresponding
+                entries in these options.
             orphan_collections (list | None): An array of additional vertex
                 collections. Documents in these collections do not have edges
                 within this graph.
@@ -1260,22 +1267,31 @@ async def create_graph(
         References:
             - `create-a-graph <https://docs.arango.ai/arangodb/stable/develop/http-api/graphs/named-graphs/#create-a-graph>`__
         """  # noqa: E501
+        if is_smart is not None or is_disjoint is not None:
+            warn(
+                "The separate is_smart and is_disjoint arguments are deprecated; "
+                "pass these flags through GraphOptions or the options "
+                "dictionary instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         params: Params = {}
         if wait_for_sync is not None:
             params["waitForSync"] = wait_for_sync
 
         data: Json = {"name": name}
         if edge_definitions is not None:
             data["edgeDefinitions"] = edge_definitions
-        if is_disjoint is not None:
-            data["isDisjoint"] = is_disjoint
-        if is_smart is not None:
-            data["isSmart"] = is_smart
         if options is not None:
             if isinstance(options, GraphOptions):
-                data["options"] = options.to_dict()
+                data["options"] = options.to_dict().copy()
             else:
-                data["options"] = options
+                data["options"] = options.copy()
+        if is_disjoint is not None:
+            data.setdefault("options", {})["isDisjoint"] = is_disjoint
+        if is_smart is not None:
+            data.setdefault("options", {})["isSmart"] = is_smart
         if orphan_collections is not None:
             data["orphanCollections"] = orphan_collections
 

arangoasync/graph.py

@@ -87,10 +87,11 @@ async def properties(self) -> Result[GraphProperties]:
         """Get the properties of the graph.
 
         Returns:
-            GraphProperties: Properties of the graph.
+            GraphProperties: Properties of the graph, including `is_smart`,
+                `is_disjoint`, and `smart_graph_attribute`.
 
         Raises:
-            GraphProperties: If the operation fails.
+            GraphPropertiesError: If the operation fails.
 
         References:
             - `get-a-graph <https://docs.arango.ai/arangodb/stable/develop/http-api/graphs/named-graphs/#get-a-graph>`__

arangoasync/typings.py

@@ -1873,6 +1873,16 @@ def is_smart(self) -> bool:
         """Check if the graph is a smart graph."""
         return cast(bool, self._data.get("isSmart", False))
 
+    @property
+    def is_disjoint(self) -> bool:
+        """Check if the graph is a Disjoint SmartGraph."""
+        return cast(bool, self._data.get("isDisjoint", False))
+
+    @property
+    def smart_graph_attribute(self) -> Optional[str]:
+        """Return the sharding attribute for a SmartGraph, if present."""
+        return cast(Optional[str], self._data.get("smartGraphAttribute"))
+
     @property
     def is_satellite(self) -> bool:
         """Check if the graph is a satellite graph."""
@@ -1963,6 +1973,10 @@ class GraphOptions(JsonWrapper):
             Enterprise Edition.
         write_concern (int | None): The write concern for new collections in the
             graph.
+        is_smart (bool | None): Create a SmartGraph when `smart_graph_attribute`
+            is set, or an EnterpriseGraph otherwise. Ignored for SatelliteGraphs.
+        is_disjoint (bool | None): Create a Disjoint SmartGraph. Only has an
+            effect with `smart_graph_attribute` or `is_smart`.
 
     References:
         - `create-a-graph <https://docs.arango.ai/arangodb/stable/develop/http-api/graphs/named-graphs/#create-a-graph>`__
@@ -1975,6 +1989,8 @@ def __init__(
         satellites: Optional[List[str]] = None,
         smart_graph_attribute: Optional[str] = None,
         write_concern: Optional[int] = None,
+        is_smart: Optional[bool] = None,
+        is_disjoint: Optional[bool] = None,
     ) -> None:
         data: Json = dict()
         if number_of_shards is not None:
@@ -1987,8 +2003,20 @@ def __init__(
             data["smartGraphAttribute"] = smart_graph_attribute
         if write_concern is not None:
             data["writeConcern"] = write_concern
+        if is_smart is not None:
+            data["isSmart"] = is_smart
+        if is_disjoint is not None:
+            data["isDisjoint"] = is_disjoint
         super().__init__(data)
 
+    @property
+    def is_smart(self) -> Optional[bool]:
+        return cast(Optional[bool], self._data.get("isSmart"))
+
+    @property
+    def is_disjoint(self) -> Optional[bool]:
+        return cast(Optional[bool], self._data.get("isDisjoint"))
+
     @property
     def number_of_shards(self) -> Optional[int]:
         return cast(int, self._data.get("numberOfShards"))

docs/graph.rst

@@ -42,6 +42,45 @@ For more information, refer to `ArangoDB Manual`_.
         # Delete the graph.
         await db.delete_graph("school")
 
+Graph creation options
+======================
+
+Pass ``GraphOptions(is_smart=True)`` as ``options`` to create an EnterpriseGraph
+without a user-defined sharding attribute. Set ``GraphOptions.smart_graph_attribute``
+to create a SmartGraph; ``is_smart=True`` is then optional. Every vertex must
+contain the named attribute.
+
+Set ``is_disjoint=True`` in ``GraphOptions`` with a smart graph attribute to create
+a Disjoint SmartGraph, which rejects edges between vertices with different
+attribute values.
+For example, on an Enterprise cluster:
+
+.. code-block:: python
+
+    from arangoasync.typings import GraphOptions
+
+    regional = await db.create_graph(
+        "regional",
+        edge_definitions=[{
+            "collection": "routes",
+            "from": ["places"],
+            "to": ["places"],
+        }],
+        options=GraphOptions(
+            number_of_shards=3,
+            smart_graph_attribute="region",
+            is_disjoint=True,
+        ),
+    )
+    properties = await regional.properties()
+    assert properties.is_disjoint is True
+    assert properties.smart_graph_attribute == "region"
+
+The separate ``create_graph(is_smart=..., is_disjoint=...)`` arguments are
+deprecated. Pass these flags through ``GraphOptions`` or a raw ``options``
+dictionary instead. Existing separate arguments still take precedence over
+matching entries in ``options``, including explicit ``False`` values.
+
 .. _edge-definitions:
 
 Edge Definitions

tests/test_graph.py

@@ -1,7 +1,9 @@
 import pytest
 
+from arangoasync.errno import INVALID_DISJOINT_SMART_EDGE
 from arangoasync.exceptions import (
     DocumentDeleteError,
+    DocumentInsertError,
     EdgeCollectionListError,
     EdgeDefinitionDeleteError,
     EdgeDefinitionListError,
@@ -19,6 +21,80 @@
 from tests.helpers import generate_col_name, generate_graph_name
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("creation_api", ["options", "legacy"])
+@pytest.mark.parametrize(
+    "smart_field,disjoint",
+    [(None, False), ("region", False), ("region", True)],
+    ids=["enterprise", "smart", "disjoint-smart"],
+)
+async def test_graph_create_smart_options(
+    db, cluster, skip_tests, smart_field, disjoint, creation_api
+):
+    if "enterprise" in skip_tests:
+        pytest.skip("Only for ArangoDB enterprise edition")
+    if not cluster:
+        pytest.skip("Only for ArangoDB clusters")
+
+    # Create an EnterpriseGraph, SmartGraph, or Disjoint SmartGraph.
+    graph_name = generate_graph_name()
+    vertices = generate_col_name()
+    edges = generate_col_name()
+    options = GraphOptions(
+        number_of_shards=3,
+        smart_graph_attribute=smart_field,
+        is_smart=True,
+        is_disjoint=disjoint,
+    )
+    kwargs = {}
+    if creation_api == "legacy":
+        options = GraphOptions(
+            number_of_shards=3,
+            smart_graph_attribute=smart_field,
+            is_smart=False,
+            is_disjoint=not disjoint,
+        )
+        kwargs = {"is_smart": True, "is_disjoint": disjoint}
+    original_options = options.to_dict().copy()
+    graph = await db.create_graph(
+        graph_name,
+        edge_definitions=[{"collection": edges, "from": [vertices], "to": [vertices]}],
+        options=options,
+        **kwargs,
+    )
+    try:
+        assert options.to_dict() == original_options
+        # Check that the server honored the graph creation options.
+        properties = await graph.properties()
+        assert properties.is_smart is True
+        assert properties.is_disjoint is disjoint
+        assert (properties.smart_graph_attribute or "") == (smart_field or "")
+        assert properties.number_of_shards == 3
+
+        # Add two vertices in one region and a third in another.
+        first = await graph.insert_vertex(vertices, {"region": "eu"})
+        second = await graph.insert_vertex(vertices, {"region": "eu"})
+        third = await graph.insert_vertex(vertices, {"region": "us"})
+        # All three graph types must allow edges within a region.
+        edge = await graph.insert_edge(
+            edges, {"_from": first["_id"], "_to": second["_id"]}
+        )
+        assert await graph.has_edge(edge["_id"])
+
+        # Only the disjoint graph must reject edges between regions.
+        cross_region_edge = {"_from": first["_id"], "_to": third["_id"]}
+        if disjoint:
+            with pytest.raises(DocumentInsertError) as err:
+                await graph.insert_edge(edges, cross_region_edge)
+            assert err.value.error_code == INVALID_DISJOINT_SMART_EDGE
+        else:
+            edge = await graph.insert_edge(edges, cross_region_edge)
+            assert await graph.has_edge(edge["_id"])
+    finally:
+        # Remove the graph and its collections even if an assertion fails.
+        await db.delete_graph(graph_name, drop_collections=True)
+
+
 @pytest.mark.asyncio
 async def test_graph_basic(db, bad_db):
     graph1_name = generate_graph_name()
@@ -60,8 +136,8 @@ async def test_graph_properties(db, bad_graph, cluster, skip_tests):
     # Create a graph
     name = generate_graph_name()
     is_smart = cluster and "enterprise" not in skip_tests
-    options = GraphOptions(number_of_shards=3)
-    graph = await db.create_graph(name, is_smart=is_smart, options=options)
+    options = GraphOptions(number_of_shards=3, is_smart=is_smart)
+    graph = await db.create_graph(name, options=options)
 
     with pytest.raises(GraphPropertiesError):
         await bad_graph.properties()

tests/test_typings.py

@@ -358,6 +358,17 @@ def test_GraphProperties():
     ]
     assert graph_properties.orphan_collections == ["orphan1", "orphan2"]
 
+    for is_disjoint in (False, True):
+        properties = GraphProperties(
+            {"isDisjoint": is_disjoint, "smartGraphAttribute": "region"}
+        )
+        assert properties.is_disjoint is is_disjoint
+        assert properties.smart_graph_attribute == "region"
+
+    properties = GraphProperties({})
+    assert properties.is_disjoint is False
+    assert properties.smart_graph_attribute is None
+
 
 def test_GraphOptions():
     graph_options = GraphOptions(
@@ -374,6 +385,18 @@ def test_GraphOptions():
     assert graph_options.smart_graph_attribute == "region"
     assert graph_options.write_concern == 1
 
+    for flag in (False, True):
+        options = GraphOptions(is_smart=flag, is_disjoint=flag)
+        assert options.is_smart is flag
+        assert options.is_disjoint is flag
+        assert options.to_dict() == {"isSmart": flag, "isDisjoint": flag}
+
+    options = GraphOptions()
+    assert options.is_smart is None
+    assert options.is_disjoint is None
+    assert "isSmart" not in options
+    assert "isDisjoint" not in options
+
 
 def test_VertexCollectionOptions():
     options = VertexCollectionOptions(