diff --git a/CHANGELOG.md b/CHANGELOG.md index a85dd12..606d83c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,68 @@ 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 — 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. 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. + + 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 + `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 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 ### Fixed diff --git a/hotdata_framework/client.py b/hotdata_framework/client.py index ff6570d..8716bca 100644 --- a/hotdata_framework/client.py +++ b/hotdata_framework/client.py @@ -77,8 +77,17 @@ _INDEX_TYPES = frozenset(get_args(IndexType)) _VECTOR_METRICS = frozenset(get_args(VectorMetric)) -_TERMINAL = frozenset({"succeeded", "failed", "cancelled"}) -_RESULT_FAILURE = frozenset({"failed", "cancelled"}) +# 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"}) @@ -918,7 +927,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 in _RUN_TERMINAL: return last time.sleep(interval_s) raise TimeoutError( 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..8ed9627 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -10,12 +10,12 @@ 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 ResultNotReadyError from hotdata.arrow import ResultsApi as ArrowResultsApi from hotdata.models.async_query_response import AsyncQueryResponse from hotdata.models.query_request import QueryRequest @@ -32,16 +32,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. @@ -130,70 +120,120 @@ 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 - ) - - 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``. - """ + arrow = ArrowResultsApi(self._runtime.api) 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") + 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( - 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) - return None + # 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) + # 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. + + 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 + last_status: str | None = None + 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) + 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. + # `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": {warning}" if warning else "") + ) + 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") + # 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 " + f"{self._QUERY_TIMEOUT_SECONDS}s (last status: {last_status})" ) - 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_client.py b/tests/test_client.py index 7ead972..741a374 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -245,20 +245,67 @@ 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_a_failed_result(): 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="failed", error_message="out of memory") with ( patch.object(client, "_results_api", return_value=FakeResultsApi()), - pytest.raises(RuntimeError, match="cancelled"), + pytest.raises(RuntimeError, match="out of memory"), ): client._wait_result_ready("res_1", timeout_s=0.1, interval_s=0) +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") + + class FakeResultsApi: + def get_result(self, result_id: str): + return SimpleNamespace(status="something_new", error_message=None) + + with ( + patch.object(client, "_results_api", return_value=FakeResultsApi()), + pytest.raises(TimeoutError, match="something_new"), + ): + client._wait_result_ready("res_1", timeout_s=0.05, 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 b7eed9b..3b54c05 100644 --- a/tests/test_managed_client.py +++ b/tests/test_managed_client.py @@ -7,11 +7,14 @@ 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.models.query_run_info import QueryRunInfo 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: @@ -28,13 +31,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 +65,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 +85,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 +102,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 +131,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 +153,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 +187,232 @@ 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_a_succeeded_run_that_saved_no_result_raises_rather_than_reading_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """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] = [] + + 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, + warning_message="result row creation failed; result not persisted", + ) + + 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) + 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="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``. @@ -415,9 +668,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 +679,291 @@ 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 + # 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( + 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", + ] + + +def test_unknown_run_status_keeps_polling_and_the_timeout_names_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """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. + """ + + 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="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() + monkeypatch.setattr(client, "_QUERY_TIMEOUT_SECONDS", 0.05) + + # 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") + + +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 + + +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 + + +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")