graphdatascience 1.14__tar.gz → 1.15__tar.gz

{graphdatascience-1.14/graphdatascience.egg-info → graphdatascience-1.15}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.14
+Version: 1.15
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -51,6 +51,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: project-url
 Dynamic: provides-extra
 Dynamic: requires-dist

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/error/cypher_warning_handler.py

@@ -24,6 +24,12 @@ def filter_id_func_deprecation_warning() -> Callable[[F], F]:
                 message=r"^The query used a deprecated function. \('id' is no longer supported\)",
             )
 
+            # since 2025.04.0
+            warnings.filterwarnings(
+                "ignore",
+                message=r"^The query used a deprecated function. \('id' has been replaced by 'elementId or an application-generated id'\)",
+            )
+
             return func(self, *args, **kwargs)
 
         return cast(F, wrapper)

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/graph_data_science.py

@@ -7,6 +7,8 @@ from typing import Any, Optional, Type, Union
 from neo4j import Driver
 from pandas import DataFrame
 
+from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
+
 from .call_builder import IndirectCallBuilder
 from .endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
 from .error.uncallable_namespace import UncallableNamespace
@@ -93,10 +95,15 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
 
         arrow_info = ArrowInfo.create(self._query_runner)
         if arrow and arrow_info.enabled and self._server_version >= ServerVersion(2, 1, 0):
+            arrow_auth = None
+            if auth is not None:
+                username, password = auth
+                arrow_auth = UsernamePasswordAuthentication(username, password)
+
             self._query_runner = ArrowQueryRunner.create(
                 self._query_runner,
                 arrow_info,
-                auth,
+                arrow_auth,
                 self._query_runner.encrypted(),
                 arrow_disable_server_verification,
                 arrow_tls_root_certs,

graphdatascience-1.15/graphdatascience/query_runner/arrow_authentication.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+class ArrowAuthentication(ABC):
+    @abstractmethod
+    def auth_pair(self) -> tuple[str, str]:
+        """Returns the auth pair used for authentication."""
+        pass
+
+
+@dataclass
+class UsernamePasswordAuthentication(ArrowAuthentication):
+    username: str
+    password: str
+
+    def auth_pair(self) -> tuple[str, str]:
+        return self.username, self.password

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/arrow_graph_constructor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import concurrent
+import logging
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
@@ -33,6 +34,7 @@ class ArrowGraphConstructor(GraphConstructor):
         )
         self._chunk_size = chunk_size
         self._min_batch_size = chunk_size * 10
+        self._logger = logging.getLogger()
 
     def run(self, node_dfs: list[DataFrame], relationship_dfs: list[DataFrame]) -> None:
         try:
@@ -60,8 +62,11 @@ class ArrowGraphConstructor(GraphConstructor):
 
             self._client.relationship_load_done(self._graph_name)
         except (Exception, KeyboardInterrupt) as e:
-            self._client.abort(self._graph_name)
-
+            try:
+                self._client.abort(self._graph_name)
+            except Exception as abort_exception:
+                if "No arrow process" not in str(abort_exception):
+                    self._logger.warning(f"error aborting graph creation: {abort_exception}")
             raise e
 
     def _partition_dfs(self, dfs: list[DataFrame]) -> list[DataFrame]:

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/arrow_query_runner.py

@@ -5,6 +5,9 @@ from typing import Any, Optional
 
 from pandas import DataFrame
 
+from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
+from graphdatascience.retry_utils.retry_config import RetryConfig
+
 from ..call_parameters import CallParameters
 from ..query_runner.arrow_info import ArrowInfo
 from ..server_version.server_version import ServerVersion
@@ -19,22 +22,24 @@ class ArrowQueryRunner(QueryRunner):
     def create(
         fallback_query_runner: QueryRunner,
         arrow_info: ArrowInfo,
-        auth: Optional[tuple[str, str]] = None,
+        arrow_authentication: Optional[ArrowAuthentication] = None,
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
+        retry_config: Optional[RetryConfig] = None,
     ) -> ArrowQueryRunner:
         if not arrow_info.enabled:
             raise ValueError("Arrow is not enabled on the server")
 
         gds_arrow_client = GdsArrowClient.create(
             arrow_info,
-            auth,
+            arrow_authentication,
             encrypted,
             disable_server_verification,
             tls_root_certs,
             connection_string_override,
+            retry_config=retry_config,
         )
 
         return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/gds_arrow_client.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import base64
 import json
