diff --git a/nodes/src/nodes/db_hotdata/IGlobal.py b/nodes/src/nodes/db_hotdata/IGlobal.py index d07b5127bb..9b9933a343 100644 --- a/nodes/src/nodes/db_hotdata/IGlobal.py +++ b/nodes/src/nodes/db_hotdata/IGlobal.py @@ -46,6 +46,12 @@ from .hotdata_client import HotdataClient +#: Smallest async_after_ms the API accepts. Anything lower is answered with +#: 400 "async_after_ms must be at least 1000", which failed every single query +#: for a pipeline configured below it. services.json advertises the same floor. +ASYNC_AFTER_MS_MIN = 1000 + + def _int_or(value: Any, default: int, *, lo: int, hi: int) -> int: try: n = int(value) @@ -94,6 +100,7 @@ class IGlobal(IGlobalBase): async_after_ms: int = 5000 def beginGlobal(self) -> None: + """Read the node config and build the REST client; the database itself is created lazily.""" if self.IEndpoint.endpoint.openMode == OPEN_MODE.CONFIG: return @@ -128,7 +135,7 @@ def beginGlobal(self) -> None: self.allow_execute = _bool_or(cfg.get('allow_execute'), False) self.allow_destructive_load = _bool_or(cfg.get('allow_destructive_load'), False) self.job_timeout_secs = _int_or(cfg.get('job_timeout_secs'), 300, lo=10, hi=3600) - self.async_after_ms = _int_or(cfg.get('async_after_ms'), 5000, lo=0, hi=60000) + self.async_after_ms = _int_or(cfg.get('async_after_ms'), 5000, lo=ASYNC_AFTER_MS_MIN, hi=60000) self.client = HotdataClient( apikey=self.apikey, diff --git a/nodes/src/nodes/db_hotdata/IInstance.py b/nodes/src/nodes/db_hotdata/IInstance.py index e1e98fa5d9..f90730034b 100644 --- a/nodes/src/nodes/db_hotdata/IInstance.py +++ b/nodes/src/nodes/db_hotdata/IInstance.py @@ -79,7 +79,13 @@ _POLL_MAX_S = 5.0 _TERMINAL_OK = frozenset({'succeeded', 'success', 'completed', 'complete', 'ready', 'finished'}) -_TERMINAL_BAD = frozenset({'failed', 'error', 'cancelled', 'canceled'}) +_TERMINAL_BAD = frozenset({'failed', 'error', 'interrupted', 'cancelled', 'canceled'}) +#: Terminal, but not the statement's fault, so raised as a plain RuntimeError and +#: never as SqlStatementError. Only `interrupted` is a documented status of a +#: query run (https://www.hotdata.dev/openapi.yaml: "terminal, and safe to +#: retry"); `cancelled` / `canceled` were never seen from this endpoint and are +#: kept from the original set as a guard. +_ENDED_NOT_BY_THE_SQL = frozenset({'interrupted', 'cancelled', 'canceled'}) #: Async job states from the loads/indexes 202 envelope. _JOB_PENDING = frozenset({'pending', 'running'}) @@ -194,6 +200,64 @@ def _to_ndjson(rows: List[Any]) -> bytes: return ('\n'.join(lines) + '\n').encode('utf-8') +class SqlStatementError(RuntimeError): + """The server ran the statement and the statement failed. + + Distinct from a transport or account failure: this one names something in the + SQL, so regenerating is worth a turn. Carried as its own type because a + deferred run reports its failure in the poll body, with no HTTP status to + classify on. + """ + + +def _run_failure_text(run: Dict[str, Any], fallback: str) -> str: + """Why a query run failed, in the server's words. + + The run body reports it as ``error_message`` ("query execution failed: Arrow + error: Divide by zero error"). This text is what gets fed back to the model + for its corrective attempt, so falling through to the bare status would send + it "failed" and leave it rewriting blind. ``error`` and ``message`` are read + too - ``error`` as either a string or the ``{"message": ...}`` object the + query endpoint uses - so a body shaped like the other endpoints still works. + """ + for key in ('error_message', 'error', 'message'): + value = run.get(key) + if isinstance(value, dict): + value = value.get('message') + if value: + return str(value) + return fallback + + +def _is_sql_fixable(error: Exception) -> bool: + """Could a *different statement* plausibly succeed where this one failed? + + Only worth another generation turn when the server rejected the SQL itself - + a missing table, a parse error, an unknown function all come back 400. A + failure that describes the run rather than the statement (401 invalid_api_key, + 404 workspace_not_found or database not found, an exhausted 429 budget, a + connection that never landed) is identical on every retry, so re-asking the + model only spends another LLM call and another round trip before reporting + the same error with the cause buried behind "could not answer after N + attempts". + + The status alone is too coarse: ``_run_sql`` drives four endpoints, and a + bad ttl on create_database, a malformed limit on get_result or a missing + scoping header all answer 400 without a single thing being wrong with the + SQL. The server marks a real statement failure by answering with a + ``query_run_id`` - it only mints one once it has created a run and executed + the statement - so that, not the bare 400, is the test. + + Defaults to *not* fixable: a failure this code cannot classify is far more + likely to be environmental than to be a statement the model can rewrite. + """ + if isinstance(error, SqlStatementError): + return True + if getattr(error, 'status_code', None) != 400: + return False + return bool(getattr(error, 'query_run_id', '')) + + def _is_missing_column(error: Exception) -> bool: """Is this the server refusing a write that omits a column the table has? @@ -309,8 +373,90 @@ def _rows_from_payload(payload: Any) -> List[Dict[str, Any]]: return [] +def _free_name(base: str, used: set[str]) -> str: + """``base``, suffixed until it is a key nothing has claimed, then reserved. + + The generated name has to be checked against the names already taken, not + just against the base: for columns ``['city', 'city', 'city_1']`` the second + ``city`` would otherwise become ``city_1`` and collide with the real third + column. + """ + name = base + suffix = 1 + while name in used: + name = f'{base}_{suffix}' + suffix += 1 + used.add(name) + return name + + +def _name_columns(columns: List[Any]) -> List[str]: + """Column names for a result, made unique and non-empty. + + A join comes back with the same name twice: ``SELECT a.city, b.city`` answers + ``['city', 'city']`` and ``SELECT *`` over a self-join answers ``['city', + 'units', 'city', 'units']`` (the planner rejects a bare ``SELECT city, city``, + so joins are how this is reached). Zipping that into a dict would keep only + the last value and silently drop a column. + """ + used: set[str] = set() + named: List[str] = [] + for index, raw in enumerate(columns): + name = str(raw).strip() if raw is not None else '' + if not name: + name = f'column_{index + 1}' + named.append(_free_name(name, used)) + return named + + +def _rows_as_objects(rows: Any, columns: Any) -> List[Any]: + """Pair positional result rows with the column names returned beside them. + + ``/v1/query`` and ``/v1/results`` both answer with ``columns`` (the names) and + ``rows`` (a list of value *lists*), never with objects. Everything downstream + is written against objects - the Markdown table, the information_schema + reshaping, the rows handed to the agent - so the names are attached here, at + the one boundary where both halves are in hand. Without this the rows keep + their positions and nothing can name a field: the table renderer falls + through to ``str(row)`` and emits Python list reprs. + + Rows that are already objects are passed through, so a server that starts + returning objects, and the tests that mock them, keep working. A row + *narrower* than its header yields an object without the trailing keys rather + than inventing nulls for values the server never sent. + """ + if not isinstance(rows, list) or not rows: + return rows if isinstance(rows, list) else [] + if not isinstance(columns, list) or not columns: + return rows + names = _name_columns(columns) + out: List[Any] = [] + for row in rows: + if isinstance(row, dict) or not isinstance(row, (list, tuple)): + out.append(row) + continue + item: Dict[str, Any] = {name: row[i] for i, name in enumerate(names) if i < len(row)} + # A row wider than its header keeps the tail rather than dropping it: + # losing a value silently is worse than naming it by position. The + # positional key goes through _free_name too, because `column_3` is + # itself a legal column name and would otherwise be overwritten. + if len(row) > len(names): + used = set(item) + for i in range(len(names), len(row)): + item[_free_name(f'column_{i + 1}', used)] = row[i] + out.append(item) + return out + + def _cell(value: Any) -> str: - """Render one table cell: pipes and newlines both break the row otherwise.""" + """Render one table cell: pipes and newlines both break the row otherwise. + + SQL NULL arrives as ``None`` and renders as an empty cell. ``str(None)`` would + put the Python word "None" in front of a reader who asked a question about + their data - the same leak as the list reprs, one value at a time. + """ + if value is None: + return '' return str(value).replace('|', '\\|').replace('\r\n', ' ').replace('\n', ' ').replace('\r', ' ') @@ -328,7 +474,9 @@ def _rows_to_markdown(rows: List[Any], limit: int = 100) -> str: if key not in columns: columns.append(key) - header = '| ' + ' | '.join(columns) + ' |' + # Header cells are escaped like body cells: `SELECT total AS "a|b"` is a + # legal alias, and an unescaped pipe in the header splits the column in two. + header = '| ' + ' | '.join(_cell(c) for c in columns) + ' |' divider = '| ' + ' | '.join('---' for _ in columns) + ' |' body = ['| ' + ' | '.join(_cell(row.get(c, '')) for c in columns) + ' |' for row in shown] table = '\n'.join([header, divider] + body) @@ -381,19 +529,26 @@ def _run_sql(self, sql: str, limit: int) -> Dict[str, Any]: # 25,000 - so the id is only offered when the caller has seen all of it. run_id = response.get('query_run_id') or response.get('id') if response.get('rows') is not None and not response.get('truncated'): - rows = response.get('rows') or [] + rows = _rows_as_objects(response.get('rows') or [], response.get('columns')) result = {'rows': rows[:limit], 'row_count': len(rows[:limit]), 'sql': sql} if response.get('result_id') and len(rows) <= limit: result['result_id'] = response['result_id'] return result - if run_id and response.get('result_id') is None: - response = self._await_run(run_id) + # A truncated response always carries a result_id, but per QueryResponse in + # https://www.hotdata.dev/openapi.yaml that id "is issued while the save is + # still in flight, so on its own it does not mean the result can be + # retrieved": a read against it answers 202 with no rows, which came back + # from here as an empty result. The run is the readiness signal, so wait + # on it whenever the server deferred or truncated. + if run_id and (response.get('result_id') is None or response.get('truncated')): + response = self._await_run(run_id, database_id) result_id = response.get('result_id') if result_id: - payload = glb.client.get_result(result_id, offset=0, limit=limit) - rows = payload.get('rows') or payload.get('data') or [] + payload = glb.client.get_result(result_id, database_id=database_id, offset=0, limit=limit) + raw = payload.get('rows') or payload.get('data') or [] + rows = _rows_as_objects(raw, payload.get('columns')) result = {'rows': rows[:limit], 'row_count': len(rows[:limit]), 'sql': sql} # get_result is explicitly a window; a full page back means there may # be more behind it, and the id would then name more than was seen. @@ -401,7 +556,7 @@ def _run_sql(self, sql: str, limit: int) -> Dict[str, Any]: result['result_id'] = result_id return result - rows = response.get('rows') or [] + rows = _rows_as_objects(response.get('rows') or [], response.get('columns')) return {'rows': rows[:limit], 'row_count': len(rows[:limit]), 'sql': sql} def _schema_via_sql(self) -> List[Dict[str, Any]]: @@ -457,17 +612,23 @@ def _table_columns(self, schema: str, table: str) -> List[str]: names.append(str(value)) return names - def _await_run(self, run_id: str) -> Dict[str, Any]: + def _await_run(self, run_id: str, database_id: str = '') -> Dict[str, Any]: """Poll a query run to a terminal state under a monotonic deadline.""" glb = self.IGlobal deadline = time.monotonic() + glb.job_timeout_secs delay = _POLL_BASE_S while True: - run = glb.client.get_query_run(run_id) + run = glb.client.get_query_run(run_id, database_id=database_id) status = str(run.get('status') or '').lower() if status in _TERMINAL_BAD: - message = run.get('error') or run.get('message') or status - raise RuntimeError(f'db_hotdata: query failed: {message}') + message = _run_failure_text(run, status) + if status in _ENDED_NOT_BY_THE_SQL: + # Not the statement's fault, so not SqlStatementError: a + # rewritten query cannot un-interrupt or un-cancel a run, and + # classifying it as fixable would spend an LLM turn finding + # that out. + raise RuntimeError(f'db_hotdata: query run ended as {status!r}: {message}') + raise SqlStatementError(f'db_hotdata: query failed: {message}') if status in _TERMINAL_OK or run.get('result_id'): return run if time.monotonic() + delay > deadline: @@ -888,6 +1049,7 @@ def get_sql(self, args: Any) -> Dict[str, Any]: ), ) def get_data(self, args: Any) -> Dict[str, Any]: + """Answer a question: generate SQL, run it, and retry a rejected statement with its error fed back.""" args = normalize_tool_input(args, tool_name='get_data') question_text = str(args.get('question') or '').strip() if not question_text: @@ -913,6 +1075,12 @@ def get_data(self, args: Any) -> Dict[str, Any]: result['attempts'] = attempt return result except Exception as e: + if not _is_sql_fixable(e): + # Raised as-is: the original carries the status and the + # server's own wording, which "could not answer after 3 + # attempts" would bury behind two pointless LLM calls. + debug(f'db_hotdata: attempt {attempt} hit a failure no rewrite can fix: {e}') + raise previous_sql, last_error = cleaned, str(e) debug(f'db_hotdata: attempt {attempt} failed: {last_error}') diff --git a/nodes/src/nodes/db_hotdata/README.md b/nodes/src/nodes/db_hotdata/README.md index 710144a0ec..cff02a7a00 100644 --- a/nodes/src/nodes/db_hotdata/README.md +++ b/nodes/src/nodes/db_hotdata/README.md @@ -36,7 +36,7 @@ An agent flow: `load_data` the rows, `build_index` on the text column, then `get | Function | Description | | ------------- | ------------------------------------------------------------------------------------------------------------- | | `load_data` | Load rows, or a previous query result by `result_id`, into a table. Creates the table if needed | -| `get_data` | Answer a plain-language question. The bound LLM writes the SQL; failures are retried with the error fed back | +| `get_data` | Answer a plain-language question. The bound LLM writes the SQL; only rejected SQL is retried, with its error | | `get_sql` | Generate the SQL for a question without running it | | `execute` | Run one raw read-only SQL statement. Gated by `allow_execute` | | `get_schema` | Live tables and columns from `information_schema` | @@ -125,6 +125,7 @@ Hotdata runs **Apache DataFusion 54 behind the PostgreSQL parser dialect**. Trea | Symptom | Cause | | -------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- | | `apikey is required` / `workspace_id is required` | Neither the config field nor the env var is set | +| `HTTP 401` / `HTTP 404` from `get_data` | Wrong API key, workspace or database ID. Raised at once: only a rejected SQL statement is retried | | `raw SQL execution is disabled` | Turn on `allow_execute`, or use `get_data` instead | | `only one statement per call` | Hotdata rejects semicolon-separated batches | | Unknown function errors | A Postgres-only function that DataFusion lacks — check `dialect` | diff --git a/nodes/src/nodes/db_hotdata/hotdata_client.py b/nodes/src/nodes/db_hotdata/hotdata_client.py index 6f8d4ebf3b..a86e55db08 100644 --- a/nodes/src/nodes/db_hotdata/hotdata_client.py +++ b/nodes/src/nodes/db_hotdata/hotdata_client.py @@ -80,6 +80,10 @@ #: which must NOT be retried. _RESOURCE_LOCKED = 'RESOURCE_LOCKED' +#: States in which GET /v1/results/{id} answers 202 with {status, result_id} and +#: no data, per the endpoint's status table in https://www.hotdata.dev/openapi.yaml. +_RESULT_NOT_READY = frozenset({'pending', 'processing'}) + def _creates_a_database(method: str, path: str) -> bool: """Is this the one request whose replay could orphan a billable resource?""" @@ -114,10 +118,22 @@ class HotdataError(RuntimeError): treat as success, and "this table is busy", which they must not. """ - def __init__(self, message: str, status_code: Optional[int] = None, error_code: str = '') -> None: + def __init__( + self, + message: str, + status_code: Optional[int] = None, + error_code: str = '', + query_run_id: str = '', + ) -> None: + """Keep the HTTP status, API error code and query run id beside the message.""" super().__init__(message) self.status_code = status_code self.error_code = error_code + #: Set when the server answered the error with a query run id, which it + #: does only once it has created a run and executed the statement. That + #: makes it the marker for "the SQL failed" as opposed to "the request + #: was wrong" - both of which are plain 400s otherwise. + self.query_run_id = query_run_id class HotdataOverloadedError(HotdataError): @@ -294,6 +310,7 @@ def _request( f'hotdata: {method} {path} failed with HTTP {status}: {_body_snippet(response)}', status_code=status, error_code=error_code, + query_run_id=_query_run_id(response), ) return _parse_json(response) @@ -361,14 +378,51 @@ def query( body['default_schema'] = default_schema return self._request('POST', '/v1/query', json_body=body) - def get_query_run(self, query_run_id: str) -> Dict[str, Any]: - return self._request('GET', f'/v1/query-runs/{query_run_id}') + def get_query_run(self, query_run_id: str, database_id: str = '') -> Dict[str, Any]: + """Read one query run. + + Scoped to a database: without ``X-Database-Id`` the server answers 400 + ``"this endpoint is scoped to a database"``, so every poll of a deferred + query failed. The id is threaded through from the caller rather than read + off the global, because the run belongs to the database it was issued + against, not to whichever database the node happens to hold now. + """ + return self._request( + 'GET', + f'/v1/query-runs/{query_run_id}', + extra_headers={'X-Database-Id': database_id} if database_id else None, + ) + + def get_result( + self, result_id: str, offset: int = 0, limit: Optional[int] = None, *, database_id: str = '' + ) -> Dict[str, Any]: + """Read a window of a stored result. Database-scoped, as ``get_query_run`` is. - def get_result(self, result_id: str, offset: int = 0, limit: Optional[int] = None) -> Dict[str, Any]: + ``database_id`` is keyword-only and sits after the original parameters, so + a positional ``get_result(result_id, offset, limit)`` call written against + the old signature still means what it meant. + """ params: Dict[str, Any] = {'offset': offset} if limit is not None: params['limit'] = limit - return self._request('GET', f'/v1/results/{result_id}', params=params) + body = self._request( + 'GET', + f'/v1/results/{result_id}', + params=params, + extra_headers={'X-Database-Id': database_id} if database_id else None, + ) + # A result still being saved answers 202 with {status, result_id} and no + # rows. _request only raises from 400 up, so without this the caller reads + # "no rows" and reports that the query matched nothing. Not ready is not + # the same as empty. + state = str(body.get('status') or '').lower() + if state in _RESULT_NOT_READY and body.get('rows') is None: + raise HotdataError( + f'hotdata: result {result_id} is not ready yet (status {state!r}); ' + 'wait for its query run to succeed before reading it', + status_code=202, + ) + return body def information_schema( self, @@ -664,6 +718,29 @@ def _error_code(response: Any) -> str: return str(body.get('code') or '') +def _query_run_id(response: Any) -> str: + """The ``query_run_id`` on an error body, or '' if there isn't one. + + Only the query endpoint sets it, and only after it has run the statement, so + its presence separates "your SQL is wrong" - which a regenerated statement + could fix - from "your request is wrong" (a bad ttl, a missing header, an + out-of-range limit), which is identical on every attempt. + """ + try: + body = response.json() + except Exception: + return '' + if not isinstance(body, dict): + return '' + # Both levels, the way _error_code reads the code: the documented error body + # is {"error": {code, message}}, so the run id is looked for inside `error` + # as well as beside it, where the live API puts it today. + error = body.get('error') + if isinstance(error, dict) and error.get('query_run_id'): + return str(error['query_run_id']) + return str(body.get('query_run_id') or '') + + def _body_snippet(response: Any, limit: int = 500) -> str: """A short, safe excerpt of an error body for log and exception text.""" try: diff --git a/nodes/src/nodes/db_hotdata/services.json b/nodes/src/nodes/db_hotdata/services.json index f0102550a3..9d23d30dab 100644 --- a/nodes/src/nodes/db_hotdata/services.json +++ b/nodes/src/nodes/db_hotdata/services.json @@ -129,7 +129,7 @@ "title": "Run asynchronously after (milliseconds)", "description": "Wait this long before a query continues as an asynchronous job.", "default": 5000, - "minimum": 0, + "minimum": 1000, "maximum": 60000 } }, diff --git a/nodes/test/test_db_hotdata.py b/nodes/test/test_db_hotdata.py index 5d8bb04717..752df5ea45 100644 --- a/nodes/test/test_db_hotdata.py +++ b/nodes/test/test_db_hotdata.py @@ -735,24 +735,34 @@ def _query(**kw): def test_execute_follows_async_run_to_a_result(): + """A deferred query is polled to completion and its result read, both scoped to the database.""" polls = [{'status': 'running'}, {'status': 'succeeded', 'result_id': 'res-1'}] + scoped = [] g = _global() g.client = SimpleNamespace( query=lambda **_kw: {'query_run_id': 'run-1'}, - get_query_run=lambda _i: polls.pop(0), - get_result=lambda _i, offset=0, limit=None: {'rows': [{'a': 1}]}, + get_query_run=lambda _i, database_id='': scoped.append(('run', database_id)) or polls.pop(0), + get_result=lambda _i, database_id='', offset=0, limit=None: ( + scoped.append(('result', database_id)) or {'rows': [{'a': 1}]} + ), create_database=lambda **_kw: {'id': 'db-1'}, ) inst = _instance(g) out = inst.execute({'sql': 'SELECT a FROM t'}) assert out['rows'] == [{'a': 1}] + # Both endpoints are database-scoped server-side; an unscoped call is a 400. + assert scoped == [('run', 'db-1'), ('run', 'db-1'), ('result', 'db-1')] def test_failed_query_run_raises_runtime_error(): + """A run that ends failed raises with the server's reason, read from ``error_message``.""" g = _global() g.client = SimpleNamespace( query=lambda **_kw: {'query_run_id': 'run-1'}, - get_query_run=lambda _i: {'status': 'failed', 'error': 'syntax error at or near "SELCT"'}, + get_query_run=lambda _i, database_id='': { + 'status': 'failed', + 'error_message': 'query execution failed: syntax error at or near "SELCT"', + }, create_database=lambda **_kw: {'id': 'db-1'}, ) inst = _instance(g) @@ -827,12 +837,20 @@ def test_get_sql_requires_a_question(): def test_get_data_retries_with_the_error_fed_back(): + """A rejected statement earns another attempt, and the model sees why the first one failed.""" calls = {'n': 0} def _query(**_kw): calls['n'] += 1 if calls['n'] == 1: - raise RuntimeError("Invalid function 'to_number'. Did you mean 'to_char'?") + # How the live API rejects a bad statement: HTTP 400 carrying a + # query_run_id, which is what marks it as the statement failing + # rather than the request being malformed. + raise client_mod.HotdataError( + "hotdata: POST /v1/query failed with HTTP 400: Invalid function 'to_number'.", + status_code=400, + query_run_id='qrun-1', + ) return {'rows': [{'a': 1}]} g = _loaded_global(max_attempts=3) @@ -847,10 +865,13 @@ def _query(**_kw): def test_get_data_gives_up_after_max_attempts(): + """Statement failures stop after max_attempts and report the last error.""" g = _loaded_global(max_attempts=2) g.client = SimpleNamespace( information_schema=lambda **_kw: {'tables': []}, - query=lambda **_kw: (_ for _ in ()).throw(RuntimeError('boom')), + query=lambda **_kw: (_ for _ in ()).throw( + client_mod.HotdataError('hotdata: HTTP 400: boom', status_code=400, query_run_id='qrun-1') + ), ) inst = _llm_instance(g, ['SELECT 1', 'SELECT 2']) with pytest.raises(RuntimeError, match='after 2 attempts'): @@ -2416,3 +2437,486 @@ def test_a_new_database_does_not_inherit_the_previous_dedup_records(): assert not out.get('deduplicated'), 'a different database must not reuse the old fingerprint' assert len(calls) == 2, 'the load must actually run against the new database' + + +# --------------------------------------------------------------------------- +# Result rows carry their column names +# +# /v1/query and /v1/results answer with `columns` (the names) and `rows` (lists +# of values), never with objects. Everything downstream is written against +# objects, so the names are attached in _run_sql. Before that, every lane +# rendered Python list reprs - "['Reno', 200]" - instead of a table. +# --------------------------------------------------------------------------- + + +def test_rows_are_named_from_the_columns_beside_them(): + """The live wire shape - `columns` beside positional `rows` - becomes a list of objects.""" + g = _loaded_global() + g.client = SimpleNamespace( + query=lambda **_kw: { + 'columns': ['city', 'units'], + 'rows': [['Reno', 200], ['Austin', 120], ['Denver', 80]], + 'row_count': 3, + 'truncated': False, + } + ) + out = _instance(g)._run_sql('SELECT city, units FROM main.sales', 100) + assert out['rows'] == [ + {'city': 'Reno', 'units': 200}, + {'city': 'Austin', 'units': 120}, + {'city': 'Denver', 'units': 80}, + ] + + +def test_named_rows_render_as_a_markdown_table(): + """The lane output the three bugs were reported against.""" + rows = iinstance_mod._rows_as_objects([['Reno', 200], ['Austin', 120]], ['city', 'units']) + table = iinstance_mod._rows_to_markdown(rows) + assert table.splitlines()[0] == '| city | units |' + assert '| Reno | 200 |' in table + assert '[' not in table, f'must not fall through to a Python list repr: {table!r}' + + +def test_a_repeated_column_name_does_not_collapse(): + """A join answers ['city', 'city']; a plain dict would keep only the last value.""" + rows = iinstance_mod._rows_as_objects([['Reno', 'Austin']], ['city', 'city']) + assert rows == [{'city': 'Reno', 'city_1': 'Austin'}] + + +def test_an_unnamed_column_still_gets_a_key(): + """An empty or null column name falls back to its position.""" + assert iinstance_mod._rows_as_objects([[1, 2]], ['', None]) == [{'column_1': 1, 'column_2': 2}] + + +def test_a_row_wider_than_its_header_keeps_the_tail(): + """Dropping a value silently is worse than naming it by position.""" + assert iinstance_mod._rows_as_objects([[1, 2, 3]], ['a']) == [{'a': 1, 'column_2': 2, 'column_3': 3}] + + +def test_rows_that_are_already_objects_pass_through(): + """Rows that arrive as objects are left exactly as they are.""" + assert iinstance_mod._rows_as_objects([{'a': 1}], ['a']) == [{'a': 1}] + + +def test_rows_without_columns_are_left_alone(): + """With no names to attach there is nothing to do; rows keep their positions.""" + assert iinstance_mod._rows_as_objects([[1]], None) == [[1]] + + +def test_information_schema_over_sql_reshapes_named_rows(): + """_schema_via_sql skips any row that is not an object, so unnamed rows + silently produced an empty schema and the LLM wrote SQL with no columns. + """ + g = _loaded_global() + g.client = SimpleNamespace( + query=lambda **_kw: { + 'columns': ['table_schema', 'table_name', 'column_name', 'data_type'], + 'rows': [['main', 'sales', 'city', 'Utf8'], ['main', 'sales', 'units', 'Int64']], + 'truncated': False, + } + ) + tables = _instance(g)._schema_via_sql() + assert tables == [ + { + 'schema': 'main', + 'table': 'sales', + 'columns': [ + {'name': 'city', 'data_type': 'Utf8'}, + {'name': 'units', 'data_type': 'Int64'}, + ], + } + ] + + +# --------------------------------------------------------------------------- +# Database-scoped reads +# --------------------------------------------------------------------------- + + +def test_get_result_and_query_run_send_the_database_header(): + """Both endpoints answer 400 without X-Database-Id.""" + c, rec = _client([_Resp(200, {'rows': []}), _Resp(200, {'status': 'succeeded'})]) + c.get_result('res-1', database_id='db-1', offset=0, limit=10) + c.get_query_run('run-1', 'db-1') + assert rec.calls[0]['headers']['X-Database-Id'] == 'db-1' + assert rec.calls[1]['headers']['X-Database-Id'] == 'db-1' + + +def test_an_unscoped_result_read_sends_no_database_header(): + """The header is omitted rather than sent empty when no id is available.""" + c, rec = _client([_Resp(200, {'rows': []})]) + c.get_result('res-1') + assert 'X-Database-Id' not in rec.calls[0]['headers'] + + +# --------------------------------------------------------------------------- +# Retry only what a different statement could fix +# --------------------------------------------------------------------------- + + +def test_get_data_does_not_retry_a_failure_sql_cannot_fix(): + """A bad workspace is 404 on every attempt. Re-asking the model spends another + LLM call and another round trip to report the same thing. + """ + calls = {'n': 0} + + def _query(**_kw): + calls['n'] += 1 + raise client_mod.HotdataError( + 'hotdata: POST /v1/query failed with HTTP 404: {"error": "workspace_not_found"}', + status_code=404, + ) + + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace(information_schema=lambda **_kw: {'tables': []}, query=_query) + inst = _llm_instance(g, ['SELECT 1', 'SELECT 2', 'SELECT 3']) + + with pytest.raises(client_mod.HotdataError, match='workspace_not_found'): + inst.get_data({'question': 'x'}) + assert calls['n'] == 1, 'must stop at the first unfixable failure' + assert len(inst.asked) == 1, 'must not spend a second LLM call' + + +def test_get_data_does_not_retry_an_auth_failure(): + """A rejected API key is a 401 on every attempt, so it is raised after the first.""" + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace( + information_schema=lambda **_kw: {'tables': []}, + query=lambda **_kw: (_ for _ in ()).throw( + client_mod.HotdataError('hotdata: HTTP 401: invalid_api_key', status_code=401) + ), + ) + inst = _llm_instance(g, ['SELECT 1', 'SELECT 2', 'SELECT 3']) + with pytest.raises(client_mod.HotdataError, match='invalid_api_key'): + inst.get_data({'question': 'x'}) + assert len(inst.asked) == 1 + + +def test_get_data_still_retries_a_failed_query_run(): + """A deferred run reports its failure in the poll body, with no HTTP status - + it is still the statement that failed, so it still earns another turn. + """ + calls = {'n': 0} + + def _query(**_kw): + calls['n'] += 1 + if calls['n'] == 1: + return {'query_run_id': 'run-1'} + return {'columns': ['a'], 'rows': [[1]], 'truncated': False} + + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace( + information_schema=lambda **_kw: {'tables': []}, + query=_query, + get_query_run=lambda _i, database_id='': {'status': 'failed', 'error': 'Invalid function'}, + ) + inst = _llm_instance(g, ['SELECT bad(a) FROM t', 'SELECT a FROM t']) + out = inst.get_data({'question': 'x'}) + assert out['rows'] == [{'a': 1}] + assert out['attempts'] == 2 + + +def test_a_generated_suffix_cannot_collide_with_a_real_column(): + """For columns ['city', 'city', 'city_1'] the generated name for the second + `city` must not land on a column the result already has. + """ + names = iinstance_mod._name_columns(['city', 'city', 'city_1']) + assert len(set(names)) == 3, f'names collide: {names}' + rows = iinstance_mod._rows_as_objects([['a', 'b', 'c']], ['city', 'city', 'city_1']) + assert list(rows[0].values()) == ['a', 'b', 'c'], f'a value was dropped: {rows}' + + +def test_an_excess_value_cannot_overwrite_a_named_column(): + """A row wider than its header must not park a value on a key the header + already used - `column_3` is a legal column name. + """ + rows = iinstance_mod._rows_as_objects([[1, 2, 3]], ['a', 'column_3']) + assert list(rows[0].values()) == [1, 2, 3], f'a value was dropped: {rows}' + + +def test_a_400_without_a_query_run_id_is_not_retried(): + """`_run_sql` drives four endpoints. A bad ttl, a malformed limit or a + missing scoping header all answer 400 with nothing wrong with the SQL, and + the server mints no run id for them. + """ + calls = {'n': 0} + + def _query(**_kw): + calls['n'] += 1 + raise client_mod.HotdataError( + "hotdata: POST /v1/databases failed with HTTP 400: expires_at 'not-a-ttl' is not an RFC 3339 timestamp", + status_code=400, + ) + + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace(information_schema=lambda **_kw: {'tables': []}, query=_query) + inst = _llm_instance(g, ['SELECT 1', 'SELECT 2', 'SELECT 3']) + + with pytest.raises(client_mod.HotdataError, match='not an RFC 3339 timestamp'): + inst.get_data({'question': 'x'}) + assert calls['n'] == 1, 'a request-shaped 400 is identical on every attempt' + assert len(inst.asked) == 1, 'must not spend a second LLM call' + + +def test_a_400_with_a_query_run_id_is_retried(): + """The server mints a run id only once it has executed the statement.""" + err = client_mod.HotdataError('boom', status_code=400, query_run_id='qrun-1') + assert iinstance_mod._is_sql_fixable(err) is True + assert iinstance_mod._is_sql_fixable(client_mod.HotdataError('boom', status_code=400)) is False + + +def test_the_client_carries_the_query_run_id_off_an_error_body(): + """A statement failure comes back as 400 with the run id beside the error.""" + c, _rec = _client( + [_Resp(400, {'error': {'code': 'BAD_REQUEST', 'message': "table 'x' not found"}, 'query_run_id': 'qrun-9'})] + ) + with pytest.raises(client_mod.HotdataError) as excinfo: + c.query(sql='SELECT * FROM x', database_id='db-1') + assert excinfo.value.query_run_id == 'qrun-9' + assert iinstance_mod._is_sql_fixable(excinfo.value) is True + + +def test_the_client_leaves_the_run_id_empty_when_the_body_has_none(): + """A request-shaped 400 carries no run id, and is classified as not fixable.""" + c, _rec = _client( + [_Resp(400, {'error': {'code': 'BAD_REQUEST', 'message': 'async_after_ms must be at least 1000'}})] + ) + with pytest.raises(client_mod.HotdataError) as excinfo: + c.query(sql='SELECT 1', database_id='db-1') + assert excinfo.value.query_run_id == '' + assert iinstance_mod._is_sql_fixable(excinfo.value) is False + + +def test_get_result_keeps_its_positional_arguments(): + """``get_result(result_id, offset, limit)`` predates database scoping. The new + parameter is keyword-only, so an old positional call cannot send its offset + as the X-Database-Id header. + """ + c, rec = _client([_Resp(200, {'rows': []})]) + c.get_result('res-1', 5, 10) + assert rec.calls[0]['params'] == {'offset': 5, 'limit': 10} + assert 'X-Database-Id' not in rec.calls[0]['headers'] + with pytest.raises(TypeError): + c.get_result('res-1', 5, 10, 'db-1') + + +def test_a_cancelled_run_is_not_treated_as_a_statement_failure(): + """A rewritten query cannot un-cancel a run, so it must not earn a retry.""" + calls = {'n': 0} + + def _query(**_kw): + calls['n'] += 1 + return {'query_run_id': 'run-1'} + + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace( + information_schema=lambda **_kw: {'tables': []}, + query=_query, + get_query_run=lambda _i, database_id='': {'status': 'cancelled'}, + ) + inst = _llm_instance(g, ['SELECT 1', 'SELECT 2', 'SELECT 3']) + with pytest.raises(RuntimeError, match='cancelled') as excinfo: + inst.get_data({'question': 'x'}) + assert not isinstance(excinfo.value, iinstance_mod.SqlStatementError) + assert calls['n'] == 1 + assert len(inst.asked) == 1 + + +def test_a_pipe_in_a_column_name_does_not_split_the_header(): + """``SELECT total AS "a|b"`` is a legal alias. Unreachable while rows were + never objects; reachable now that the header is built from real names. + """ + table = iinstance_mod._rows_to_markdown([{'a|b': 1}]) + header, divider, body = table.splitlines() + assert header == '| a\\|b |' + assert divider == '| --- |' + assert body == '| 1 |' + + +def test_a_null_renders_as_an_empty_cell(): + """SQL NULL arrives as None; the Python word must not reach the reader.""" + table = iinstance_mod._rows_to_markdown([{'city': 'Reno', 'units': None}]) + assert table.splitlines()[2] == '| Reno | |' + assert 'None' not in table + + +def test_a_row_narrower_than_its_header_omits_the_missing_keys(): + """No nulls are invented for values the server never sent.""" + assert iinstance_mod._rows_as_objects([[1]], ['a', 'b']) == [{'a': 1}] + + +def test_the_config_schema_advertises_the_floor_the_node_enforces(): + """The form must not accept a value the node silently overrides, and neither + may drift from what the API accepts. + """ + services = json.loads((_NODE_DIR / 'services.json').read_text(encoding='utf-8')) + field = services['fields']['hotdata.async_after_ms'] + assert field['minimum'] == iglobal_mod.ASYNC_AFTER_MS_MIN == 1000 + assert field['default'] >= field['minimum'] + + +def test_a_self_join_keeps_every_value(): + """The shape the live API returns for ``SELECT * FROM sales a JOIN sales b``: + every name twice. This is how repeated names are actually reached - the + planner rejects a bare ``SELECT city, city``. + """ + rows = iinstance_mod._rows_as_objects([['Reno', 200, 'Reno', 200]], ['city', 'units', 'city', 'units']) + assert rows == [{'city': 'Reno', 'units': 200, 'city_1': 'Reno', 'units_1': 200}] + + +def test_a_failed_run_feeds_the_real_reason_back_to_the_model(): + """The live poll body reports the cause as ``error_message``. Reading only + ``error``/``message`` fell through to the status, so the corrective attempt + was handed "failed" and rewrote blind. + """ + calls = {'n': 0} + + def _query(**_kw): + calls['n'] += 1 + if calls['n'] == 1: + return {'query_run_id': 'run-1'} + return {'columns': ['a'], 'rows': [[1]], 'truncated': False} + + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace( + information_schema=lambda **_kw: {'tables': []}, + query=_query, + get_query_run=lambda _i, database_id='': { + 'status': 'failed', + 'error_message': 'query execution failed: Arrow error: Divide by zero error', + }, + ) + inst = _llm_instance(g, ['SELECT 1 / 0 FROM t', 'SELECT 1 FROM t']) + out = inst.get_data({'question': 'x'}) + assert out['attempts'] == 2 + assert 'Divide by zero' in inst.asked[1].all_text(), 'the model must see why the first statement failed' + + +def test_run_failure_text_reads_every_shape_the_api_uses(): + """``error_message`` first, then ``error`` as text or object, then ``message``, then the fallback.""" + text = iinstance_mod._run_failure_text + assert text({'error_message': 'boom'}, 'failed') == 'boom' + assert text({'error': 'boom'}, 'failed') == 'boom' + assert text({'error': {'code': 'BAD_REQUEST', 'message': 'boom'}}, 'failed') == 'boom' + assert text({'message': 'boom'}, 'failed') == 'boom' + assert text({'status': 'failed'}, 'failed') == 'failed' + + +# --------------------------------------------------------------------------- +# Review follow-ups, read against https://www.hotdata.dev/openapi.yaml +# --------------------------------------------------------------------------- + + +def _bounded_poll(body, most=3): + """A get_query_run stub that fails the test if it is polled past `most`. + + `time.sleep` is a no-op under test, so a status the poll loop does not + recognise would otherwise spin for the whole of job_timeout_secs. + """ + seen = {'n': 0} + + def _poll(_i, database_id=''): + """Return the same body, a bounded number of times.""" + seen['n'] += 1 + assert seen['n'] <= most, 'the poll loop did not treat this status as terminal' + return body + + return _poll + + +def test_a_truncated_result_is_read_only_after_its_run_succeeds(): + """A truncated body always carries a result_id, so the old `result_id is None` + test never waited for it - and that id is issued while the save is still in + flight. The read must follow the run. + """ + events = [] + polls = [{'status': 'running'}, {'status': 'succeeded', 'result_id': 'res-from-run'}] + + def _get_query_run(_i, database_id=''): + """Record the poll, then answer with the next queued run body.""" + events.append('poll') + return polls.pop(0) + + def _get_result(result_id, database_id='', offset=0, limit=None): + """Record which result id was read, and when.""" + events.append(f'read:{result_id}') + return {'status': 'ready', 'columns': ['n'], 'rows': [[1], [2]]} + + g = _loaded_global() + g.client = SimpleNamespace( + query=lambda **_kw: { + 'query_run_id': 'run-1', + 'result_id': 'res-from-body', + 'truncated': True, + 'columns': ['n'], + 'rows': [[1]], + }, + get_query_run=_get_query_run, + get_result=_get_result, + ) + out = _instance(g)._run_sql('SELECT n FROM big', 2) + assert events == ['poll', 'poll', 'read:res-from-run'], events + assert out['rows'] == [{'n': 1}, {'n': 2}] + + +def test_a_result_that_is_not_ready_raises_instead_of_reading_as_empty(): + """202 {status, result_id} carries no rows. It means "not ready", and must + never come back as a query that matched nothing. + """ + for state in ('processing', 'pending'): + c, _rec = _client([_Resp(202, {'status': state, 'result_id': 'res-1'}, headers={'Retry-After': '1'})]) + with pytest.raises(client_mod.HotdataError, match='not ready') as excinfo: + c.get_result('res-1', database_id='db-1') + assert excinfo.value.status_code == 202 + assert iinstance_mod._is_sql_fixable(excinfo.value) is False + + +def test_a_ready_result_with_no_rows_is_still_an_empty_result(): + """The guard keys on the not-ready status, so a genuinely empty result passes.""" + c, _rec = _client([_Resp(200, {'status': 'ready', 'columns': ['n'], 'rows': []})]) + assert c.get_result('res-1', database_id='db-1')['rows'] == [] + + +def test_an_interrupted_run_ends_the_poll_at_once(): + """`interrupted` is a documented terminal status. Unknown to the poll loop, it + read as "still running" and was polled until job_timeout_secs ran out, then + reported as a timeout that never happened. + """ + g = _loaded_global(job_timeout_secs=300) + g.client = SimpleNamespace( + query=lambda **_kw: {'query_run_id': 'run-1'}, + get_query_run=_bounded_poll({'status': 'interrupted'}, most=1), + ) + with pytest.raises(RuntimeError, match='interrupted') as excinfo: + _instance(g)._run_sql('SELECT a FROM t', 10) + assert not isinstance(excinfo.value, iinstance_mod.SqlStatementError) + assert 'did not finish within' not in str(excinfo.value) + + +def test_an_interrupted_run_is_not_handed_to_the_model_to_rewrite(): + """A rewritten statement cannot un-interrupt a run, so it must not cost an + LLM turn. + """ + g = _loaded_global(max_attempts=3) + g.client = SimpleNamespace( + information_schema=lambda **_kw: {'tables': []}, + query=lambda **_kw: {'query_run_id': 'run-1'}, + get_query_run=_bounded_poll({'status': 'interrupted'}), + ) + inst = _llm_instance(g, ['SELECT 1', 'SELECT 2', 'SELECT 3']) + with pytest.raises(RuntimeError, match='interrupted'): + inst.get_data({'question': 'x'}) + assert len(inst.asked) == 1 + + +def test_the_query_run_id_is_read_from_inside_the_error_object_too(): + """The documented error body is {"error": {code, message}}. The live API puts + the run id beside `error`; were it to move inside, every retry would stop. + """ + c, _rec = _client( + [_Resp(400, {'error': {'code': 'BAD_REQUEST', 'message': "table 'x' not found", 'query_run_id': 'qrun-7'}})] + ) + with pytest.raises(client_mod.HotdataError) as excinfo: + c.query(sql='SELECT * FROM x', database_id='db-1') + assert excinfo.value.query_run_id == 'qrun-7' + assert iinstance_mod._is_sql_fixable(excinfo.value) is True