Fix imports

Author: DarthMax

examples/graph-analytics-serverless-spark.ipynb

@@ -215,25 +215,28 @@
    ]
   },
   {
-   "metadata": {},
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "The user wants to add a 1-second delay (sleep) within the loop that waits for the import job to finish. This requires importing the `time` module and adding `time.sleep(1)` inside the `while` loop at the end of the cell.\n",
     "\n"
    ]
   },
   {
-   "metadata": {},
    "cell_type": "markdown",
-   "source": "<llm-snippet-file>graph-analytics-serverless-spark.ipynb</llm-snippet-file>\n"
+   "metadata": {},
+   "source": [
+    "<llm-snippet-file>graph-analytics-serverless-spark.ipynb</llm-snippet-file>\n"
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import time\n",
+    "\n",
     "import pandas as pd\n",
     "import pyarrow\n",
     "from pyspark.sql import functions\n",
@@ -280,8 +283,8 @@
    ]
   },
   {
-   "metadata": {},
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Running Algorithms\n",
     "\n",

src/graphdatascience/arrow_client/v1/gds_arrow_client.py

@@ -15,8 +15,8 @@
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
 from graphdatascience.arrow_client.v1.data_mapper_utils import deserialize_single
 
+from ...procedure_surface.arrow.error_handler import handle_flight_error
 from ...semantic_version.semantic_version import SemanticVersion
-from ..error_handler import handle_flight_error
 
 
 class GdsArrowClient:

src/graphdatascience/procedure_surface/arrow/error_handler.py

@@ -0,0 +1,24 @@
+import re
+
+from neo4j.exceptions import ClientError
+from pyarrow import flight
+
+
+def handle_flight_error(e: Exception) -> None:
+    if isinstance(e, flight.FlightServerError | flight.FlightInternalError | ClientError):
+        original_message = e.args[0] if len(e.args) > 0 else e.message
+        improved_message = original_message.replace(
+            "Flight RPC failed with message: org.apache.arrow.flight.FlightRuntimeException: ", ""
+        )
+        improved_message = improved_message.replace(
+            "Flight returned internal error, with message: org.apache.arrow.flight.FlightRuntimeException: ", ""
+        )
+        improved_message = improved_message.replace(
+            "Failed to invoke procedure `gds.arrow.project`: Caused by: org.apache.arrow.flight.FlightRuntimeException: ",
+            "",
+        )
+        improved_message = re.sub(r"(\. )?gRPC client debug context: .+$", "", improved_message)
+
+        raise flight.FlightServerError(improved_message)
+    else:
+        raise e

src/graphdatascience/procedure_surface/cypher/catalog_cypher_endpoints.py

@@ -42,7 +42,7 @@ def __init__(self, cypher_runner: Neo4jQueryRunner, arrow_client: GdsArrowClient
     def get(self, graph_name: str) -> GraphV2:
         if not self.list(graph_name):
             raise ValueError(f"A graph with name '{graph_name}' does not exist in the catalog.")
-        return get_graph(graph_name, self.cypher_runner)
+        return get_graph(graph_name, self._cypher_runner)
 
     def construct(
         self,

src/graphdatascience/query_runner/session_query_runner.py

@@ -13,8 +13,8 @@
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.server_version.server_version import ServerVersion
 
-from ..arrow_client.error_handler import handle_flight_error
 from ..call_parameters import CallParameters
+from ..procedure_surface.arrow.error_handler import handle_flight_error
 from ..session.dbms.protocol_resolver import ProtocolVersionResolver
 from .protocol.project_protocols import ProjectProtocol
 from .protocol.write_protocols import WriteProtocol

tests/unit/arrow_client/V1/test_gds_arrow_client.py

@@ -15,8 +15,8 @@
 )
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
-from graphdatascience.arrow_client.error_handler import handle_flight_error
 from graphdatascience.arrow_client.v1.gds_arrow_client import GdsArrowClient
+from graphdatascience.procedure_surface.arrow.error_handler import handle_flight_error
 from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.query_runner.arrow_info import ArrowInfo