|
10 | 10 | import random |
11 | 11 | import time |
12 | 12 | from collections.abc import Callable |
13 | | -from typing import Any, Protocol, TypeVar |
| 13 | +from typing import Any, TypeVar |
14 | 14 |
|
15 | 15 | import pyarrow as pa |
16 | 16 | from hotdata.api.query_api import QueryApi |
17 | 17 | from hotdata.api.query_runs_api import QueryRunsApi |
18 | | -from hotdata.api.results_api import ResultsApi |
| 18 | +from hotdata.arrow import ResultNotReadyError |
19 | 19 | from hotdata.arrow import ResultsApi as ArrowResultsApi |
20 | 20 | from hotdata.models.async_query_response import AsyncQueryResponse |
21 | 21 | from hotdata.models.query_request import QueryRequest |
|
32 | 32 | T = TypeVar("T") |
33 | 33 |
|
34 | 34 |
|
35 | | -class _StatusResponse(Protocol): |
36 | | - """Async resources (query runs, results) expose a status and error message.""" |
37 | | - |
38 | | - status: str |
39 | | - error_message: str | None |
40 | | - |
41 | | - |
42 | | -S = TypeVar("S", bound=_StatusResponse) |
43 | | - |
44 | | - |
45 | 35 | class ManagedDatabaseClient: |
46 | 36 | """Managed-database client with bounded retries over hotdata-framework. |
47 | 37 |
|
@@ -130,70 +120,120 @@ def _fetch_result_arrow(self, result_id: str, *, database_id: str) -> pa.Table: |
130 | 120 | 0.6.0 SDK exposes (and requires) ``x_database_id`` on the Arrow |
131 | 121 | helper directly. |
132 | 122 | """ |
133 | | - return ArrowResultsApi(self._runtime.api).get_result_arrow( |
134 | | - result_id, x_database_id=database_id |
135 | | - ) |
136 | | - |
137 | | - def _poll( |
138 | | - self, |
139 | | - fetch: Callable[[], S], |
140 | | - *, |
141 | | - is_ready: Callable[[S], bool], |
142 | | - describe: str, |
143 | | - ) -> S: |
144 | | - """Poll ``fetch`` until ``is_ready`` is satisfied, or raise on failure/timeout. |
145 | | -
|
146 | | - ``failed``/``cancelled`` statuses raise ``RuntimeError``; exceeding |
147 | | - :attr:`_QUERY_TIMEOUT_SECONDS` raises ``TimeoutError``. |
148 | | - """ |
| 123 | + arrow = ArrowResultsApi(self._runtime.api) |
149 | 124 | deadline = time.monotonic() + self._QUERY_TIMEOUT_SECONDS |
150 | | - while time.monotonic() < deadline: |
151 | | - obj = fetch() |
152 | | - if obj.status in ("failed", "cancelled"): |
153 | | - raise RuntimeError(obj.error_message or f"{describe} {obj.status}") |
154 | | - if is_ready(obj): |
155 | | - return obj |
156 | | - time.sleep(self._POLL_INTERVAL_SECONDS) |
157 | | - raise TimeoutError(f"{describe} timed out after {self._QUERY_TIMEOUT_SECONDS}s") |
| 125 | + while True: |
| 126 | + try: |
| 127 | + return arrow.get_result_arrow(result_id, x_database_id=database_id) |
| 128 | + except ResultNotReadyError: |
| 129 | + # Waiting on the run should already have made this unreachable: |
| 130 | + # a run reports `succeeded` only once its result is saved and |
| 131 | + # ready. Tolerating it anyway costs nothing and removes the need |
| 132 | + # to take that ordering on trust. The Arrow endpoint answers a |
| 133 | + # result that is not ready with a small refusal rather than with |
| 134 | + # data, so waiting here is cheap in the way waiting on the JSON |
| 135 | + # result body -- which is what this change removed -- is not. |
| 136 | + if time.monotonic() >= deadline: |
| 137 | + raise |
| 138 | + time.sleep(self._POLL_INTERVAL_SECONDS) |
158 | 139 |
|
159 | 140 | def _query_database_scoped(self, sql: str, *, database_id: str) -> str | None: |
160 | 141 | raw = QueryApi(self._runtime.api).query( |
161 | | - QueryRequest(sql=sql), |
| 142 | + # Asked asynchronously because this caller wants a result id, not |
| 143 | + # rows. A synchronous submit always builds an inline preview of the |
| 144 | + # result and sends it -- megabytes, on a path that then reads the |
| 145 | + # whole result as Arrow anyway and never looks at the preview. The |
| 146 | + # async reply carries a run id and nothing else, and there is no way |
| 147 | + # to suppress the preview on a synchronous one. |
| 148 | + # |
| 149 | + # It also settles the types: the preview is JSON, which has no Arrow |
| 150 | + # schema and renders non-finite floats as null, so it could not have |
| 151 | + # substituted for the Arrow fetch even when it holds every row. |
| 152 | + # |
| 153 | + # `var_async` is the generated SDK's spelling of the wire field |
| 154 | + # `async`, which is a Python keyword and so cannot be the attribute |
| 155 | + # name. |
| 156 | + QueryRequest(sql=sql, var_async=True), |
162 | 157 | x_database_id=database_id, |
163 | 158 | ) |
164 | | - if isinstance(raw, QueryResponse): |
165 | | - # A synchronous response still persists its full result out-of-band |
166 | | - # under ``result_id``; that result may be ``processing`` when the |
167 | | - # inline preview returns, so wait for ``ready`` before the caller |
168 | | - # fetches it as Arrow. |
169 | | - return self._wait_result_ready(raw.result_id, database_id=database_id) |
170 | | - if isinstance(raw, AsyncQueryResponse): |
171 | | - run_result = self._await_query_run(raw.query_run_id, database_id=database_id) |
172 | | - return self._wait_result_ready(run_result, database_id=database_id) |
173 | | - return None |
| 159 | + # Both reply shapes carry `query_run_id`, and the run is the readiness |
| 160 | + # signal for either -- a synchronous reply (which `async_after_ms` can |
| 161 | + # still produce) returns rows inline but goes on saving the full result |
| 162 | + # in the background, so it is not the finish line either. |
| 163 | + if isinstance(raw, (QueryResponse, AsyncQueryResponse)): |
| 164 | + return self._await_query_run(raw.query_run_id, database_id=database_id) |
| 165 | + # Returning nothing here would read as an empty table: `fetch_table` |
| 166 | + # answers `None`, `fetch_table_rows` turns that into `[]`, and a |
| 167 | + # read-modify-write load would write only its new batch over rows it |
| 168 | + # believed were not there. A reply shape this client does not know is a |
| 169 | + # reason to stop, not to report emptiness. `HotdataClient` raises on the |
| 170 | + # same condition. |
| 171 | + raise RuntimeError(f"Unexpected query response type: {type(raw)!r}") |
174 | 172 |
|
175 | 173 | def _await_query_run(self, query_run_id: str, *, database_id: str) -> str | None: |
| 174 | + """Wait for a query run to finish; return the result id it produced. |
| 175 | +
|
| 176 | + The run is the whole wait. A run turns `succeeded` only after its result |
| 177 | + has been saved and is `ready`, so `succeeded` needs no second check |
| 178 | + against the result -- and asking the result endpoint instead would mean |
| 179 | + downloading the entire result to read one field, which the server |
| 180 | + refuses outright (413/429) once the result is large enough. |
| 181 | +
|
| 182 | + `result_id` comes off the run rather than off the query reply because a |
| 183 | + `succeeded` run reports none when every row came back inline but the |
| 184 | + result could not be saved for later retrieval. |
| 185 | + """ |
176 | 186 | runs = QueryRunsApi(self._runtime.api) |
177 | | - run = self._poll( |
| 187 | + deadline = time.monotonic() + self._QUERY_TIMEOUT_SECONDS |
| 188 | + last_status: str | None = None |
| 189 | + while time.monotonic() < deadline: |
178 | 190 | # Runs (like results) of database-scoped queries are database-scoped. |
179 | | - lambda: runs.get_query_run(query_run_id, x_database_id=database_id), |
180 | | - is_ready=lambda r: r.status == "succeeded", |
181 | | - describe="Query", |
182 | | - ) |
183 | | - return run.result_id |
184 | | - |
185 | | - def _wait_result_ready(self, result_id: str | None, *, database_id: str) -> str | None: |
186 | | - if result_id is None: |
187 | | - return None |
188 | | - results = ResultsApi(self._runtime.api) |
189 | | - self._poll( |
190 | | - # The stored result of a database-scoped query 400s without the |
191 | | - # database scope. |
192 | | - lambda: results.get_result(result_id, x_database_id=database_id), |
193 | | - is_ready=lambda r: r.status == "ready", |
194 | | - describe=f"Result {result_id}", |
| 191 | + run = runs.get_query_run(query_run_id, x_database_id=database_id) |
| 192 | + last_status = run.status |
| 193 | + if run.status == "succeeded": |
| 194 | + if run.result_id is None: |
| 195 | + # A run succeeds with no result id when its rows were |
| 196 | + # returned inline but the result could not be saved for |
| 197 | + # later retrieval. Returning nothing here would surface as |
| 198 | + # an empty table -- `fetch_table` answers `None`, and |
| 199 | + # `fetch_table_rows` turns that into `[]`, which is the same |
| 200 | + # answer it gives for a table that does not exist. A |
| 201 | + # read-modify-write load would then read no existing rows |
| 202 | + # and write only its new batch, dropping what was there. |
| 203 | + # Terminal rather than transient: re-running the query |
| 204 | + # cannot save a result that was already discarded. |
| 205 | + # `getattr` because this runs while building an error: if |
| 206 | + # the field ever goes away, losing the explanation is a far |
| 207 | + # better outcome than an AttributeError replacing the raise. |
| 208 | + warning = getattr(run, "warning_message", None) |
| 209 | + raise RuntimeError( |
| 210 | + f"Query run {query_run_id} succeeded but its result was not " |
| 211 | + f"saved, so the table cannot be read" |
| 212 | + + (f": {warning}" if warning else "") |
| 213 | + ) |
| 214 | + return run.result_id |
| 215 | + if run.status == "interrupted": |
| 216 | + # Terminal, but the server lost the run rather than rejecting |
| 217 | + # the query, so it is the one failure here worth re-running. |
| 218 | + # Raised pre-classified: `classify_sdk_error` cannot tell this |
| 219 | + # apart from an ordinary RuntimeError and would call it terminal. |
| 220 | + raise HotdataTransientError( |
| 221 | + run.error_message or f"Query run {query_run_id} was interrupted" |
| 222 | + ) |
| 223 | + if run.status == "failed": |
| 224 | + raise RuntimeError(run.error_message or f"Query run {query_run_id} failed") |
| 225 | + # Any other status keeps polling, including one this client has never |
| 226 | + # seen. Treating an unrecognised status as terminal is the cheaper |
| 227 | + # failure to diagnose and by far the more expensive one to suffer: a |
| 228 | + # single status added upstream would then fail every query at once, |
| 229 | + # where waiting costs one slow call. What made `interrupted` |
| 230 | + # expensive was not the waiting, it was that the timeout never said |
| 231 | + # which status it had waited on -- so the message now carries it. |
| 232 | + time.sleep(self._POLL_INTERVAL_SECONDS) |
| 233 | + raise TimeoutError( |
| 234 | + f"Query run {query_run_id} did not finish within " |
| 235 | + f"{self._QUERY_TIMEOUT_SECONDS}s (last status: {last_status})" |
195 | 236 | ) |
196 | | - return result_id |
197 | 237 |
|
198 | 238 | def fetch_table_rows(self, *, database: str, schema: str, table: str) -> list[dict[str, Any]]: |
199 | 239 | result = self.fetch_table(database=database, schema=schema, table=table) |
|
0 commit comments