From 6c959b0fc9e827aed9189159ed6071636260d95e Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 1 Sep 2026 15:11:01 +0530 Subject: [PATCH 1/6] fix(managed): wait on the query run, not on the result body MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reading a managed table made three calls and used one. `POST /v1/query` returned an inline preview of the rows, `GET /v1/results/{id}` was polled until the result was `ready`, and the result was then fetched as Arrow. Only the Arrow copy was ever read. The readiness poll was the expensive one. `limit` on that endpoint defaults to unbounded, so polling a ready result downloads the whole result body to read one status field. It is also the wrong endpoint to lean on as a table grows: a JSON body over the per-fetch memory budget is refused with 413, and one that would fit alone but not alongside concurrent JSON fetches with 429, so the check fails on exactly the largest tables. The query is now submitted with `async`, so the server returns a run id instead of a preview, and readiness comes from `GET /v1/query-runs/{id}`, which carries no rows at any size. `result_id` is read off the run rather than off the query reply, because a run can succeed having saved nothing and the run is what reports that. Arrow stays the only path the data travels, so column types come from the server's schema rather than being inferred from a JSON preview that has none. That costs one extra round trip on a query that would have answered synchronously, against not transferring the result twice. Two status bugs surfaced while rewriting the poll. It treated `failed` and `cancelled` as the terminal failures: `cancelled` is not a status this API returns, and `interrupted` — which it does return, for a run whose server was replaced before finishing — matched neither branch, so the poll ran to its five-minute timeout rather than failing fast. An interrupted run is safe to re-run, so it is now raised as transient. `classify_sdk_error` grew a passthrough for already-classified errors, without which it demoted that transient error to terminal and cost the retry. `_wait_result_ready` and the generic `_poll` it shared go away, and the plain `ResultsApi` is no longer imported here at all — the only result endpoint this module can now reach is the streaming Arrow one. --- CHANGELOG.md | 39 +++ hotdata_framework/errors.py | 6 + hotdata_framework/managed_client.py | 114 ++++----- tests/test_managed_client.py | 382 +++++++++++++++++++++++++--- 4 files changed, 445 insertions(+), 96 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a85dd12..a00e77d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,45 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- fix(managed): wait on the query run instead of downloading the result to check it. + + Reading a managed table made three calls and used one. `POST /v1/query` returned + an inline preview of the rows, `GET /v1/results/{id}` was polled until the + result was `ready`, and the result was then fetched as Arrow. Only the Arrow + copy was used. + + The readiness poll was the expensive one. `limit` on that endpoint defaults to + unbounded, so polling a ready result downloads the entire result body to read + one status field. It is also the wrong endpoint to lean on as a table grows: + a JSON body over the instance's per-fetch memory budget is refused with 413, + and one that would fit alone but not alongside concurrent JSON fetches with + 429 — so the readiness check starts failing on exactly the largest tables. + + The query is now submitted with `async`, so the server returns a run id rather + than a preview, and readiness comes from `GET /v1/query-runs/{id}`, which + carries no rows at any size. `result_id` is read off the run rather than off + the query reply, because a run can succeed having saved nothing and the run is + what reports that. Arrow stays the only path the data travels, so column types + come from the server's schema rather than being inferred from JSON. + + Costs one extra round trip on a query that would have answered synchronously, + in exchange for not transferring the result twice. + +- fix(managed): recognise `interrupted`, and drop a run status the API never sends. + + The query-run poll treated `failed` and `cancelled` as the terminal failures. + `cancelled` is not a status this API returns. `interrupted` is — a run whose + server was replaced before it finished — and it matched neither branch, so the + poll ran to its five-minute timeout and raised `TimeoutError` instead of + failing fast. + + An interrupted run is safe to re-run, so it is now raised as transient and the + surrounding retry re-submits the query. `classify_sdk_error` passes an + already-classified error through unchanged, rather than demoting a + caller-raised transient error to terminal. + ## [0.13.0] - 2026-08-27 ### Fixed diff --git a/hotdata_framework/errors.py b/hotdata_framework/errors.py index 0f77e43..15f1233 100644 --- a/hotdata_framework/errors.py +++ b/hotdata_framework/errors.py @@ -112,6 +112,12 @@ def _error_class(status_code: int, code: str | None) -> type[HotdataError]: def classify_sdk_error(error: Exception) -> HotdataError: + if isinstance(error, HotdataError): + # Already classified. A caller that read transience off a typed status + # -- an interrupted query run, say -- knows more than this function can + # recover from the exception, and the fallback below would demote it to + # terminal and cost the retry. + return error if isinstance(error, TimeoutError): return HotdataTransientError(str(error)) if isinstance(error, ConnectionError): diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index ed4b9b6..a94b3fc 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -10,12 +10,11 @@ import random import time from collections.abc import Callable -from typing import Any, Protocol, TypeVar +from typing import Any, TypeVar import pyarrow as pa from hotdata.api.query_api import QueryApi from hotdata.api.query_runs_api import QueryRunsApi -from hotdata.api.results_api import ResultsApi from hotdata.arrow import ResultsApi as ArrowResultsApi from hotdata.models.async_query_response import AsyncQueryResponse from hotdata.models.query_request import QueryRequest @@ -32,16 +31,6 @@ T = TypeVar("T") -class _StatusResponse(Protocol): - """Async resources (query runs, results) expose a status and error message.""" - - status: str - error_message: str | None - - -S = TypeVar("S", bound=_StatusResponse) - - class ManagedDatabaseClient: """Managed-database client with bounded retries over hotdata-framework. @@ -134,66 +123,67 @@ def _fetch_result_arrow(self, result_id: str, *, database_id: str) -> pa.Table: result_id, x_database_id=database_id ) - def _poll( - self, - fetch: Callable[[], S], - *, - is_ready: Callable[[S], bool], - describe: str, - ) -> S: - """Poll ``fetch`` until ``is_ready`` is satisfied, or raise on failure/timeout. - - ``failed``/``cancelled`` statuses raise ``RuntimeError``; exceeding - :attr:`_QUERY_TIMEOUT_SECONDS` raises ``TimeoutError``. - """ - deadline = time.monotonic() + self._QUERY_TIMEOUT_SECONDS - while time.monotonic() < deadline: - obj = fetch() - if obj.status in ("failed", "cancelled"): - raise RuntimeError(obj.error_message or f"{describe} {obj.status}") - if is_ready(obj): - return obj - time.sleep(self._POLL_INTERVAL_SECONDS) - raise TimeoutError(f"{describe} timed out after {self._QUERY_TIMEOUT_SECONDS}s") - def _query_database_scoped(self, sql: str, *, database_id: str) -> str | None: raw = QueryApi(self._runtime.api).query( - QueryRequest(sql=sql), + # Asked asynchronously because this caller wants a result id, not + # rows. A synchronous submit always builds an inline preview of the + # result and sends it -- megabytes, on a path that then reads the + # whole result as Arrow anyway and never looks at the preview. The + # async reply carries a run id and nothing else, and there is no way + # to suppress the preview on a synchronous one. + # + # It also settles the types: the preview is JSON, which has no Arrow + # schema and renders non-finite floats as null, so it could not have + # substituted for the Arrow fetch even when it holds every row. + # + # `var_async` is the generated SDK's spelling of the wire field + # `async`, which is a Python keyword and so cannot be the attribute + # name. + QueryRequest(sql=sql, var_async=True), x_database_id=database_id, ) - if isinstance(raw, QueryResponse): - # A synchronous response still persists its full result out-of-band - # under ``result_id``; that result may be ``processing`` when the - # inline preview returns, so wait for ``ready`` before the caller - # fetches it as Arrow. - return self._wait_result_ready(raw.result_id, database_id=database_id) - if isinstance(raw, AsyncQueryResponse): - run_result = self._await_query_run(raw.query_run_id, database_id=database_id) - return self._wait_result_ready(run_result, database_id=database_id) + # Both reply shapes carry `query_run_id`, and the run is the readiness + # signal for either -- a synchronous reply (which `async_after_ms` can + # still produce) returns rows inline but goes on saving the full result + # in the background, so it is not the finish line either. + if isinstance(raw, (QueryResponse, AsyncQueryResponse)): + return self._await_query_run(raw.query_run_id, database_id=database_id) return None def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None: + """Wait for a query run to finish; return the result id it produced. + + The run is the whole wait. A run turns `succeeded` only after its result + has been saved and is `ready`, so `succeeded` needs no second check + against the result -- and asking the result endpoint instead would mean + downloading the entire result to read one field, which the server + refuses outright (413/429) once the result is large enough. + + `result_id` comes off the run rather than off the query reply because a + `succeeded` run reports none when every row came back inline but the + result could not be saved for later retrieval. + """ runs = QueryRunsApi(self._runtime.api) - run = self._poll( + deadline = time.monotonic() + self._QUERY_TIMEOUT_SECONDS + while time.monotonic() < deadline: # Runs (like results) of database-scoped queries are database-scoped. - lambda: runs.get_query_run(query_run_id, x_database_id=database_id), - is_ready=lambda r: r.status == "succeeded", - describe="Query", - ) - return run.result_id - - def _wait_result_ready(self, result_id: str | None, *, database_id: str) -> str | None: - if result_id is None: - return None - results = ResultsApi(self._runtime.api) - self._poll( - # The stored result of a database-scoped query 400s without the - # database scope. - lambda: results.get_result(result_id, x_database_id=database_id), - is_ready=lambda r: r.status == "ready", - describe=f"Result {result_id}", + run = runs.get_query_run(query_run_id, x_database_id=database_id) + if run.status == "succeeded": + return run.result_id + if run.status == "interrupted": + # Terminal, but the server lost the run rather than rejecting + # the query, so it is the one failure here worth re-running. + # Raised pre-classified: `classify_sdk_error` cannot tell this + # apart from an ordinary RuntimeError and would call it terminal. + raise HotdataTransientError( + run.error_message or f"Query run {query_run_id} was interrupted" + ) + if run.status == "failed": + raise RuntimeError(run.error_message or f"Query run {query_run_id} failed") + time.sleep(self._POLL_INTERVAL_SECONDS) + raise TimeoutError( + f"Query run {query_run_id} did not finish within {self._QUERY_TIMEOUT_SECONDS}s" ) - return result_id def fetch_table_rows(self, *, database: str, schema: str, table: str) -> list[dict[str, Any]]: result = self.fetch_table(database=database, schema=schema, table=table) diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index b7eed9b..d3542aa 100644 --- a/tests/test_managed_client.py +++ b/tests/test_managed_client.py @@ -7,6 +7,7 @@ import pyarrow as pa import pytest +from hotdata.models.async_query_response import AsyncQueryResponse from hotdata.models.query_response import QueryResponse from hotdata.rest import ApiException @@ -28,13 +29,29 @@ def _query_response(result_id: str) -> QueryResponse: ) -def test_fetch_table_waits_for_ready_before_arrow(monkeypatch: pytest.MonkeyPatch) -> None: - """A synchronous ``QueryResponse`` persists its full result out-of-band, and - that result can still be ``processing`` when the inline preview returns. +def _async_query_response() -> AsyncQueryResponse: + return AsyncQueryResponse( + query_run_id="qr", + status="running", + status_url="/v1/query-runs/qr", + ) - ``fetch_table`` must poll the result to ``ready`` before fetching it as - Arrow. The earlier bug returned the ``result_id`` immediately on the sync - path, so Arrow was fetched against a ``processing`` result and failed. + +def test_fetch_table_waits_on_the_query_run_not_the_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Readiness is the query run's answer to give, not the result endpoint's. + + A synchronous ``QueryResponse`` returns its rows inline but goes on saving + the full result in the background, so the reply is not the finish line and + Arrow cannot be fetched yet. The run turns ``succeeded`` only once that + result is saved and ready, which makes it a sufficient readiness signal on + its own. + + The earlier version polled ``GET /results/{id}`` instead, which is the + result *data* endpoint: a ready JSON reply carries the whole result, so the + check downloaded all of it to read one status field, and the server refuses + that outright once the result is large enough. """ calls: list[str] = [] @@ -46,16 +63,16 @@ def query(self, request: object, *, x_database_id: str) -> QueryResponse: calls.append("query") return _query_response("rslt1") - statuses = iter(["processing", "processing", "ready"]) + statuses = iter(["running", "running", "succeeded"]) - class FakeResultsApi: + class FakeQueryRunsApi: def __init__(self, api: object) -> None: pass - def get_result(self, result_id: str, **kwargs: Any) -> Any: + def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any: status = next(statuses) - calls.append(f"get_result:{status}") - return SimpleNamespace(status=status, result_id=result_id, error_message=None) + calls.append(f"get_query_run:{status}") + return SimpleNamespace(status=status, result_id="rslt1", error_message=None) class FakeArrowResultsApi: def __init__(self, api: object) -> None: @@ -66,7 +83,7 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: return pa.table({"id": [1, 2]}) monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) - monkeypatch.setattr(mc, "ResultsApi", FakeResultsApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) monkeypatch.setattr(mc, "ArrowResultsApi", FakeArrowResultsApi) monkeypatch.setattr(mc.time, "sleep", lambda _seconds: None) @@ -83,10 +100,23 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: assert table is not None assert table.num_rows == 2 - # The result was polled to readiness, and Arrow was fetched only afterwards. - assert "get_result:processing" in calls - assert "get_result:ready" in calls - assert calls.index("arrow") > calls.index("get_result:ready") + # The run was polled while still running, and Arrow was fetched only once + # it had succeeded. + assert "get_query_run:running" in calls + assert calls.index("arrow") > calls.index("get_query_run:succeeded") + + +def test_readiness_never_touches_the_json_results_endpoint() -> None: + """The generated ``ResultsApi`` -- the JSON result *data* endpoint -- must + not be reachable from this module at all. + + Its absence is the fix: readiness comes from the query run, and the only + result endpoint this client is allowed to call is the streaming Arrow one, + imported under a distinct name. A future edit that reintroduces the plain + ``ResultsApi`` here is reintroducing a whole-result download per poll. + """ + assert not hasattr(mc, "ResultsApi") + assert hasattr(mc, "ArrowResultsApi") def _fake_runtime() -> SimpleNamespace: @@ -99,19 +129,18 @@ def _fake_runtime() -> SimpleNamespace: ) -def test_fetch_table_carries_database_scope_on_result_reads( +def test_fetch_table_carries_database_scope_on_reads( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Results (and runs) of a database-scoped query are database-scoped: - the results endpoints 400 with "X-Database-Id header is required" when - the scope is missing. ``fetch_table`` must carry the database id on the - result poll and the Arrow fetch, not only on the query submit — the - hotdata 0.6.0 SDK exposes ``x_database_id`` on all three. + """Runs (and results) of a database-scoped query are database-scoped: the + endpoints 400 with "X-Database-Id header is required" when the scope is + missing. ``fetch_table`` must carry the database id on the run poll and the + Arrow fetch, not only on the query submit. Regression: reruns/append loads against an existing synced table failed with an opaque ``400: Bad Request`` because both reads omitted the scope. """ - result_scopes: list[str | None] = [] + run_scopes: list[str | None] = [] arrow_scopes: list[str | None] = [] class FakeQueryApi: @@ -122,26 +151,26 @@ def query(self, request: object, *, x_database_id: str) -> QueryResponse: assert x_database_id == "db1" return _query_response("rslt1") - class FakeResultsApi: + class FakeQueryRunsApi: def __init__(self, api: object) -> None: pass - def get_result(self, result_id: str, *, x_database_id: str | None = None) -> Any: - result_scopes.append(x_database_id) - return SimpleNamespace(status="ready", result_id=result_id, error_message=None) + # x_database_id is REQUIRED on this endpoint -- mirroring that here + # makes this test fail if a caller ever drops the scope again. + def get_query_run(self, query_run_id: str, *, x_database_id: str) -> Any: + run_scopes.append(x_database_id) + return SimpleNamespace(status="succeeded", result_id="rslt1", error_message=None) class FakeArrowResultsApi: def __init__(self, api: object) -> None: pass - # x_database_id is REQUIRED in the 0.6.0 SDK — mirroring that here - # makes this test fail if a caller ever drops the scope again. def get_result_arrow(self, result_id: str, *, x_database_id: str) -> pa.Table: arrow_scopes.append(x_database_id) return pa.table({"id": [1]}) monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) - monkeypatch.setattr(mc, "ResultsApi", FakeResultsApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) monkeypatch.setattr(mc, "ArrowResultsApi", FakeArrowResultsApi) client = mc.ManagedDatabaseClient( @@ -156,10 +185,175 @@ def get_result_arrow(self, result_id: str, *, x_database_id: str) -> pa.Table: table = client.fetch_table(database="mydb", schema="public", table="orders") assert table is not None - assert result_scopes == ["db1"] + assert run_scopes == ["db1"] assert arrow_scopes == ["db1"] +def test_interrupted_query_run_is_retried_not_waited_out( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``interrupted`` is terminal but safe to retry: the server lost the run + rather than rejecting the query. + + It has to be raised as transient so the surrounding retry re-runs the + query. The earlier poll recognised only ``failed`` and a ``cancelled`` + status the API never sends, so an interrupted run matched neither -- it + spun for the full five-minute timeout and then failed. + """ + calls: list[str] = [] + statuses = iter(["interrupted", "succeeded"]) + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: object, *, x_database_id: str) -> QueryResponse: + calls.append("query") + return _query_response("rslt1") + + class FakeQueryRunsApi: + def __init__(self, api: object) -> None: + pass + + def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any: + status = next(statuses) + calls.append(f"get_query_run:{status}") + return SimpleNamespace( + status=status, + result_id="rslt1", + error_message="instance lost" if status == "interrupted" else None, + ) + + class FakeArrowResultsApi: + def __init__(self, api: object) -> None: + pass + + def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: + calls.append("arrow") + return pa.table({"id": [1]}) + + 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=2, + retry_backoff_seconds=0.0, + ) + client._runtime = _fake_runtime() + + table = client.fetch_table(database="mydb", schema="public", table="orders") + + assert table is not None + # The interrupted run re-ran the query rather than being waited out. + assert calls.count("query") == 2 + assert calls == [ + "query", + "get_query_run:interrupted", + "query", + "get_query_run:succeeded", + "arrow", + ] + + +def test_failed_query_run_is_not_retried(monkeypatch: pytest.MonkeyPatch) -> None: + """``failed`` means the query itself failed, so every retry reaches the same + answer. It must surface the run's own message rather than being re-run.""" + calls: list[str] = [] + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: object, *, x_database_id: str) -> QueryResponse: + calls.append("query") + return _query_response("rslt1") + + class FakeQueryRunsApi: + def __init__(self, api: object) -> None: + pass + + def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any: + return SimpleNamespace( + status="failed", result_id=None, error_message="no such column: nope" + ) + + monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) + 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=3, + retry_backoff_seconds=0.0, + ) + client._runtime = _fake_runtime() + + with pytest.raises(HotdataTerminalError, match="no such column"): + client.fetch_table(database="mydb", schema="public", table="orders") + + assert calls.count("query") == 1 + + +def test_result_id_is_read_off_the_run_not_the_query_reply( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A ``succeeded`` run reports ``result_id: null`` when every row came back + inline but the result could not be saved for later retrieval. + + The query reply's own ``result_id`` is optimistic -- it names a result that + may never resolve -- so the run's is the one to believe. Believing the reply + means fetching Arrow against an id the server will 404. + """ + arrow_calls: list[str] = [] + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: object, *, x_database_id: str) -> QueryResponse: + # The reply hands out an id... + return _query_response("rslt1") + + class FakeQueryRunsApi: + def __init__(self, api: object) -> None: + pass + + def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any: + # ...that the run says was never saved. + return SimpleNamespace(status="succeeded", result_id=None, error_message=None) + + class FakeArrowResultsApi: + def __init__(self, api: object) -> None: + pass + + def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: + arrow_calls.append(result_id) + return pa.table({"id": [1]}) + + monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) + monkeypatch.setattr(mc, "ArrowResultsApi", FakeArrowResultsApi) + + 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() + + assert client.fetch_table(database="mydb", schema="public", table="orders") is None + assert arrow_calls == [] + + def _load_recording_runtime(calls: list[str], uploads: list[str] | None = None) -> SimpleNamespace: """A runtime whose ``load_managed_table`` records each mode and always fails with a transient error, so retry behaviour is observable via ``calls``. @@ -415,9 +609,7 @@ def load_managed_table( ) -> SimpleNamespace: captured["mode"] = mode captured["key"] = key - return SimpleNamespace( - connection_id="c", schema_name=schema, table_name=table, row_count=0 - ) + return SimpleNamespace(connection_id="c", schema_name=schema, table_name=table, row_count=0) client = _managed_client(max_retries=1) runtime = _fake_runtime() @@ -428,3 +620,125 @@ def load_managed_table( "db", "orders", schema="public", upload_id="u1", mode="delete", key=["id"] ) assert captured == {"mode": "delete", "key": ["id"]} + + +def test_query_is_submitted_async_so_no_preview_is_built( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``fetch_table`` wants a result id, not rows. + + A synchronous submit always serialises an inline preview of the result into + its reply, and there is no request field that suppresses it -- so the only + way not to be sent megabytes this path never reads is to ask + asynchronously. The async reply carries a run id and nothing else. + """ + requests: list[Any] = [] + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: Any, *, x_database_id: str) -> AsyncQueryResponse: + requests.append(request) + 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: + return pa.table({"id": [1]}) + + monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) + monkeypatch.setattr(mc, "ArrowResultsApi", FakeArrowResultsApi) + + 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 len(requests) == 1 + assert requests[0].var_async is True + + +def test_async_reply_is_followed_through_to_arrow( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The async reply carries no ``result_id`` at all -- only a run id. + + So the run is not merely the cheapest way to learn the result is ready, it + is the only way to learn the result's id in the first place. + """ + calls: list[str] = [] + statuses = iter(["running", "succeeded"]) + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: object, *, x_database_id: str) -> AsyncQueryResponse: + calls.append("query") + 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: + assert query_run_id == "qr" + status = next(statuses) + calls.append(f"get_query_run:{status}") + return SimpleNamespace( + status=status, + result_id="rslt-from-run" if status == "succeeded" else None, + error_message=None, + ) + + class FakeArrowResultsApi: + def __init__(self, api: object) -> None: + pass + + def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: + calls.append(f"arrow:{result_id}") + return pa.table({"id": [1, 2, 3]}) + + 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 == 3 + # The id Arrow was fetched with came off the run, not the query reply. + assert calls == [ + "query", + "get_query_run:running", + "get_query_run:succeeded", + "arrow:rslt-from-run", + ] From 21730a37b810fe9cf3629ff9d13c6a07441aa9b7 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 1 Sep 2026 15:21:43 +0530 Subject: [PATCH 2/6] fix(managed): enumerate in-flight statuses, not terminal ones Follow-up from review on the query-run wait. Both polls listed the statuses that mean finished and treated everything else as "keep waiting". That arrangement is what let `interrupted` be polled for the full five-minute timeout: a status absent from the list was indistinguishable from one still running. Listing the in-flight side instead -- `running` for a query run, `pending`/`processing` for a result -- makes an unrecognised status raise on the first pass with its own name in the message, so the next status the API adds costs one round trip rather than a timeout. That also fixes both status bugs in `HotdataClient`, which still had `cancelled` in its terminal sets and lacked `interrupted` entirely, so `execute_sql` kept the behaviour this branch documents as a bug. Doing both keeps the changelog entry true for the whole package rather than for one of its two clients. The async-submit test asserted the model attribute, which would still pass if the `async` alias were missing or misspelled -- the server would then ignore the field, answer synchronously with a full preview, and silently restore what this branch removes. It now asserts the serialised wire key. --- CHANGELOG.md | 28 +++++++++------ hotdata_framework/client.py | 15 +++++--- hotdata_framework/managed_client.py | 13 +++++-- tests/test_client.py | 38 +++++++++++++++++++-- tests/test_managed_client.py | 53 ++++++++++++++++++++++++++++- 5 files changed, 127 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a00e77d..1b52d31 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,16 +35,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - fix(managed): recognise `interrupted`, and drop a run status the API never sends. - The query-run poll treated `failed` and `cancelled` as the terminal failures. - `cancelled` is not a status this API returns. `interrupted` is — a run whose - server was replaced before it finished — and it matched neither branch, so the - poll ran to its five-minute timeout and raised `TimeoutError` instead of - failing fast. - - An interrupted run is safe to re-run, so it is now raised as transient and the - surrounding retry re-submits the query. `classify_sdk_error` passes an - already-classified error through unchanged, rather than demoting a - caller-raised transient error to terminal. + Both `ManagedDatabaseClient` and `HotdataClient` treated `failed` and + `cancelled` as the terminal run failures. `cancelled` is not a status this API + returns. `interrupted` is — a run whose server was replaced before it finished + — and it matched neither, so an interrupted run was polled for the full + five-minute timeout and then raised `TimeoutError`: a retryable condition + hidden behind a long wait and an error naming the wrong problem. + + On `ManagedDatabaseClient` an interrupted run is now raised as transient, so + the surrounding retry re-submits the query. That needed `classify_sdk_error` to + pass an already-classified error through unchanged rather than demoting a + 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. ## [0.13.0] - 2026-08-27 diff --git a/hotdata_framework/client.py b/hotdata_framework/client.py index ff6570d..cb93369 100644 --- a/hotdata_framework/client.py +++ b/hotdata_framework/client.py @@ -77,8 +77,15 @@ _INDEX_TYPES = frozenset(get_args(IndexType)) _VECTOR_METRICS = frozenset(get_args(VectorMetric)) -_TERMINAL = frozenset({"succeeded", "failed", "cancelled"}) -_RESULT_FAILURE = frozenset({"failed", "cancelled"}) +# 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"}) # Jobs have no "cancelled" state; "partially_succeeded" carries an error_message. _JOB_TERMINAL = frozenset({"succeeded", "partially_succeeded", "failed"}) @@ -918,7 +925,7 @@ def _poll_query_run( last = None while time.monotonic() < deadline: last = runs.get_query_run(query_run_id) - if last.status in _TERMINAL: + if last.status not in _RUN_IN_FLIGHT: return last time.sleep(interval_s) raise TimeoutError( @@ -1027,7 +1034,7 @@ def _wait_result_ready( last = results.get_result(result_id) if last.status == "ready": return last - if last.status in _RESULT_FAILURE: + if last.status not in _RESULT_IN_FLIGHT: raise RuntimeError(last.error_message or f"Result {last.status}") time.sleep(interval_s) raise TimeoutError( diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index a94b3fc..b96e898 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -40,6 +40,11 @@ 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 @@ -178,8 +183,12 @@ 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" ) - if run.status == "failed": - raise RuntimeError(run.error_message or f"Query run {query_run_id} failed") + # 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}") time.sleep(self._POLL_INTERVAL_SECONDS) raise TimeoutError( f"Query run {query_run_id} did not finish within {self._QUERY_TIMEOUT_SECONDS}s" diff --git a/tests/test_client.py b/tests/test_client.py index 7ead972..a155b56 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -245,20 +245,52 @@ 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_cancelled(): +def test_wait_result_ready_raises_on_any_status_that_is_not_in_flight(): + """A result neither `ready` nor still being saved is terminal. + + 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. + """ client = HotdataClient("k", "ws", host="https://api.hotdata.dev") class FakeResultsApi: def get_result(self, result_id: str): - return SimpleNamespace(status="cancelled", error_message=None) + return SimpleNamespace(status="something_new", error_message=None) with ( patch.object(client, "_results_api", return_value=FakeResultsApi()), - pytest.raises(RuntimeError, match="cancelled"), + pytest.raises(RuntimeError, match="something_new"), ): client._wait_result_ready("res_1", timeout_s=0.1, interval_s=0) +def test_poll_query_run_returns_promptly_on_interrupted(): + """`interrupted` is terminal, so the poll must stop on it. + + It was absent from the terminal set, so an interrupted run was polled for the + full timeout and then raised `TimeoutError` -- hiding a retryable condition + behind a five-minute wait, behind an error naming the wrong problem. + """ + client = HotdataClient("k", "ws", host="https://api.hotdata.dev") + calls: list[str] = [] + + class FakeQueryRunsApi: + def get_query_run(self, query_run_id: str): + calls.append(query_run_id) + return SimpleNamespace( + status="interrupted", error_message="instance lost", result_id=None + ) + + with patch.object(client, "_query_runs_api", return_value=FakeQueryRunsApi()): + run = client._poll_query_run("qrun_1", timeout_s=30.0, interval_s=0) + + assert run.status == "interrupted" + # One request, not a timeout's worth. + assert len(calls) == 1 + + def test_connection_id_by_name_raises_on_duplicate_names(): client = HotdataClient("k", "ws", host="https://api.hotdata.dev") listing = SimpleNamespace( diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index d3542aa..df4c3c6 100644 --- a/tests/test_managed_client.py +++ b/tests/test_managed_client.py @@ -673,7 +673,12 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: assert table is not None assert len(requests) == 1 - assert requests[0].var_async is True + # The attribute alone is not the contract: the field only reaches the server + # as `async`. Were the alias missing or misspelled, the attribute assertion + # would still pass, the server would ignore the field and answer + # synchronously with a full preview -- restoring the exact behaviour this + # change removes, silently. + assert requests[0].to_dict()["async"] is True def test_async_reply_is_followed_through_to_arrow( @@ -742,3 +747,49 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: "get_query_run:succeeded", "arrow:rslt-from-run", ] + + +def test_unknown_run_status_fails_fast_instead_of_polling_to_the_timeout( + 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. + """ + calls: 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: + calls.append("get_query_run") + return SimpleNamespace(status="evicted", result_id=None, error_message=None) + + monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) + 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() + + with pytest.raises(HotdataTerminalError, match="evicted"): + client.fetch_table(database="mydb", schema="public", table="orders") + + # One poll, not a timeout's worth. + assert len(calls) == 1 From 258859a36c8f9719a04d76a1e6856df2ddec7fdb Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 1 Sep 2026 15:28:18 +0530 Subject: [PATCH 3/6] 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. --- CHANGELOG.md | 18 ++++-- hotdata_framework/client.py | 24 ++++---- hotdata_framework/managed_client.py | 45 ++++++++++----- tests/test_client.py | 31 ++++++++--- tests/test_managed_client.py | 85 +++++++++++++++++++++++++---- 5 files changed, 152 insertions(+), 51 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b52d31..5c368c2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/hotdata_framework/client.py b/hotdata_framework/client.py index cb93369..8716bca 100644 --- a/hotdata_framework/client.py +++ b/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( diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index b96e898..db68d12 100644 --- a/hotdata_framework/managed_client.py +++ b/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]]: diff --git a/tests/test_client.py b/tests/test_client.py index a155b56..741a374 100644 --- a/tests/test_client.py +++ b/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(): diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index df4c3c6..5d7d3c8 100644 --- a/tests/test_managed_client.py +++ b/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 From b017161dd790596f193a42b591f786bc9a666c80 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 1 Sep 2026 15:34:19 +0530 Subject: [PATCH 4/6] fix(managed): raise when a succeeded run saved no result A run succeeds with no result id when its rows came back inline but the result could not be saved for later retrieval. `fetch_table` answered `None` for that, `fetch_table_rows` turns `None` into `[]`, and `[]` is also the answer both give for a table that is not synced -- so the two states were indistinguishable to a caller. A read-modify-write load would therefore read no existing rows and write only its new batch, dropping every row already in the table. Silent data loss is the worst outcome on offer, and raising costs nothing if the state never occurs. Terminal rather than transient: re-running the query cannot save a result that was already discarded. The run's `warning_message` goes into the error, since it is the only thing that explains why. The test this replaces asserted the old behaviour -- it pinned `fetch_table` returning `None` for exactly this case, which committed the data-loss path to the suite rather than catching it. --- CHANGELOG.md | 9 +++- hotdata_framework/managed_client.py | 16 ++++++ tests/test_managed_client.py | 77 +++++++++++++++++++++++++---- 3 files changed, 90 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c368c2..83b1ed4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,8 +27,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 than a preview, and readiness comes from `GET /v1/query-runs/{id}`, which carries no rows at any size. `result_id` is read off the run rather than off the query reply, because a run can succeed having saved nothing and the run is - what reports that. Arrow stays the only path the data travels, so column types - come from the server's schema rather than being inferred from JSON. + what reports that — and that case now raises rather than reading as an empty + table. `fetch_table` answered `None` for it, which `fetch_table_rows` turns + into `[]`, the same answer both give for a table that is not synced. A + read-modify-write load would have read no existing rows and written only its + new batch, dropping every row already there. Arrow stays the only path the + data travels, so column types come from the server's schema rather than being + inferred from JSON. Costs one extra round trip on a query that would have answered synchronously, in exchange for not transferring the result twice. diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index db68d12..3152cfa 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -185,6 +185,22 @@ def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None run = runs.get_query_run(query_run_id, x_database_id=database_id) last_status = run.status if run.status == "succeeded": + if run.result_id is None: + # A run succeeds with no result id when its rows were + # returned inline but the result could not be saved for + # later retrieval. Returning nothing here would surface as + # an empty table -- `fetch_table` answers `None`, and + # `fetch_table_rows` turns that into `[]`, which is the same + # answer it gives for a table that does not exist. A + # read-modify-write load would then read no existing rows + # and write only its new batch, dropping what was there. + # Terminal rather than transient: re-running the query + # cannot save a result that was already discarded. + raise RuntimeError( + f"Query run {query_run_id} succeeded but its result was not " + f"saved, so the table cannot be read" + + (f": {run.warning_message}" if run.warning_message else "") + ) return run.result_id if run.status == "interrupted": # Terminal, but the server lost the run rather than rejecting diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index 5d7d3c8..81ac2f0 100644 --- a/tests/test_managed_client.py +++ b/tests/test_managed_client.py @@ -302,15 +302,21 @@ def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any: assert calls.count("query") == 1 -def test_result_id_is_read_off_the_run_not_the_query_reply( +def test_a_succeeded_run_that_saved_no_result_raises_rather_than_reading_empty( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A ``succeeded`` run reports ``result_id: null`` when every row came back - inline but the result could not be saved for later retrieval. - - The query reply's own ``result_id`` is optimistic -- it names a result that - may never resolve -- so the run's is the one to believe. Believing the reply - means fetching Arrow against an id the server will 404. + """The one state where an empty answer would be a wrong answer. + + A run succeeds with no `result_id` when its rows came back inline but the + result could not be saved. Answering `None` puts that on the same footing as + a table that does not exist -- `fetch_table_rows` turns both into `[]` -- so + a read-modify-write load would read no existing rows and write only its new + batch, dropping the rows already in the table. Silent data loss is the worst + outcome available here, so this raises, and raises terminally: re-running + cannot save a result that was already discarded. + + The query reply carries an id of its own, and it must not be believed over + the run's. """ arrow_calls: list[str] = [] @@ -328,7 +334,12 @@ def __init__(self, api: object) -> None: def get_query_run(self, query_run_id: str, **kwargs: Any) -> Any: # ...that the run says was never saved. - return SimpleNamespace(status="succeeded", result_id=None, error_message=None) + return SimpleNamespace( + status="succeeded", + result_id=None, + error_message=None, + warning_message="result row creation failed; result not persisted", + ) class FakeArrowResultsApi: def __init__(self, api: object) -> None: @@ -341,20 +352,66 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: 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, + max_retries=3, retry_backoff_seconds=0.0, ) client._runtime = _fake_runtime() - assert client.fetch_table(database="mydb", schema="public", table="orders") is None + with pytest.raises(HotdataTerminalError, match="result was not saved"): + client.fetch_table(database="mydb", schema="public", table="orders") + + # The reply's id was never used. assert arrow_calls == [] +def test_fetch_table_rows_cannot_turn_an_unsaved_result_into_no_rows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`fetch_table_rows` is where an empty answer does the damage. + + It maps `None` to `[]`, which is also its answer for a table that is not + synced, so the unsaved-result case must not be able to reach that mapping. + """ + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: object, *, x_database_id: str) -> QueryResponse: + return _query_response("rslt1") + + 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=None, error_message=None, warning_message=None + ) + + monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) + monkeypatch.setattr(mc, "QueryRunsApi", FakeQueryRunsApi) + 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() + + with pytest.raises(HotdataTerminalError): + client.fetch_table_rows(database="mydb", schema="public", table="orders") + + def _load_recording_runtime(calls: list[str], uploads: list[str] | None = None) -> SimpleNamespace: """A runtime whose ``load_managed_table`` records each mode and always fails with a transient error, so retry behaviour is observable via ``calls``. From 19ae3d4cd9470760c70ed16be953744fd1c15bfd Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 1 Sep 2026 15:38:55 +0530 Subject: [PATCH 5/6] fix(managed): read the run's warning defensively, and pin the fields `warning_message` is read while building an error, which is the worst place to assume an attribute: if an SDK release dropped it, the `AttributeError` would replace a message naming the problem with one naming nothing. `getattr` keeps the raise intact and loses only the explanation. Pinning the assumption is the better half of the fix, though. Every fake in the managed-client tests is a `SimpleNamespace`, so nothing there would notice a field being renamed or dropped -- the fakes would keep answering and the suite would keep passing, with the cost landing at runtime. A test now asserts the generated `QueryRunInfo` carries all four fields this client reads off a run, not just the one that prompted this. --- hotdata_framework/managed_client.py | 6 +++++- tests/test_managed_client.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index 3152cfa..aac3d47 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -196,10 +196,14 @@ def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None # and write only its new batch, dropping what was there. # Terminal rather than transient: re-running the query # cannot save a result that was already discarded. + # `getattr` because this runs while building an error: if + # the field ever goes away, losing the explanation is a far + # better outcome than an AttributeError replacing the raise. + warning = getattr(run, "warning_message", None) raise RuntimeError( f"Query run {query_run_id} succeeded but its result was not " f"saved, so the table cannot be read" - + (f": {run.warning_message}" if run.warning_message else "") + + (f": {warning}" if warning else "") ) return run.result_id if run.status == "interrupted": diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index 81ac2f0..a606aa9 100644 --- a/tests/test_managed_client.py +++ b/tests/test_managed_client.py @@ -10,6 +10,7 @@ from hotdata.arrow import ResultNotReadyError from hotdata.models.async_query_response import AsyncQueryResponse from hotdata.models.query_response import QueryResponse +from hotdata.models.query_run_info import QueryRunInfo from hotdata.rest import ApiException import hotdata_framework.managed_client as mc @@ -911,3 +912,17 @@ def get_result_arrow(self, result_id: str, **kwargs: Any) -> pa.Table: assert table is not None assert table.num_rows == 2 assert len(attempts) == 3 + + +def test_query_run_model_carries_every_field_this_client_reads() -> None: + """Pins the attributes read off a query run against the generated model. + + Every fake in this file is a `SimpleNamespace`, so nothing else here would + notice if one of these fields were renamed or dropped by an SDK release -- + the fakes would keep answering and the tests would keep passing. The cost + lands at runtime, and worst on `warning_message`, which is read while + building an error: losing it turns a message naming the problem into an + `AttributeError` naming nothing. + """ + for field in ("status", "result_id", "error_message", "warning_message"): + assert field in QueryRunInfo.model_fields, field From 516103649ba3dc6bd87f974ed33ada885b78cf46 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 1 Sep 2026 15:42:36 +0530 Subject: [PATCH 6/6] fix(managed): raise on a query reply shape this client cannot read The last route by which "I do not know" could read as "no rows". If an SDK release adds a third reply model for the query endpoint, neither `isinstance` branch matches, `_query_database_scoped` returns `None`, `fetch_table` returns `None`, and `fetch_table_rows` turns that into `[]` -- the same answer it gives for a table that is not synced. A merge or append load would then write only its new batch over rows it believed were absent. Same collapse the previous commit closed for a succeeded run that saved nothing, by the other path into it. `HotdataClient` already raised here, so this also settles a disagreement between the two clients. A `None` from `fetch_table` now carries exactly one meaning: the table is not synced. --- CHANGELOG.md | 8 +++--- hotdata_framework/managed_client.py | 8 +++++- tests/test_managed_client.py | 41 +++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 83b1ed4..606d83c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,9 +31,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 table. `fetch_table` answered `None` for it, which `fetch_table_rows` turns into `[]`, the same answer both give for a table that is not synced. A read-modify-write load would have read no existing rows and written only its - new batch, dropping every row already there. Arrow stays the only path the - data travels, so column types come from the server's schema rather than being - inferred from JSON. + new batch, dropping every row already there. A reply shape this client does not + recognise raises for the same reason, as `HotdataClient` already did — so a + `None` from `fetch_table` now means one thing only: the table is not synced. + Arrow stays the only path the data travels, so column types come from the + server's schema rather than being inferred from JSON. Costs one extra round trip on a query that would have answered synchronously, in exchange for not transferring the result twice. diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index aac3d47..8ed9627 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -162,7 +162,13 @@ def _query_database_scoped(self, sql: str, *, database_id: str) -> str | None: # in the background, so it is not the finish line either. if isinstance(raw, (QueryResponse, AsyncQueryResponse)): return self._await_query_run(raw.query_run_id, database_id=database_id) - return None + # Returning nothing here would read as an empty table: `fetch_table` + # answers `None`, `fetch_table_rows` turns that into `[]`, and a + # read-modify-write load would write only its new batch over rows it + # believed were not there. A reply shape this client does not know is a + # reason to stop, not to report emptiness. `HotdataClient` raises on the + # same condition. + raise RuntimeError(f"Unexpected query response type: {type(raw)!r}") def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None: """Wait for a query run to finish; return the result id it produced. diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index a606aa9..3b54c05 100644 --- a/tests/test_managed_client.py +++ b/tests/test_managed_client.py @@ -926,3 +926,44 @@ def test_query_run_model_carries_every_field_this_client_reads() -> None: """ for field in ("status", "result_id", "error_message", "warning_message"): assert field in QueryRunInfo.model_fields, field + + +def test_an_unknown_query_reply_shape_raises_rather_than_reading_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The last route by which "I do not know" could have read as "no rows". + + If an SDK release adds a third reply model for the query endpoint, neither + `isinstance` branch matches. Falling through to `None` would surface as an + empty table -- and `fetch_table_rows` maps `None` to `[]`, the same answer it + gives for a table that is not synced -- so a merge or append load would drop + every row already there. `HotdataClient` raises on this condition; now both + do. + + After this, a `None` from `fetch_table` means one thing only: the table is + not synced. + """ + + class FakeQueryApi: + def __init__(self, api: object) -> None: + pass + + def query(self, request: object, *, x_database_id: str) -> Any: + # A shape from neither branch -- a future reply model, as far as + # this client is concerned. + return SimpleNamespace(something_new="?") + + monkeypatch.setattr(mc, "QueryApi", FakeQueryApi) + 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() + + with pytest.raises(HotdataTerminalError, match="Unexpected query response type"): + client.fetch_table_rows(database="mydb", schema="public", table="orders")