diff --git a/packages/gooddata-flight-server/src/gooddata_flight_server/server/flight_rpc/server_methods.py b/packages/gooddata-flight-server/src/gooddata_flight_server/server/flight_rpc/server_methods.py index 958979732..eb06ab7ec 100644 --- a/packages/gooddata-flight-server/src/gooddata_flight_server/server/flight_rpc/server_methods.py +++ b/packages/gooddata-flight-server/src/gooddata_flight_server/server/flight_rpc/server_methods.py @@ -101,6 +101,8 @@ def do_get_task_result( f"While the result exists, it is of an unexpected type: {type(result).__name__} ", ).to_internal_error() + finalizer = FlightServerMethods.call_finalizer_middleware(context) + rlock, data = result.acquire_data() def _on_end(_: pyarrow.ArrowException | None) -> None: @@ -121,7 +123,6 @@ def _on_end(_: pyarrow.ArrowException | None) -> None: # log and sink these Exceptions - not much to do _LOGGER.error("do_get_close_failed", exc_info=True) - finalizer = FlightServerMethods.call_finalizer_middleware(context) finalizer.register_on_end(_on_end) if isinstance(data, pyarrow.Table): diff --git a/packages/gooddata-flight-server/src/gooddata_flight_server/tasks/thread_task_executor.py b/packages/gooddata-flight-server/src/gooddata_flight_server/tasks/thread_task_executor.py index bc2773278..6428fab3d 100644 --- a/packages/gooddata-flight-server/src/gooddata_flight_server/tasks/thread_task_executor.py +++ b/packages/gooddata-flight-server/src/gooddata_flight_server/tasks/thread_task_executor.py @@ -158,6 +158,7 @@ class _TaskExecution: "_result_future", "_lock", "_completed", + "_done", "_stats", ) @@ -189,6 +190,8 @@ def __init__( # all these are protected using the lock self._result_future: Future[Union[TaskResult, TaskError]] | None = None self._completed: threading.Condition = threading.Condition(self._lock) + # indicates the task actually finished + self._done: bool = False @property def task(self) -> Task: @@ -235,6 +238,7 @@ def on_result_done(self, fut: Future) -> None: with self._lock: execution_result = self._cb.process_task_result(self, self._result_future) + self._done = True self._completed.notify_all() self._complete_execution_span(execution_result) @@ -291,7 +295,7 @@ def cancel(self) -> bool: def wait_for_completion(self, timeout: float | None = None) -> None: with self._lock: - completed = self._completed.wait(timeout=timeout) + completed = self._completed.wait_for(lambda: self._done, timeout=timeout) if not completed: raise TaskWaitTimeoutError(task_id=self._task.task_id, cmd=self._task.cmd)