Add safeguard to raise exception on termination

CI showed the closing the put-stream does not always throw an exception

Author: FlorentinD

src/graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -506,8 +506,9 @@ def upload_batch(p: RecordBatch) -> None:
         with put_stream:
             for partition in batches:
                 if termination_flag is not None and termination_flag.is_set():
-                    self.abort_job(job_id)  # closing the put_stream will raise an error
-                    break
+                    self.abort_job(job_id)
+                    # closing the put_stream should raise an error. this is a safeguard to always signal the termination to the user.
+                    raise RuntimeError(f"Upload for job '{job_id}' was aborted via termination flag.")
 
                 upload_batch(partition)
                 ack_stream.read()

tests/integrationV2/arrow_client/v2/test_gds_arrow_client_v2.py

@@ -4,7 +4,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-from pyarrow.flight import FlightCancelledError
 from testcontainers.core.network import Network
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
@@ -107,9 +106,7 @@ def test_project_from_triplets(arrow_client: AuthenticatedArrowClient, gds_arrow
         assert G.name() == graph_name
 
 
-def test_project_from_triplets_interrupted(
-    arrow_client: AuthenticatedArrowClient, gds_arrow_client: GdsArrowClient
-) -> None:
+def test_project_from_triplets_interrupted(gds_arrow_client: GdsArrowClient) -> None:
     df = pd.DataFrame(
         {"sourceNode": np.array([1, 2, 3], dtype=np.int64), "targetNode": np.array([4, 5, 6], dtype=np.int64)}
     )
@@ -118,7 +115,7 @@ def test_project_from_triplets_interrupted(
     termination_flag.set()
 
     job_id = gds_arrow_client.create_graph_from_triplets("triplets")
-    with pytest.raises(FlightCancelledError, match=".*Arrow process 'triplets' was aborted.*"):
+    with pytest.raises(Exception, match=".*was aborted.*"):
         gds_arrow_client.upload_triplets(job_id, df, termination_flag=termination_flag)