fix(managed): keep waiting on an unknown status, and name it

Reverses the status inversion from the previous commit, and takes the
belt-and-braces route on the run-readiness question rather than arguing
it a third time.

The inversion traded one failure mode for a worse one. Enumerating the
in-flight statuses does make a missing status surface in one round trip
instead of a timeout -- but it also means a status added upstream fails
every query the moment it appears, where the arrangement it replaced cost
one slow call. That asymmetry was not worth the diagnosis, and the actual
defect in the `interrupted` case was never the waiting: it was that the
timeout said only "did not finish", never which status it had been
waiting on. Both timeouts now carry it, in both clients, so the next
missing status costs one slow call and names itself in the error.

Separately, `_fetch_result_arrow` now waits out a result that reports
itself not ready. A run reports `succeeded` only once its result is saved
and ready, so this is unreachable, and the review was right that nothing
in this package demonstrated that. Tolerating it is cheap and settles the
question in code: the Arrow endpoint answers a result which is not ready
with a small refusal rather than with data, so waiting there costs one
tiny request -- which is the whole difference from waiting on the JSON
result body, and the reason this is not a return to what was removed.

Author: anoop-narang

CHANGELOG.md

@@ -33,6 +33,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Costs one extra round trip on a query that would have answered synchronously,
   in exchange for not transferring the result twice.
 
+  The Arrow fetch now also waits out a result that reports itself not ready, in
+  case that ordering ever stops holding. It should be unreachable, and it is
+  cheap to keep: that endpoint answers a result which is not ready with a small
+  refusal rather than with data, which is exactly what made waiting on the JSON
+  result body expensive and waiting here not.
+
 - fix(managed): recognise `interrupted`, and drop a run status the API never sends.
 
   Both `ManagedDatabaseClient` and `HotdataClient` treated `failed` and
@@ -48,11 +54,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   caller-raised transient error to terminal. `HotdataClient.execute_sql` now
   fails fast on it with the run's own message.
 
-  Both polls now enumerate the statuses that mean *still in flight* rather than
-  the ones that mean *finished*, for query runs and for results alike. Listing
-  the terminal side treats an unrecognised status as "keep waiting", which is
-  precisely how `interrupted` came to be waited out; listing the in-flight side
-  makes the next status the API adds fail on the first pass, naming itself.
+  Both polls keep enumerating the statuses that mean *finished*, and an
+  unrecognised status still waits. Calling an unknown status terminal would make
+  the omission easier to diagnose and much worse to live with: one status added
+  upstream would fail every read at once, where waiting costs a single slow call.
+  What made `interrupted` expensive was not the waiting — it was that the
+  timeout never said which status it had been waiting on. Both timeouts now name
+  it.
 
 ## [0.13.0] - 2026-08-27
 

hotdata_framework/client.py

@@ -77,15 +77,17 @@
 _INDEX_TYPES = frozenset(get_args(IndexType))
 _VECTOR_METRICS = frozenset(get_args(VectorMetric))
 
-# Enumerate the in-flight statuses, not the terminal ones. A poll that lists
-# what is terminal treats anything it does not recognise as "keep waiting", so a
-# status the API adds later -- or one that was simply missed -- costs the full
-# timeout before surfacing. Listing what is still in flight makes an unknown
-# status raise on the first pass, naming itself. `interrupted` was missed by
-# exactly the other arrangement, and `cancelled` was listed here without being a
-# status this API sends.
-_RUN_IN_FLIGHT = frozenset({"running"})
-_RESULT_IN_FLIGHT = frozenset({"pending", "processing"})
+# Query-run statuses that mean the run is over. `interrupted` belongs here --
+# omitting it is what made an interrupted run wait out the full timeout -- and
+# `cancelled`, listed here for a long time, is not a status this API sends.
+#
+# Enumerating the terminal side rather than the in-flight side is deliberate. An
+# unrecognised status then keeps polling and costs one slow call, where treating
+# it as terminal would fail every query the moment a status is added upstream.
+# The timeout names the status it last saw, so a missing one is diagnosable
+# without being dangerous.
+_RUN_TERMINAL = frozenset({"succeeded", "failed", "interrupted"})
+_RESULT_FAILURE = frozenset({"failed"})
 # Jobs have no "cancelled" state; "partially_succeeded" carries an error_message.
 _JOB_TERMINAL = frozenset({"succeeded", "partially_succeeded", "failed"})
 
