diff --git a/CONTRACT.md b/CONTRACT.md new file mode 100644 index 0000000..ae60b53 --- /dev/null +++ b/CONTRACT.md @@ -0,0 +1,92 @@ +# hotdata-runtime Contract + +`hotdata-runtime` is the framework-agnostic runtime contract for Hotdata integrations. + +## Scope + +This package provides shared primitives for: + +- Environment and workspace resolution +- Query execution and polling +- Normalized tabular result handling +- Basic workspace health checks + +## Public Runtime Contract + +The supported import surface is: + +- `HotdataClient` +- `QueryResult` +- `from_env` +- `workspace_health_lines` +- `default_api_key` +- `default_host` +- `default_session_id` +- `explicit_workspace_id` +- `list_workspaces` +- `normalize_host` +- `pick_workspace` +- `resolve_workspace_selection` +- `ResultSummary` +- `RunHistoryItem` +- `WorkspaceSelection` + +Adapters should import from `hotdata_runtime` and treat this surface as the stable API. + +## Semantic Guarantees + +### `HotdataClient` + +- Represents runtime context: API key, host, workspace, optional session. +- `from_env()` resolves runtime context from env vars and selected workspace. +- `execute_sql(sql)` returns `QueryResult` or raises `RuntimeError`/`TimeoutError`. +- `get_result(result_id)` returns a ready `QueryResult` and waits for readiness when needed. +- `connections()` returns the connections API wrapper for adapter UI/status features. +- `query_runs()` returns the query-runs API wrapper for adapter history views. +- `results()` returns the results API wrapper for adapter result pickers. +- `list_recent_results(...)` returns normalized `ResultSummary` entries. +- `list_run_history(limit=...)` returns normalized `RunHistoryItem` entries. +- `list_qualified_table_names(...)` returns sorted fully qualified table names. +- `columns_for_qualified(qualified, connection_id=...)` resolves table columns, and + adapters should pass `connection_id` when known. + +### `QueryResult` + +- Canonical tabular result model with `columns`, `rows`, and `row_count`. +- Carries server identifiers and execution metadata when available. +- `to_pandas()` converts to a DataFrame with stable column ordering. +- `to_records(max_rows=...)` returns row dicts keyed by column names. +- `metadata_dict()` returns normalized result metadata for adapter rendering. + +### Env Resolution + +- `default_api_key()` reads `HOTDATA_API_KEY`. +- `default_host()` reads `HOTDATA_API_URL` (default: `https://api.hotdata.dev`) and normalizes it. +- `default_session_id()` reads `HOTDATA_SANDBOX`. +- `explicit_workspace_id()` reads `HOTDATA_WORKSPACE` (workspace public id). +- `pick_workspace()` prefers explicit env workspace, then active workspace, then first workspace. +- `resolve_workspace_selection()` is the canonical workspace selection algorithm. It returns `WorkspaceSelection` with selected workspace id, selection source, and discovered workspaces when auto-selected. + +## Adapter Responsibilities + +Framework packages (Jupyter, Marimo, LangChain, LangGraph, LlamaIndex, Streamlit) own: + +- Framework-native lifecycle and state management +- Rendering/UI concerns +- Tool/agent wrappers and callback integration + +They should not duplicate runtime env/workspace/query semantics. + +## Runtime Non-Goals + +`hotdata-runtime` does not define framework UI primitives and does not require framework dependencies. + +## Versioning Policy + +- Backward-incompatible contract changes require a major version bump. +- Additive contract changes are minor versions. +- Bug fixes that preserve contract semantics are patch versions. + +## Enforcement + +Contract stability is enforced by tests that verify the public export surface and key behavioral invariants. diff --git a/README.md b/README.md index e833ef8..f4e41ad 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,19 @@ Shared runtime primitives for Hotdata integrations: workspace/session semantics, execution context, query state, run history, and replayable result handles. Framework packages (Marimo, Jupyter, Streamlit, LangGraph) depend on this package. +Runtime boundary and guarantees are defined in `CONTRACT.md`. + +## Features + +- **Environment-driven client setup** — create clients from `HOTDATA_API_KEY`, optional `HOTDATA_API_URL`, `HOTDATA_WORKSPACE`, and `HOTDATA_SANDBOX`. +- **Workspace resolution** — choose an explicit workspace from env, otherwise discover workspaces and select the active workspace or first available workspace. +- **Sandbox/session propagation** — pass sandbox session context through the SDK via `X-Session-Id`. +- **HTTP resilience** — configure SDK retries for transient connection failures and retry SQL execution on stale pooled sockets. +- **SQL execution helper** — run SQL through `POST /v1/query`, poll async query runs when needed, and return a `QueryResult`. +- **Result utilities** — convert query results to records, pandas DataFrames, or metadata dictionaries for adapter display layers. +- **History helpers** — list recent results and query run history with normalized dataclasses. +- **Health helpers** — build compact API/workspace health summaries for UI integrations. + Install: ```bash @@ -9,6 +22,12 @@ uv pip install hotdata-runtime # or: pip install hotdata-runtime ``` +Example: + +```bash +python examples/basic_usage.py +``` + Development (uses **uv**; creates `.venv/` in this repo): ```bash diff --git a/examples/basic_usage.py b/examples/basic_usage.py new file mode 100644 index 0000000..4f93bd0 --- /dev/null +++ b/examples/basic_usage.py @@ -0,0 +1,25 @@ +"""Basic hotdata-runtime usage.""" + +from hotdata_runtime import from_env + + +def main() -> None: + client = from_env() + result = client.execute_sql("SELECT 1 AS ok") + + print("result metadata:", result.metadata_dict()) + print("records:", result.to_records(max_rows=5)) + + print("recent results:") + for item in client.list_recent_results(limit=5, offset=0): + print(item.to_dict()) + + print("run history:") + for item in client.list_run_history(limit=5): + print(item.to_dict()) + + client.close() + + +if __name__ == "__main__": + main() diff --git a/hotdata_runtime/__init__.py b/hotdata_runtime/__init__.py index 94d84a4..ef6bd9c 100644 --- a/hotdata_runtime/__init__.py +++ b/hotdata_runtime/__init__.py @@ -2,7 +2,12 @@ from importlib.metadata import PackageNotFoundError, version -from hotdata_runtime.client import HotdataClient, from_env +from hotdata_runtime.client import ( + HotdataClient, + ResultSummary, + RunHistoryItem, + from_env, +) from hotdata_runtime.env import ( default_api_key, default_host, @@ -11,6 +16,8 @@ list_workspaces, normalize_host, pick_workspace, + resolve_workspace_selection, + WorkspaceSelection, ) from hotdata_runtime.health import workspace_health_lines from hotdata_runtime.result import QueryResult @@ -33,4 +40,8 @@ "list_workspaces", "normalize_host", "pick_workspace", + "resolve_workspace_selection", + "ResultSummary", + "RunHistoryItem", + "WorkspaceSelection", ] diff --git a/hotdata_runtime/client.py b/hotdata_runtime/client.py index ea73bfc..c648f78 100644 --- a/hotdata_runtime/client.py +++ b/hotdata_runtime/client.py @@ -1,8 +1,12 @@ from __future__ import annotations +from dataclasses import asdict, dataclass import time from typing import Any, Iterator +from urllib3.exceptions import HTTPError as Urllib3HTTPError +from urllib3.exceptions import ProtocolError + from hotdata import ApiClient, Configuration from hotdata.api.connections_api import ConnectionsApi from hotdata.api.information_schema_api import InformationSchemaApi @@ -22,9 +26,33 @@ normalize_host, pick_workspace, ) +from hotdata_runtime.http import default_http_retries from hotdata_runtime.result import QueryResult _TERMINAL = frozenset({"succeeded", "failed", "cancelled"}) +_RESULT_FAILURE = frozenset({"failed", "cancelled"}) + + +@dataclass(frozen=True) +class ResultSummary: + result_id: str + status: str + created_at: str | None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class RunHistoryItem: + query_run_id: str + status: str + created_at: str | None + execution_time_ms: int | None + result_id: str | None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) class HotdataClient: @@ -47,6 +75,7 @@ def __init__( api_key=api_key, workspace_id=workspace_id, session_id=session_id, + retries=default_http_retries(), ) self._api = ApiClient(self._config) @@ -54,9 +83,7 @@ def __init__( def from_env(cls) -> HotdataClient: api_key = default_api_key() if not api_key: - raise RuntimeError( - "HOTDATA_API_KEY or HOTDATA_TOKEN must be set." - ) + raise RuntimeError("HOTDATA_API_KEY must be set.") host = default_host() session = default_session_id() workspace_id = pick_workspace(api_key, host, session) @@ -108,6 +135,39 @@ def query_runs(self) -> QueryRunsApi: def results(self) -> ResultsApi: return self._results_api() + def list_recent_results( + self, + *, + limit: int = 50, + offset: int = 0, + ) -> list[ResultSummary]: + listing = self.results().list_results(limit=limit, offset=offset) + return [ + ResultSummary( + result_id=r.id, + status=r.status, + created_at=r.created_at, + ) + for r in listing.results + ] + + def list_run_history( + self, + *, + limit: int = 20, + ) -> list[RunHistoryItem]: + listing = self.query_runs().list_query_runs(limit=limit) + return [ + RunHistoryItem( + query_run_id=r.id, + status=r.status, + created_at=r.created_at, + execution_time_ms=r.execution_time_ms, + result_id=r.result_id, + ) + for r in listing.query_runs + ] + def iter_tables( self, *, @@ -143,9 +203,26 @@ def list_qualified_table_names( def connection_id_by_name(self) -> dict[str, str]: listing = self.connections().list_connections() - return {c.name: c.id for c in listing.connections} + id_map: dict[str, str] = {} + duplicate_names: set[str] = set() + for c in listing.connections: + if c.name in id_map and id_map[c.name] != c.id: + duplicate_names.add(c.name) + id_map[c.name] = c.id + if duplicate_names: + names = ", ".join(sorted(duplicate_names)) + raise RuntimeError( + f"Duplicate connection names found: {names}. " + "Use an explicit connection_id." + ) + return id_map - def columns_for_qualified(self, qualified: str) -> list[TableInfo]: + def columns_for_qualified( + self, + qualified: str, + *, + connection_id: str | None = None, + ) -> list[TableInfo]: parts = qualified.split(".") if len(parts) < 3: raise ValueError( @@ -156,10 +233,12 @@ def columns_for_qualified(self, qualified: str) -> list[TableInfo]: parts[1], ".".join(parts[2:]), ) - id_map = self.connection_id_by_name() - conn_id = id_map.get(conn_name) - if not conn_id: - raise KeyError(f"Unknown connection {conn_name!r}") + conn_id = connection_id + if conn_id is None: + id_map = self.connection_id_by_name() + conn_id = id_map.get(conn_name) + if not conn_id: + raise KeyError(f"Unknown connection {conn_name!r}") resp = self._information_schema().information_schema( connection_id=conn_id, var_schema=schema_name, @@ -206,9 +285,9 @@ def _wait_result_ready( last = results.get_result(result_id) if last.status == "ready": return last - if last.status == "failed": + if last.status in _RESULT_FAILURE: raise RuntimeError( - last.error_message or "Result persistence failed" + last.error_message or f"Result {last.status}" ) time.sleep(interval_s) raise TimeoutError( @@ -217,6 +296,18 @@ def _wait_result_ready( ) def execute_sql(self, sql: str) -> QueryResult: + last_err: BaseException | None = None + for attempt in range(3): + try: + return self._execute_sql_once(sql) + except (ProtocolError, ConnectionResetError, Urllib3HTTPError) as e: + last_err = e + if attempt == 2: + raise + time.sleep(0.2 * (2**attempt)) + raise last_err # pragma: no cover + + def _execute_sql_once(self, sql: str) -> QueryResult: q = self._query_api() try: raw = q.query(QueryRequest(sql=sql)) diff --git a/hotdata_runtime/env.py b/hotdata_runtime/env.py index 4496a8b..8b0a84b 100644 --- a/hotdata_runtime/env.py +++ b/hotdata_runtime/env.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from dataclasses import dataclass from urllib.parse import urlparse from hotdata import ApiClient, Configuration @@ -18,15 +19,11 @@ def normalize_host(url: str) -> str: def default_api_key() -> str: - return os.environ.get("HOTDATA_API_KEY", "") or os.environ.get( - "HOTDATA_TOKEN", "" - ) + return os.environ.get("HOTDATA_API_KEY", "") def explicit_workspace_id() -> str | None: - return os.environ.get("HOTDATA_WORKSPACE") or os.environ.get( - "HOTDATA_WORKSPACE_ID" - ) + return os.environ.get("HOTDATA_WORKSPACE") def default_host() -> str: @@ -50,13 +47,35 @@ def list_workspaces(api_key: str, host: str, session_id: str | None): return listing.workspaces -def pick_workspace(api_key: str, host: str, session_id: str | None) -> str: +@dataclass(frozen=True) +class WorkspaceSelection: + workspace_id: str + source: str + workspaces: list + + +def resolve_workspace_selection( + api_key: str, host: str, session_id: str | None +) -> WorkspaceSelection: explicit = explicit_workspace_id() if explicit: - return explicit + return WorkspaceSelection( + workspace_id=explicit, + source="explicit_env", + workspaces=[], + ) workspaces = list_workspaces(api_key, host, session_id) if not workspaces: raise RuntimeError("No Hotdata workspaces found for this API key.") active = [w for w in workspaces if w.active] chosen = active[0] if active else workspaces[0] - return chosen.public_id + return WorkspaceSelection( + workspace_id=chosen.public_id, + source="active" if active else "first", + workspaces=workspaces, + ) + + +def pick_workspace(api_key: str, host: str, session_id: str | None) -> str: + selection = resolve_workspace_selection(api_key, host, session_id) + return selection.workspace_id diff --git a/hotdata_runtime/http.py b/hotdata_runtime/http.py new file mode 100644 index 0000000..1e7458b --- /dev/null +++ b/hotdata_runtime/http.py @@ -0,0 +1,19 @@ +"""HTTP client defaults for Hotdata SDK :class:`~hotdata.Configuration`.""" + +from __future__ import annotations + +from urllib3.util.retry import Retry + + +def default_http_retries() -> Retry: + """Retry transient connection failures (e.g. stale pooled sockets).""" + return Retry( + total=3, + connect=3, + read=3, + backoff_factor=0.2, + status_forcelist=(502, 503, 504), + allowed_methods=frozenset( + ["GET", "HEAD", "POST", "PUT", "DELETE", "OPTIONS", "PATCH"] + ), + ) diff --git a/hotdata_runtime/result.py b/hotdata_runtime/result.py index 64e700b..a8a18e8 100644 --- a/hotdata_runtime/result.py +++ b/hotdata_runtime/result.py @@ -20,6 +20,25 @@ class QueryResult: warning: str | None = None error_message: str | None = None + def to_records( + self, + *, + max_rows: int | None = None, + ) -> list[dict[str, Any]]: + rows = self.rows if max_rows is None else self.rows[:max_rows] + return [dict(zip(self.columns, row)) for row in rows] + + def metadata_dict(self) -> dict[str, Any]: + return { + "row_count": self.row_count, + "column_count": len(self.columns), + "result_id": self.result_id, + "query_run_id": self.query_run_id, + "execution_time_ms": self.execution_time_ms, + "warning": self.warning, + "error_message": self.error_message, + } + def to_pandas(self): # type: ignore[no-untyped-def] import pandas as pd diff --git a/tests/test_client.py b/tests/test_client.py index 1b420ec..fc5ccdb 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -6,7 +6,7 @@ import pytest -from hotdata_runtime.env import normalize_host, pick_workspace +from hotdata_runtime.env import normalize_host, pick_workspace, resolve_workspace_selection from hotdata_runtime.client import HotdataClient @@ -30,15 +30,23 @@ def test_pick_workspace_prefers_env(monkeypatch: pytest.MonkeyPatch): assert pick_workspace("k", "https://api.hotdata.dev", None) == "ws_explicit" -def test_pick_workspace_prefers_workspace_id_env(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv("HOTDATA_WORKSPACE", raising=False) - monkeypatch.setenv("HOTDATA_WORKSPACE_ID", "ws_from_id") - assert pick_workspace("k", "https://api.hotdata.dev", None) == "ws_from_id" +def test_resolve_workspace_selection_prefers_env_without_listing( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("HOTDATA_WORKSPACE", "ws_explicit") + with patch("hotdata_runtime.env.list_workspaces") as listing: + resolved = resolve_workspace_selection( + "k", "https://api.hotdata.dev", None + ) + listing.assert_not_called() + assert resolved.workspace_id == "ws_explicit" + assert resolved.source == "explicit_env" + assert resolved.workspaces == [] + def test_pick_workspace_chooses_first_active(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("HOTDATA_WORKSPACE", raising=False) - monkeypatch.delenv("HOTDATA_WORKSPACE_ID", raising=False) items = [ SimpleNamespace(public_id="ws_1", active=False), @@ -54,7 +62,6 @@ def test_pick_workspace_chooses_first_active(monkeypatch: pytest.MonkeyPatch): def test_pick_workspace_falls_back_to_first(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("HOTDATA_WORKSPACE", raising=False) - monkeypatch.delenv("HOTDATA_WORKSPACE_ID", raising=False) items = [ SimpleNamespace(public_id="ws_1", active=False), @@ -67,9 +74,155 @@ def test_pick_workspace_falls_back_to_first(monkeypatch: pytest.MonkeyPatch): assert pick_workspace("k", "https://api.hotdata.dev", None) == "ws_1" +def test_resolve_workspace_selection_source_first(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("HOTDATA_WORKSPACE", raising=False) + items = [ + SimpleNamespace(public_id="ws_1", active=False), + SimpleNamespace(public_id="ws_2", active=False), + ] + listing = SimpleNamespace(workspaces=items) + with patch("hotdata_runtime.env.WorkspacesApi") as Api: + Api.return_value.list_workspaces.return_value = listing + resolved = resolve_workspace_selection( + "k", "https://api.hotdata.dev", None + ) + assert resolved.workspace_id == "ws_1" + assert resolved.source == "first" + assert resolved.workspaces == items + + +def test_resolve_workspace_selection_returns_workspaces_and_source( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv("HOTDATA_WORKSPACE", raising=False) + + items = [ + SimpleNamespace(public_id="ws_1", active=False), + SimpleNamespace(public_id="ws_2", active=True), + ] + listing = SimpleNamespace(workspaces=items) + + with patch("hotdata_runtime.env.WorkspacesApi") as Api: + Api.return_value.list_workspaces.return_value = listing + resolved = resolve_workspace_selection( + "k", "https://api.hotdata.dev", None + ) + assert resolved.workspace_id == "ws_2" + assert resolved.source == "active" + assert resolved.workspaces == items + + def test_list_qualified_table_names_passes_connection_id(): client = HotdataClient("k", "ws", host="https://api.hotdata.dev") with patch.object(client, "iter_tables", return_value=iter([])) as it: client.list_qualified_table_names(limit=5, connection_id="conn_a") it.assert_called_once() assert it.call_args.kwargs["connection_id"] == "conn_a" + + +def test_wait_result_ready_raises_on_cancelled(): + 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) + + with patch.object(client, "_results_api", return_value=FakeResultsApi()): + with pytest.raises(RuntimeError, match="cancelled"): + client._wait_result_ready("res_1", timeout_s=0.1, interval_s=0) + + +def test_connection_id_by_name_raises_on_duplicate_names(): + client = HotdataClient("k", "ws", host="https://api.hotdata.dev") + listing = SimpleNamespace( + connections=[ + SimpleNamespace(name="warehouse", id="conn_1"), + SimpleNamespace(name="warehouse", id="conn_2"), + ] + ) + + class FakeConnectionsApi: + def list_connections(self): + return listing + + with patch.object(client, "connections", return_value=FakeConnectionsApi()): + with pytest.raises(RuntimeError, match="Duplicate connection names"): + client.connection_id_by_name() + + +def test_columns_for_qualified_prefers_explicit_connection_id(): + client = HotdataClient("k", "ws", host="https://api.hotdata.dev") + col = SimpleNamespace(name="a", data_type="INTEGER", nullable=True) + table = SimpleNamespace(columns=[col]) + response = SimpleNamespace(tables=[table]) + + class FakeInformationSchemaApi: + def __init__(self): + self.kwargs = None + + def information_schema(self, **kwargs): + self.kwargs = kwargs + return response + + fake_api = FakeInformationSchemaApi() + with patch.object(client, "_information_schema", return_value=fake_api), patch.object( + client, "connection_id_by_name" + ) as id_map: + cols = client.columns_for_qualified( + "warehouse.public.orders", + connection_id="conn_explicit", + ) + id_map.assert_not_called() + assert cols == [col] + assert fake_api.kwargs["connection_id"] == "conn_explicit" + + +def test_list_recent_results_returns_normalized_summaries(): + client = HotdataClient("k", "ws", host="https://api.hotdata.dev") + listing = SimpleNamespace( + results=[ + SimpleNamespace(id="res_1", status="ready", created_at="2026-01-01T00:00:00Z"), + SimpleNamespace(id="res_2", status="failed", created_at=None), + ] + ) + + class FakeResultsApi: + def list_results(self, *, limit: int, offset: int): + return listing + + with patch.object(client, "results", return_value=FakeResultsApi()): + out = client.list_recent_results(limit=10, offset=2) + assert [r.result_id for r in out] == ["res_1", "res_2"] + assert out[0].status == "ready" + assert out[0].to_dict()["created_at"] == "2026-01-01T00:00:00Z" + + +def test_list_run_history_returns_normalized_items(): + client = HotdataClient("k", "ws", host="https://api.hotdata.dev") + listing = SimpleNamespace( + query_runs=[ + SimpleNamespace( + id="run_1", + status="succeeded", + created_at="2026-01-01T00:00:00Z", + execution_time_ms=7, + result_id="res_1", + ), + ] + ) + + class FakeRunsApi: + def __init__(self): + self.kwargs = None + + def list_query_runs(self, *, limit: int): + self.kwargs = {"limit": limit} + return listing + + fake_runs = FakeRunsApi() + with patch.object(client, "query_runs", return_value=fake_runs): + out = client.list_run_history(limit=5) + assert [r.query_run_id for r in out] == ["run_1"] + assert out[0].execution_time_ms == 7 + assert out[0].to_dict()["result_id"] == "res_1" + assert fake_runs.kwargs == {"limit": 5} diff --git a/tests/test_contract.py b/tests/test_contract.py new file mode 100644 index 0000000..edfd55c --- /dev/null +++ b/tests/test_contract.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from dataclasses import fields +from unittest.mock import patch + +import hotdata_runtime as hr +from hotdata_runtime.client import HotdataClient +from hotdata_runtime.result import QueryResult + + +def test_public_exports_contract(): + assert hr.__all__ == [ + "__version__", + "HotdataClient", + "QueryResult", + "workspace_health_lines", + "default_api_key", + "default_host", + "default_session_id", + "explicit_workspace_id", + "from_env", + "list_workspaces", + "normalize_host", + "pick_workspace", + "resolve_workspace_selection", + "ResultSummary", + "RunHistoryItem", + "WorkspaceSelection", + ] + + +def test_module_from_env_delegates_to_client_classmethod(): + sentinel = HotdataClient("k", "ws", host="https://api.hotdata.dev") + with patch.object(HotdataClient, "from_env", return_value=sentinel) as m: + got = hr.from_env() + m.assert_called_once_with() + assert got is sentinel + + +def test_query_result_contract_fields(): + assert [f.name for f in fields(QueryResult)] == [ + "columns", + "rows", + "row_count", + "result_id", + "query_run_id", + "execution_time_ms", + "warning", + "error_message", + ] diff --git a/tests/test_result.py b/tests/test_result.py new file mode 100644 index 0000000..3c1ebac --- /dev/null +++ b/tests/test_result.py @@ -0,0 +1,37 @@ +from hotdata_runtime.result import QueryResult + + +def _result() -> QueryResult: + return QueryResult( + columns=["a", "b"], + rows=[[1, "x"], [2, "y"]], + row_count=2, + result_id="res_1", + query_run_id="run_1", + execution_time_ms=12, + warning="warn", + error_message=None, + ) + + +def test_to_records_returns_row_dicts(): + records = _result().to_records() + assert records == [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}] + + +def test_to_records_honors_max_rows(): + records = _result().to_records(max_rows=1) + assert records == [{"a": 1, "b": "x"}] + + +def test_metadata_dict_contains_normalized_fields(): + meta = _result().metadata_dict() + assert meta == { + "row_count": 2, + "column_count": 2, + "result_id": "res_1", + "query_run_id": "run_1", + "execution_time_ms": 12, + "warning": "warn", + "error_message": None, + }