Implement graph.drop for arrow endpoints

Author: DarthMax

graphdatascience/procedure_surface/api/catalog_endpoints.py

@@ -26,6 +26,19 @@ def list(self, G: Optional[Union[Graph, str]] = None) -> List[GraphListResult]:
         """
         pass
 
+    @abstractmethod
+    def drop(self, G: Union[Graph, str], fail_if_missing: Optional[bool] = None) -> GraphListResult:
+        """Drop a graph from the graph catalog.
+
+        Args:
+            G (Union[Graph, str]): Graph object or name to drop.
+            fail_if_missing (Optional[bool], optional): Whether to fail if the graph is missing. Defaults to None.
+
+        Returns:
+              GraphListResult: Graph metadata object containing information like
+                               graph name, node count, relationship count, etc.
+        """
+
     @abstractmethod
     def filter(
         self,
@@ -68,7 +81,7 @@ class GraphListResult(BaseModel):
     modification_time: datetime
     graph_schema: dict[str, Any] = Field(alias="schema")
     schema_with_orientation: dict[str, Any]
-    degree_distribution: dict[str, Any]
+    degree_distribution: Optional[dict[str, Any]] = None
 
     @field_validator("creation_time", "modification_time", mode="before")
     @classmethod

graphdatascience/procedure_surface/arrow/catalog_arrow_endpoints.py

@@ -9,7 +9,7 @@
 
 from graphdatascience import Graph, QueryRunner
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
-from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize, deserialize_single
 from graphdatascience.arrow_client.v2.job_client import JobClient
 from graphdatascience.procedure_surface.api.catalog_endpoints import (
     CatalogEndpoints,
@@ -86,6 +86,12 @@ def project(
 
         return ProjectionResult(**JobClient.get_summary(self._arrow_client, job_id))
 
+    def drop(self, G: Union[Graph, str], fail_if_missing: Optional[bool] = None) -> GraphListResult:
+        graph_name = G if isinstance(G, str) else G.name()
+        config = ConfigConverter.convert_to_gds_config(graphName=graph_name, failIfMissing=fail_if_missing)
+        result = self._arrow_client.do_action_with_retry("v2/graph.drop", json.dumps(config).encode("utf-8"))
+        return GraphListResult(**deserialize_single(result))
+
     def filter(
         self,
         G: Graph,

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_catalog_arrow_endpoints.py

@@ -3,6 +3,7 @@
 from typing import Generator
 
 import pytest
+from pyarrow import ArrowKeyError
 
 from graphdatascience import Graph, QueryRunner
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
@@ -29,7 +30,6 @@ def catalog_endpoints(arrow_client: AuthenticatedArrowClient) -> Generator[Catal
 
 
 def test_list_with_graph(catalog_endpoints: CatalogArrowEndpoints, sample_graph: Graph) -> None:
-    """Test listing graphs with a specific graph."""
     results = catalog_endpoints.list(G=sample_graph)
 
     assert len(results) == 1
@@ -46,7 +46,7 @@ def test_list_with_graph(catalog_endpoints: CatalogArrowEndpoints, sample_graph:
     assert "KiB" in result.memory_usage
     assert result.size_in_bytes > 20000
     assert result.modification_time < datetime.datetime.now(datetime.timezone.utc)
-    assert "p50" in result.degree_distribution
+    assert "p50" in result.degree_distribution  # type: ignore
 
 
 def test_list_without_graph(
@@ -60,6 +60,18 @@ def test_list_without_graph(
     assert set(g.graph_name for g in result) == {sample_graph.name(), g2.name()}
 
 
+def test_drop(catalog_endpoints: CatalogArrowEndpoints, sample_graph: Graph) -> None:
+    res = catalog_endpoints.drop(sample_graph)
+
+    assert res.graph_name == sample_graph.name()
+    assert len(catalog_endpoints.list()) == 0
+
+
+def test_drop_nonexistent(catalog_endpoints: CatalogArrowEndpoints) -> None:
+    with pytest.raises(ArrowKeyError, match="does not exist on database"):
+        catalog_endpoints.drop("nonexistent", fail_if_missing=True)
+
+
 def test_projection(arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner) -> None:
     try:
         endpoints = CatalogArrowEndpoints(arrow_client, query_runner)
@@ -75,7 +87,7 @@ def test_projection(arrow_client: AuthenticatedArrowClient, query_runner: QueryR
 
         assert len(endpoints.list("g")) == 1
     finally:
-        arrow_client.do_action("v2/graph.drop", json.dumps({"graphName": "g"}).encode("utf-8"))
+        endpoints.drop("g")
 
 
 def test_graph_filter(catalog_endpoints: CatalogArrowEndpoints, sample_graph: Graph) -> None:
@@ -90,6 +102,4 @@ def test_graph_filter(catalog_endpoints: CatalogArrowEndpoints, sample_graph: Gr
         assert result.graph_name == "filtered"
         assert result.project_millis >= 0
     finally:
-        catalog_endpoints._arrow_client.do_action(
-            "v2/graph.drop", json.dumps({"graphName": "filtered"}).encode("utf-8")
-        )
+        catalog_endpoints.drop("filtered")