+import logging
 import re
 import time
 import warnings
@@ -12,11 +13,31 @@ from typing import Any, Callable, Dict, Iterable, Optional, Type, Union
 import pandas
 import pyarrow
 from neo4j.exceptions import ClientError
-from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Table, chunked_array, flight
+from pyarrow import Array, ChunkedArray, DictionaryArray, RecordBatch, Schema, Table, chunked_array, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import (
+    ClientMiddleware,
+    ClientMiddlewareFactory,
+    FlightDescriptor,
+    FlightInternalError,
+    FlightMetadataReader,
+    FlightStreamWriter,
+    FlightTimedOutError,
+    FlightUnavailableError,
+)
 from pyarrow.types import is_dictionary
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
+from tenacity import (
+    retry,
+    retry_any,
+    retry_if_exception_type,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_exponential,
+)
+
+from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication, UsernamePasswordAuthentication
+from graphdatascience.retry_utils.retry_config import RetryConfig
+from graphdatascience.retry_utils.retry_utils import before_log
 
 from ..semantic_version.semantic_version import SemanticVersion
 from ..version import __version__
@@ -28,11 +49,12 @@ class GdsArrowClient:
     @staticmethod
     def create(
         arrow_info: ArrowInfo,
-        auth: Optional[tuple[str, str]] = None,
+        auth: Optional[Union[ArrowAuthentication, tuple[str, str]]] = None,
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
+        retry_config: Optional[RetryConfig] = None,
     ) -> GdsArrowClient:
         connection_string: str
         if connection_string_override is not None:
@@ -44,8 +66,20 @@ class GdsArrowClient:
 
         arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.versions)
 
