Fix write-back progress bar for write-back to self-managed dbs

Author: FlorentinD

graphdatascience/query_runner/protocol/write_protocols.py

@@ -149,7 +149,7 @@ def run_write_back(
         terminationFlag: TerminationFlag,
     ) -> DataFrame:
         def is_not_completed(result: DataFrame) -> bool:
-            status: str = result.squeeze()["status"]
+            status: str = result.iloc[0]["status"]
             return status != Status.COMPLETED.name
 
         logger = logging.getLogger()
@@ -175,9 +175,12 @@ def write_fn(progress_bar: Optional[TqdmProgressBar]) -> DataFrame:
                 mode=QueryMode.WRITE,
                 custom_error=False,
             )
+            result_row = result.iloc[0].to_dict()
+            # for self-managed dbs the endpoint doesn't return the progress yet
+            progress = result_row.get("progress", 0.0) * 100
 
             if progress_bar:
-                progress_bar.update(status=result.squeeze()["status"], progress=result.squeeze()["progress"] * 100)
+                progress_bar.update(status=result_row["status"], progress=progress)
 
             return result
 

graphdatascience/tests/unit/query_runner/test_write_protocols.py

@@ -34,6 +34,30 @@ def test_write_back_v3_progress_logging() -> None:
         assert any(["Write-Back (graph: myGraph): 100%|##########| 100.0/100" in line for line in bar_output])
 
 
+def test_write_back_v3_progress_logging_without_bar() -> None:
+    # for self-managed dbs the endpoint doesnt return the progress yet
+    with StringIO() as pbarOutputStream:
+        qr = CollectingQueryRunner(ServerVersion(0, 0, 0))
+        qr.add__mock_result("gds.arrow.write.v3", DataFrame([{"status": Status.COMPLETED.name}]))
+
+        wp = RemoteWriteBackV3(progress_bar_options={"file": pbarOutputStream, "mininterval": 0})
+
+        wp.run_write_back(
+            query_runner=qr,
+            parameters=CallParameters(graphName="myGraph", jobId="myJob"),
+            log_progress=True,
+            terminationFlag=TerminationFlagNoop(),
+            yields=None,
+        )
+
+        bar_output = pbarOutputStream.getvalue().split("\r")
+
+        assert any(
+            ["Write-Back (graph: myGraph):   0%|          | 0.0/100 [00:00<?, ?%/s]" in line for line in bar_output]
+        ), bar_output
+        assert any(["Write-Back (graph: myGraph): 100%|##########| 100.0/100" in line for line in bar_output])
+
+
 def test_write_back_v3_progress_logging_aborted() -> None:
     with StringIO() as pbarOutputStream:
         qr = CollectingQueryRunner(ServerVersion(0, 0, 0))