Set final status to failed on exception

Author: FlorentinD

graphdatascience/query_runner/progress/query_progress_logger.py

@@ -81,7 +81,7 @@ def _log(
                     continue
 
         if pbar is not None:
-            self._finish_pbar(pbar)
+            self._finish_pbar(future, pbar)
         # TODO show as cancelled if future was interrupted
 
     def _init_pbar(self, task_with_progress: TaskWithProgress) -> tqdm:  # type: ignore
@@ -119,7 +119,11 @@ def _update_pbar(self, pbar: tqdm, task_with_progress: TaskWithProgress) -> None
         else:
             pbar.refresh()
 
-    def _finish_pbar(self, pbar: tqdm) -> None:  # type: ignore
+    def _finish_pbar(self, future: Future[Any], pbar: tqdm) -> None:  # type: ignore
+        if future.exception():
+            pbar.set_postfix_str("status: FAILED", refresh=True)
+            return
+
         if pbar.total is not None:
             pbar.update(pbar.total - pbar.n)
         pbar.set_postfix_str("status: FINISHED", refresh=True)

graphdatascience/tests/unit/query_runner/progress/test_query_progress_logger.py

@@ -1,5 +1,6 @@
 import re
 import time
+from concurrent.futures import ThreadPoolExecutor
 from io import StringIO
 from typing import Optional
 
@@ -127,7 +128,10 @@ def simple_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
             running_output,
         ), running_output
 
-        qpl._finish_pbar(pbar)
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(lambda: None)
+            qpl._finish_pbar(future, pbar)
+
         finished_output = pbarOutputStream.getvalue().split("\r")[-1]
         assert re.match(
             r"test task: 100%\|##########\| 100.0/100 \[00:00<00:00, \d+.\d+%/s, status: FINISHED\]", finished_output
@@ -149,7 +153,10 @@ def simple_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
 
         qpl._update_pbar(pbar, TaskWithProgress("test task", "n/a", "PENDING", ""))
         qpl._update_pbar(pbar, TaskWithProgress("test task", "", "RUNNING", "root 1/1::leaf"))
-        qpl._finish_pbar(pbar)
+
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(lambda: None)
+            qpl._finish_pbar(future, pbar)
 
         assert pbarOutputStream.getvalue().rstrip() == "".join(
             [
@@ -160,6 +167,34 @@ def simple_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
         )
 
 
+def test_progress_bar_with_failing_query() -> None:
+    def simple_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
+        raise NotImplementedError("Should not be called!")
+
+    def failing_runnable() -> DataFrame:
+        raise NotImplementedError("Should not be called!")
+
+    with StringIO() as pbarOutputStream:
+        qpl = QueryProgressLogger(
+            simple_run_cypher,
+            lambda: ServerVersion(3, 0, 0),
+            progress_bar_options={"file": pbarOutputStream, "mininterval": 100},
+        )
+
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(failing_runnable)
+
+            pbar = qpl._init_pbar(TaskWithProgress("test task", "n/a", "PENDING", ""))
+            qpl._finish_pbar(future, pbar)
+
+            assert pbarOutputStream.getvalue().rstrip() == "".join(
+                [
+                    "\rtest task [elapsed: 00:00 ]",
+                    "\rtest task [elapsed: 00:00 , status: FAILED]",
+                ]
+            )
+
+
 def test_uses_static_store() -> None:
     def fake_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
         return DataFrame([{"progress": "n/a", "taskName": "Test task", "status": "RUNNING"}])