+        if retry_config is None:
+            retry_config = RetryConfig(
+                retry=retry_any(
+                    retry_if_exception_type(FlightTimedOutError),
+                    retry_if_exception_type(FlightUnavailableError),
+                    retry_if_exception_type(FlightInternalError),
+                ),
+                stop=(stop_after_delay(10) | stop_after_attempt(5)),
+                wait=wait_exponential(multiplier=1, min=1, max=10),
+            )
+
         return GdsArrowClient(
             host,
+            retry_config,
             int(port),
             auth,
             encrypted,
@@ -57,8 +91,9 @@ class GdsArrowClient:
     def __init__(
         self,
         host: str,
+        retry_config: RetryConfig,
         port: int = 8491,
-        auth: Optional[tuple[str, str]] = None,
+        auth: Optional[Union[ArrowAuthentication, tuple[str, str]]] = None,
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
@@ -73,8 +108,8 @@ class GdsArrowClient:
             The host address of the GDS Arrow server
         port: int
             The host port of the GDS Arrow server (default is 8491)
-        auth: Optional[tuple[str, str]]
-            A tuple containing the username and password for authentication
+        auth: Optional[Union[ArrowAuthentication, tuple[str, str]]]
+            Either an implementation of ArrowAuthentication providing a pair to be used for basic authentication, or a username, password tuple
         encrypted: bool
             A flag that indicates whether the connection should be encrypted (default is False)
         disable_server_verification: bool
@@ -85,17 +120,25 @@ class GdsArrowClient:
             The version of the Arrow endpoint to use (default is ArrowEndpointVersion.V1)
         user_agent: Optional[str]
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
+        retry_config: Optional[RetryConfig]
+            The retry configuration to use for the Arrow requests send by the client.
         """
         self._arrow_endpoint_version = arrow_endpoint_version
         self._host = host
         self._port = port
-        self._auth = auth
+        self._auth = None
         self._encrypted = encrypted
         self._disable_server_verification = disable_server_verification
         self._tls_root_certs = tls_root_certs
         self._user_agent = user_agent
+        self._retry_config = retry_config
+        self._logger = logging.getLogger("gds_arrow_client")
 
         if auth:
+            if not isinstance(auth, ArrowAuthentication):
+                username, password = auth
+                auth = UsernamePasswordAuthentication(username, password)
+            self._auth = auth
             self._auth_middleware = AuthMiddleware(auth)
 
         self._flight_client = self._instantiate_flight_client()
@@ -140,9 +183,22 @@ class GdsArrowClient:
         Optional[str]
             a token from the server and returns it.
         """
-        if self._auth:
+
+        @retry(
+            reraise=True,
+            before=before_log("Request token", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def auth_with_retry() -> None:
             client = self._client()
-            client.authenticate_basic_token(self._auth[0], self._auth[1])
+            if self._auth:
+                auth_pair = self._auth.auth_pair()
+                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+
+        if self._auth:
+            auth_with_retry()
             return self._auth_middleware.token()
         else:
             return "IGNORED"
@@ -193,7 +249,7 @@ class GdsArrowClient:
         if node_labels:
             config["node_labels"] = node_labels
 
-        return self._do_get(database, graph_name, proc, concurrency, config)
+        return self._do_get_with_retry(database, graph_name, proc, concurrency, config)
 
     def get_node_labels(self, graph_name: str, database: str, concurrency: Optional[int] = None) -> pandas.DataFrame:
         """
@@ -213,7 +269,7 @@ class GdsArrowClient:
         DataFrame
             The requested nodes as a DataFrame
         """
-        return self._do_get(database, graph_name, "gds.graph.nodeLabels.stream", concurrency, {})
+        return self._do_get_with_retry(database, graph_name, "gds.graph.nodeLabels.stream", concurrency, {})
 
     def get_relationships(
         self, graph_name: str, database: str, relationship_types: list[str], concurrency: Optional[int] = None
@@ -237,7 +293,7 @@ class GdsArrowClient:
         DataFrame
             The requested relationships as a DataFrame
         """
-        return self._do_get(
+        return self._do_get_with_retry(
             database,
             graph_name,
             "gds.graph.relationships.stream",
@@ -285,7 +341,7 @@ class GdsArrowClient:
         if relationship_types:
             config["relationship_types"] = relationship_types
 
-        return self._do_get(database, graph_name, proc, concurrency, config)
+        return self._do_get_with_retry(database, graph_name, proc, concurrency, config)
 
     def create_graph(
         self,
@@ -573,19 +629,32 @@ class GdsArrowClient:
 
     def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str, Any]:
         action_type = self._versioned_action_type(action_type)
+        client = self._client()
 
-        try:
-            client = self._client()
-            result = client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
+        @retry(
+            reraise=True,
+            before=before_log("Send action", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def send_with_retry() -> dict[str, Any]:
+            try:
+                result = client.do_action(
+                    action=flight.Action(action_type, json.dumps(meta_data).encode("utf-8")),
+                    options=flight.FlightCallOptions(timeout=20.0),
+                )
 
-            # Consume result fully to sanity check and avoid cancelled streams
-            collected_result = list(result)
-            assert len(collected_result) == 1
+                # Consume result fully to sanity check and avoid cancelled streams
+                collected_result = list(result)
+                assert len(collected_result) == 1
 
-            return json.loads(collected_result[0].body.to_pybytes().decode())  # type: ignore
-        except Exception as e:
-            self.handle_flight_error(e)
-            raise e  # unreachable
+                return json.loads(collected_result[0].body.to_pybytes().decode())  # type: ignore
+            except Exception as e:
+                self.handle_flight_error(e)
+                raise e  # unreachable
+
+        return send_with_retry()
 
     def _upload_data(
         self,
@@ -605,17 +674,26 @@ class GdsArrowClient:
         flight_descriptor = self._versioned_flight_descriptor({"name": graph_name, "entity_type": entity_type})
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
 
-        client = self._client()
-        put_stream, ack_stream = client.do_put(upload_descriptor, batches[0].schema)
+        @retry(
+            reraise=True,
+            before=before_log("Do put", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def safe_do_put(
+            upload_descriptor: FlightDescriptor, schema: Schema
+        ) -> tuple[FlightStreamWriter, FlightMetadataReader]:
+            return self._client().do_put(upload_descriptor, schema)  # type: ignore
+
+        put_stream, ack_stream = safe_do_put(upload_descriptor, batches[0].schema)
 
         @retry(
-            stop=(stop_after_delay(10) | stop_after_attempt(5)),
-            wait=wait_exponential(multiplier=1, min=1, max=10),
-            retry=(
-                retry_if_exception_type(flight.FlightUnavailableError)
-                | retry_if_exception_type(flight.FlightTimedOutError)
-                | retry_if_exception_type(flight.FlightInternalError)
-            ),
+            reraise=True,
+            before=before_log("Upload batch", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
         )
         def upload_batch(p: RecordBatch) -> None:
             put_stream.write_batch(p)
@@ -629,6 +707,26 @@ class GdsArrowClient:
         except Exception as e:
             GdsArrowClient.handle_flight_error(e)
 
+    def _do_get_with_retry(
+        self,
+        database: str,
+        graph_name: str,
+        procedure_name: str,
+        concurrency: Optional[int],
+        configuration: dict[str, Any],
+    ) -> pandas.DataFrame:
+        @retry(
+            reraise=True,
+            before=before_log("Do get", self._logger, logging.DEBUG),
+            retry=self._retry_config.retry,
+            stop=self._retry_config.stop,
+            wait=self._retry_config.wait,
+        )
+        def safe_do_get() -> pandas.DataFrame:
+            return self._do_get(database, graph_name, procedure_name, concurrency, configuration)
+
+        return safe_do_get()
+
     def _do_get(
         self,
         database: str,
@@ -795,7 +893,7 @@ class AuthFactory(ClientMiddlewareFactory):  # type: ignore
 
 
 class AuthMiddleware(ClientMiddleware):  # type: ignore
-    def __init__(self, auth: tuple[str, str], *args: Any, **kwargs: Any) -> None:
+    def __init__(self, auth: ArrowAuthentication, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._auth = auth
         self._token: Optional[str] = None
@@ -829,15 +927,15 @@ class AuthMiddleware(ClientMiddleware):  # type: ignore
 
     def sending_headers(self) -> dict[str, str]:
         token = self.token()
-        if not token:
-            username, password = self._auth
-            auth_token = f"{username}:{password}"
-            auth_token = "Basic " + base64.b64encode(auth_token.encode("utf-8")).decode("ASCII")
-            # There seems to be a bug, `authorization` must be lower key
-            return {"authorization": auth_token}
-        else:
+        if token is not None:
             return {"authorization": "Bearer " + token}
 
+        auth_pair = self._auth.auth_pair()
+        auth_token = f"{auth_pair[0]}:{auth_pair[1]}"
+        auth_token = "Basic " + base64.b64encode(auth_token.encode("utf-8")).decode("ASCII")
+        # There seems to be a bug, `authorization` must be lower key
+        return {"authorization": auth_token}
+
 
 @dataclass(repr=True, frozen=True)
 class NodeLoadDoneResult:

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/neo4j_query_runner.py

@@ -57,19 +57,9 @@ class Neo4jQueryRunner(QueryRunner):
             query_runner = Neo4jQueryRunner(
                 endpoint, auto_close=False, bookmarks=bookmarks, database=database, show_progress=show_progress
             )
-
         else:
             raise ValueError(f"Invalid endpoint type: {type(endpoint)}")
 
-        if Neo4jQueryRunner._NEO4J_DRIVER_VERSION >= SemanticVersion(5, 21, 0):
-            notifications_logger = logging.getLogger("neo4j.notifications")
-            # the client does not expose YIELD fields so we just skip these warnings for now
-            notifications_logger.addFilter(
-                lambda record: (
-                    "The query used a deprecated field from a procedure" in record.msg and "by 'gds." in record.msg
-                )
-            )
-
         return query_runner
 
     @staticmethod
@@ -90,19 +80,12 @@ class Neo4jQueryRunner(QueryRunner):
             show_progress=show_progress,
             bookmarks=None,
             config=driver_config,
+            # we need to explicitly set this as the default value is None
+            # database in the session is always neo4j
             database="neo4j",
             instance_description="GDS Session",
         )
 
-        if Neo4jQueryRunner._NEO4J_DRIVER_VERSION >= SemanticVersion(5, 21, 0):
-            notifications_logger = logging.getLogger("neo4j.notifications")
-            # the client does not expose YIELD fields so we just skip these warnings for now
-            notifications_logger.addFilter(
-                lambda record: (
-                    "The query used a deprecated field from a procedure" in record.msg and "by 'gds." in record.msg
-                )
-            )
-
         return query_runner
 
     @staticmethod
@@ -167,12 +150,7 @@ class Neo4jQueryRunner(QueryRunner):
                 else:
                     raise e
 
-            # Though pandas support may be experimental in the `neo4j` package, it should always
-            # be supported in the `graphdatascience` package.
-            warnings.filterwarnings(
-                "ignore",
-                message=r"^pandas support is experimental and might be changed or removed in future versions$",
-            )
+            self.__configure_warnings_filter()
 
             df = result.to_df()
 
@@ -181,19 +159,10 @@ class Neo4jQueryRunner(QueryRunner):
             else:
                 self._last_bookmarks = session.last_bookmarks()
 
-            if (
-                Neo4jQueryRunner._NEO4J_DRIVER_VERSION >= SemanticVersion(5, 21, 0)
-                and result._warn_notification_severity == "WARNING"
-            ):
-                # the client does not expose YIELD fields so we just skip these warnings for now
-                warnings.filterwarnings(
-                    "ignore", message=r".*The query used a deprecated field from a procedure\. .* by 'gds.* "
-                )
-            else:
-                notifications = result.consume().notifications
-                if notifications:
-                    for notification in notifications:
-                        self._forward_cypher_warnings(notification)
+            notifications = result.consume().notifications
+            if notifications:
+                for notification in notifications:
+                    self._forward_cypher_warnings(notification)
 
             return df
 
@@ -266,7 +235,7 @@ class Neo4jQueryRunner(QueryRunner):
         # (see https://neo4j.com/docs/status-codes/current/notifications/ for more details)
         severity = notification["severity"]
         if severity == "WARNING":
-            if "query used a deprecated field from a procedure" in notification["description"]:
+            if "deprecated field" in notification["description"] and "procedure" in notification["description"]:
                 # the client does not expose YIELD fields so we just skip these warnings for now
                 return
 
@@ -373,6 +342,27 @@ class Neo4jQueryRunner(QueryRunner):
         if retrys == retry_config.max_retries:
             raise UnableToConnectError(f"Unable to connect to the {self._instance_description}") from exception
 
+    def __configure_warnings_filter(self) -> None:
+        if Neo4jQueryRunner._NEO4J_DRIVER_VERSION >= SemanticVersion(5, 21, 0):
+            notifications_logger = logging.getLogger("neo4j.notifications")
+            # the client does not expose YIELD fields so we just skip these warnings for now
+            notifications_logger.addFilter(
+                lambda record: (
+                    "The query used a deprecated field from a procedure" in record.msg and "by 'gds." in record.msg
+                )
+            )
+            notifications_logger.addFilter(
+                lambda record: "The procedure has a deprecated field" in record.msg and "gds." in record.msg
+            )
+        warnings.filterwarnings(
+            "ignore",
+            message=r"^pandas support is experimental and might be changed or removed in future versions$",
+        )
+        # neo4j 2025.04
+        warnings.filterwarnings("ignore", message=r".*The procedure has a deprecated field.*by 'gds.*")
+        # neo4j driver 4.4
+        warnings.filterwarnings("ignore", message=r".*The query used a deprecated field from a procedure.*by 'gds.*")
+
     class ConnectivityRetriesConfig(NamedTuple):
         max_retries: int = 600
         wait_time: int = 1

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/protocol/project_protocols.py

@@ -7,9 +7,9 @@ from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.retry_utils.retry_utils import before_log
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 
@@ -133,6 +133,7 @@ class ProjectProtocolV3(ProjectProtocol):
         logger = getLogger()
 
         @retry(
+            reraise=True,
             before=before_log(f"Projection (graph: `{params['graph_name']}`)", logger, DEBUG),
             retry=retry_if_result(is_not_done),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/protocol/write_protocols.py

@@ -7,9 +7,9 @@ from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.termination_flag import TerminationFlag
+from graphdatascience.retry_utils.retry_utils import before_log
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 
@@ -142,6 +142,7 @@ class RemoteWriteBackV3(WriteProtocol):
         logger = logging.getLogger()
 
         @retry(
+            reraise=True,
             retry=retry_if_result(is_not_completed),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
             before=before_log(

{graphdatascience-1.14 → graphdatascience-1.15}/graphdatascience/query_runner/session_query_runner.py

@@ -202,12 +202,13 @@ class SessionQueryRunner(QueryRunner):
             return write_protocol.run_write_back(self._db_query_runner, write_back_params, yields, terminationFlag)
 
         try:
-            if self._resolve_show_progress(logging):
-                database_write_result = self._progress_logger.run_with_progress_logging(
-                    run_write_back, job_id, database
-                )
-            else:
-                database_write_result = run_write_back()
+            # Skipping progress for now as export has a different jobId
+            # if self._resolve_show_progress(logging):
+            #     database_write_result = self._progress_logger.run_with_progress_logging(
+            #         run_write_back, job_id, database
+            #     )
+            # else:
+            database_write_result = run_write_back()
         except Exception as e:
             # catch the case nothing was needed to write-back (empty graph)
             # once we have the Arrow Endpoints V2, we could catch by first checking the jobs summary