@@ -925,7 +927,7 @@ def _poll_query_run(
         last = None
         while time.monotonic() < deadline:
             last = runs.get_query_run(query_run_id)
-            if last.status not in _RUN_IN_FLIGHT:
+            if last.status in _RUN_TERMINAL:
                 return last
             time.sleep(interval_s)
         raise TimeoutError(
@@ -1034,7 +1036,7 @@ def _wait_result_ready(
             last = results.get_result(result_id)
             if last.status == "ready":
                 return last
-            if last.status not in _RESULT_IN_FLIGHT:
+            if last.status in _RESULT_FAILURE:
                 raise RuntimeError(last.error_message or f"Result {last.status}")
             time.sleep(interval_s)
         raise TimeoutError(

hotdata_framework/managed_client.py

@@ -15,6 +15,7 @@
 import pyarrow as pa
 from hotdata.api.query_api import QueryApi
 from hotdata.api.query_runs_api import QueryRunsApi
+from hotdata.arrow import ResultNotReadyError
 from hotdata.arrow import ResultsApi as ArrowResultsApi
 from hotdata.models.async_query_response import AsyncQueryResponse
 from hotdata.models.query_request import QueryRequest
@@ -40,11 +41,6 @@ class ManagedDatabaseClient:
     database lifecycle.
     """
 
-    # The only status a query run reports while still in flight. Listed this way
-    # round so an unknown status raises immediately rather than being polled to
-    # the timeout; see `_await_query_run`.
-    _RUN_IN_FLIGHT = frozenset({"running"})
-
     _QUERY_TIMEOUT_SECONDS = 300.0
     _POLL_INTERVAL_SECONDS = 0.4
     _MAX_BACKOFF_SECONDS = 30.0
@@ -124,9 +120,22 @@ def _fetch_result_arrow(self, result_id: str, *, database_id: str) -> pa.Table:
         0.6.0 SDK exposes (and requires) ``x_database_id`` on the Arrow
         helper directly.
         """
-        return ArrowResultsApi(self._runtime.api).get_result_arrow(
-            result_id, x_database_id=database_id
-        )
+        arrow = ArrowResultsApi(self._runtime.api)
+        deadline = time.monotonic() + self._QUERY_TIMEOUT_SECONDS
+        while True:
+            try:
+                return arrow.get_result_arrow(result_id, x_database_id=database_id)
+            except ResultNotReadyError:
+                # Waiting on the run should already have made this unreachable:
+                # a run reports `succeeded` only once its result is saved and
+                # ready. Tolerating it anyway costs nothing and removes the need
+                # to take that ordering on trust. The Arrow endpoint answers a
+                # result that is not ready with a small refusal rather than with
+                # data, so waiting here is cheap in the way waiting on the JSON
+                # result body -- which is what this change removed -- is not.
+                if time.monotonic() >= deadline:
+                    raise
+                time.sleep(self._POLL_INTERVAL_SECONDS)
 
     def _query_database_scoped(self, sql: str, *, database_id: str) -> str | None:
         raw = QueryApi(self._runtime.api).query(
@@ -170,9 +179,11 @@ def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None
         """
         runs = QueryRunsApi(self._runtime.api)
         deadline = time.monotonic() + self._QUERY_TIMEOUT_SECONDS
+        last_status: str | None = None
         while time.monotonic() < deadline:
             # Runs (like results) of database-scoped queries are database-scoped.
             run = runs.get_query_run(query_run_id, x_database_id=database_id)
+            last_status = run.status
             if run.status == "succeeded":
                 return run.result_id
             if run.status == "interrupted":
@@ -183,15 +194,19 @@ def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None
                 raise HotdataTransientError(
                     run.error_message or f"Query run {query_run_id} was interrupted"
                 )
-            # Anything not still in flight is terminal, whether or not this
-            # client has heard of it. Enumerating the terminal statuses instead
-            # would poll an unrecognised one to the timeout -- which is exactly
-            # how `interrupted` came to be waited out for five minutes.
-            if run.status not in self._RUN_IN_FLIGHT:
-                raise RuntimeError(run.error_message or f"Query run {query_run_id} {run.status}")
+            if run.status == "failed":
+                raise RuntimeError(run.error_message or f"Query run {query_run_id} failed")
+            # Any other status keeps polling, including one this client has never
+            # seen. Treating an unrecognised status as terminal is the cheaper
+            # failure to diagnose and by far the more expensive one to suffer: a
+            # single status added upstream would then fail every query at once,
+            # where waiting costs one slow call. What made `interrupted`
+            # expensive was not the waiting, it was that the timeout never said
+            # which status it had waited on -- so the message now carries it.
             time.sleep(self._POLL_INTERVAL_SECONDS)
         raise TimeoutError(
-            f"Query run {query_run_id} did not finish within {self._QUERY_TIMEOUT_SECONDS}s"
+            f"Query run {query_run_id} did not finish within "
+            f"{self._QUERY_TIMEOUT_SECONDS}s (last status: {last_status})"
         )
 
     def fetch_table_rows(self, *, database: str, schema: str, table: str) -> list[dict[str, Any]]:

tests/test_client.py

@@ -245,13 +245,28 @@ def test_list_qualified_table_names_passes_connection_id():
     assert it.call_args.kwargs["connection_id"] == "conn_a"
 
 
-def test_wait_result_ready_raises_on_any_status_that_is_not_in_flight():
-    """A result neither `ready` nor still being saved is terminal.
+def test_wait_result_ready_raises_on_a_failed_result():
+    client = HotdataClient("k", "ws", host="https://api.hotdata.dev")
+
+    class FakeResultsApi:
+        def get_result(self, result_id: str):
+            return SimpleNamespace(status="failed", error_message="out of memory")
+
+    with (
+        patch.object(client, "_results_api", return_value=FakeResultsApi()),
+        pytest.raises(RuntimeError, match="out of memory"),
+    ):
+        client._wait_result_ready("res_1", timeout_s=0.1, interval_s=0)
 
-    The poll enumerates the in-flight statuses rather than the terminal ones, so
-    a status this client has never heard of raises on the first pass instead of
-    being waited out. `cancelled` -- which this poll listed as terminal, and
-    which the API does not send -- is as good a stand-in as any.
+
+def test_unknown_result_status_times_out_and_names_the_status():
+    """An unrecognised status keeps polling rather than being called terminal.
+
+    Failing fast on an unknown status would be easier to debug, and far worse to
+    live with: one status added upstream would fail every read at once, where
+    waiting costs a single slow call. The timeout names what it waited on, which
+    is what makes the omission findable -- and what was missing when
+    `interrupted` went unrecognised.
     """
     client = HotdataClient("k", "ws", host="https://api.hotdata.dev")
 
@@ -261,9 +276,9 @@ def get_result(self, result_id: str):
 
     with (
         patch.object(client, "_results_api", return_value=FakeResultsApi()),
-        pytest.raises(RuntimeError, match="something_new"),
+        pytest.raises(TimeoutError, match="something_new"),
     ):
-        client._wait_result_ready("res_1", timeout_s=0.1, interval_s=0)
+        client._wait_result_ready("res_1", timeout_s=0.05, interval_s=0)
 
 
 def test_poll_query_run_returns_promptly_on_interrupted():

tests/test_managed_client.py

@@ -7,12 +7,13 @@
 
 import pyarrow as pa
 import pytest
+from hotdata.arrow import ResultNotReadyError
 from hotdata.models.async_query_response import AsyncQueryResponse
 from hotdata.models.query_response import QueryResponse
 from hotdata.rest import ApiException
 
 import hotdata_framework.managed_client as mc
-from hotdata_framework.errors import HotdataTerminalError
+from hotdata_framework.errors import HotdataTerminalError, HotdataTransientError
 
 
 def _query_response(result_id: str) -> QueryResponse:
@@ -749,16 +750,19 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table:
     ]
 
 
-def test_unknown_run_status_fails_fast_instead_of_polling_to_the_timeout(
+def test_unknown_run_status_keeps_polling_and_the_timeout_names_it(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """Only `running` means "still in flight"; anything else is terminal.
-
-    Enumerating the terminal statuses is what let `interrupted` be polled for
-    five minutes, so the test runs the other way round: a status this client has
-    never seen raises on the first pass and names itself.
+    """An unrecognised run status waits rather than being called terminal.
+
+    Treating an unknown status as terminal is the cheaper failure to diagnose
+    and much the more expensive one to suffer: a single status added upstream
+    would fail every read at once, where waiting costs one slow call. So the
+    poll enumerates what it knows is terminal, and the timeout carries the
+    status it last saw -- which is precisely what was missing while
+    `interrupted` went unrecognised, and what would have made that a one-line
+    diagnosis instead of a mystery.
     """
-    calls: list[str] = []
 
     class FakeQueryApi:
         def __init__(self, api: object) -> None:
@@ -772,7 +776,6 @@ def __init__(self, api: object) -> None:
             pass
 
         def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any:
-            calls.append("get_query_run")
             return SimpleNamespace(status="evicted", result_id=None, error_message=None)
 
     monkeypatch.setattr(mc, "QueryApi", FakeQueryApi)
@@ -787,9 +790,67 @@ def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any:
         retry_backoff_seconds=0.0,
     )
     client._runtime = _fake_runtime()
+    monkeypatch.setattr(client, "_QUERY_TIMEOUT_SECONDS", 0.05)
 
-    with pytest.raises(HotdataTerminalError, match="evicted"):
+    # TimeoutError classifies as transient, so the retry wrapper re-raises it
+    # as such once the budget is spent -- the status still has to reach the text.
+    with pytest.raises(HotdataTransientError, match="evicted"):
         client.fetch_table(database="mydb", schema="public", table="orders")
 
-    # One poll, not a timeout's worth.
-    assert len(calls) == 1
+
+def test_arrow_fetch_waits_out_a_result_that_is_not_ready_yet(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Belt and braces over the run wait.
+
+    A run reports `succeeded` only once its result is saved and ready, so this
+    should not happen. Tolerating it costs nothing and removes the need to take
+    that ordering on trust: the Arrow endpoint answers a result that is not ready
+    with a small refusal rather than with data, so waiting here is cheap in the
+    way waiting on the JSON result body is not.
+    """
+    attempts: list[str] = []
+
+    class FakeQueryApi:
+        def __init__(self, api: object) -> None:
+            pass
+
+        def query(self, request: object, *, x_database_id: str) -> AsyncQueryResponse:
+            return _async_query_response()
+
+    class FakeQueryRunsApi:
+        def __init__(self, api: object) -> None:
+            pass
+
+        def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any:
+            return SimpleNamespace(status="succeeded", result_id="rslt1", error_message=None)
+
+    class FakeArrowResultsApi:
+        def __init__(self, api: object) -> None:
+            pass
+
+        def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table:
+            attempts.append(result_id)
+            if len(attempts) < 3:
+                raise ResultNotReadyError(status="processing", result_id=result_id)
+            return pa.table({"id": [1, 2]})
+
+    monkeypatch.setattr(mc, "QueryApi", FakeQueryApi)
+    monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi)
+    monkeypatch.setattr(mc, "ArrowResultsApi", FakeArrowResultsApi)
+    monkeypatch.setattr(mc.time, "sleep", lambda _seconds: None)
+
+    client = mc.ManagedDatabaseClient(
+        api_key="k",
+        workspace_id="w",
+        api_base_url="https://example.test",
+        max_retries=1,
+        retry_backoff_seconds=0.0,
+    )
+    client._runtime = _fake_runtime()
+
+    table = client.fetch_table(database="mydb", schema="public", table="orders")
+
+    assert table is not None
+    assert table.num_rows == 2
+    assert len(attempts) == 3