fix: scope results and query runs to a database

Author: zfarrell

CHANGELOG.md

@@ -11,6 +11,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - feat: support async table loads and append mode
 - feat(uploads): support streaming uploads with on-demand part URLs
+- **Breaking:** results and query runs are now scoped to a database via the
+  required `X-Database-Id` header. The generated `ResultsApi` / `QueryRunsApi`
+  methods (`get_result`, `list_results`, `get_query_run`, `list_query_runs`)
+  gain a required `x_database_id` argument, and the hand-written ergonomic
+  helpers match: `hotdata.arrow.ResultsApi.get_result_arrow` /
+  `stream_result_arrow` and `hotdata.query.QueryApi.wait_for_result` take the
+  database id, and `QueryApi.query`'s truncation auto-follow forwards the
+  query's database scope (the `X-Database-Id` header, or the request-body
+  `database_id` when no header is set) to its result and query-run fetches.
 
 ## [0.5.0] - 2026-06-28
 

README.md

@@ -77,17 +77,21 @@ from hotdata.arrow import ResultsApi
 with ApiClient(Configuration(api_key="...", workspace_id="...")) as client:
     results = ResultsApi(client)
 
+    # Results are scoped to a database via the required `X-Database-Id` header,
+    # so pass the id of the database the query ran in.
+    database_id = "your_database_id"
+
     # Buffered: returns a pyarrow.Table.
-    table = results.get_result_arrow(result_id)
+    table = results.get_result_arrow(result_id, database_id)
 
     # Streaming: yields a pyarrow.RecordBatchStreamReader without
     # materializing the full table in memory.
-    with results.stream_result_arrow(result_id) as reader:
+    with results.stream_result_arrow(result_id, database_id) as reader:
         for batch in reader:
             ...
 ```
 
-Both methods accept `offset` and `limit` for pagination. They raise `hotdata.arrow.ResultNotReadyError` if the result is still pending or processing — poll `results.get_result(result_id)` until `status == "ready"` first.
+Both methods accept `offset` and `limit` for pagination. They raise `hotdata.arrow.ResultNotReadyError` if the result is still pending or processing — poll `results.get_result(result_id, database_id)` until `status == "ready"` first.
 
 ## File uploads
 

hotdata/arrow.py

@@ -87,6 +87,7 @@ class ResultsApi(_GeneratedResultsApi):
     def get_result_arrow(
         self,
         id: str,
+        x_database_id: str,
         *,
         offset: Optional[int] = None,
         limit: Optional[int] = None,
@@ -99,14 +100,18 @@ def get_result_arrow(
         iterate batches without materializing the whole table.
 
         :param id: Result ID.
+        :param x_database_id: Database the result belongs to. Results are scoped
+            to a database via the required ``X-Database-Id`` header — pass the
+            same database the query ran in.
         :param offset: Rows to skip (default: 0).
         :param limit: Maximum rows to return (default: unbounded).
         :raises ResultNotReadyError: result is still pending or processing.
         :raises hotdata.exceptions.ApiException: for other HTTP errors
             (400 invalid params, 404 not found, 409 failed result).
         """
         ipc = _import_pyarrow()
-        response = self._call_arrow(id=id, offset=offset, limit=limit,
+        response = self._call_arrow(id=id, x_database_id=x_database_id,
+                                    offset=offset, limit=limit,
                                     _request_timeout=_request_timeout)
         try:
             return ipc.open_stream(response).read_all()
@@ -117,6 +122,7 @@ def get_result_arrow(
     def stream_result_arrow(
         self,
         id: str,
+        x_database_id: str,
         *,
         offset: Optional[int] = None,
         limit: Optional[int] = None,
@@ -132,15 +138,22 @@ def stream_result_arrow(
 
         Example::
 
-            with results.stream_result_arrow(result_id) as reader:
+            with results.stream_result_arrow(result_id, database_id) as reader:
                 for batch in reader:
                     process(batch)
 
+        :param id: Result ID.
+        :param x_database_id: Database the result belongs to. Results are scoped
+            to a database via the required ``X-Database-Id`` header — pass the
+            same database the query ran in.
+        :param offset: Rows to skip (default: 0).
+        :param limit: Maximum rows to return (default: unbounded).
         :raises ResultNotReadyError: result is still pending or processing.
         :raises hotdata.exceptions.ApiException: for other HTTP errors.
         """
         ipc = _import_pyarrow()
-        response = self._call_arrow(id=id, offset=offset, limit=limit,
+        response = self._call_arrow(id=id, x_database_id=x_database_id,
+                                    offset=offset, limit=limit,
                                     _request_timeout=_request_timeout)
         try:
             yield ipc.open_stream(response)
@@ -151,16 +164,20 @@ def _call_arrow(
         self,
         *,
         id: str,
+        x_database_id: str,
         offset: Optional[int],
         limit: Optional[int],
         _request_timeout: Any,
     ) -> Any:
         # Build the request via the generator's private serialize helper so
         # path/query/auth handling stays in lockstep with the generated client.
         # Override only what we need: the Accept header and the format query.
+        # `GET /v1/results/{id}` is database-scoped, so the required
+        # X-Database-Id header flows through the generated serializer too.
         headers: Dict[str, Any] = {"Accept": ARROW_STREAM_MEDIA_TYPE}
         params = self._get_result_serialize(
             id=id,
+            x_database_id=x_database_id,
             offset=offset,
             limit=limit,
             format=ResultsFormatQuery.ARROW,

hotdata/query.py

@@ -337,11 +337,20 @@ def query(
             return response
         if not follow or not response.truncated:
             return response
-        return self._materialize_full(response, row_guard, byte_guard)
+        # Auto-follow re-fetches the persisted result and the query run, both of
+        # which are scoped to a database via the required X-Database-Id header.
+        # /v1/query takes exactly one database source, so the effective database
+        # is the header value when present, otherwise the request-body
+        # database_id.
+        effective_db = x_database_id or query_request.database_id or ""
+        return self._materialize_full(
+            response, effective_db, row_guard, byte_guard
+        )
 
     def wait_for_result(
         self,
         result_id: str,
+        x_database_id: str,
         *,
         results_api: Optional[_ResultsApi] = None,
     ) -> Any:
@@ -355,6 +364,10 @@ def wait_for_result(
         every poll, materializing the whole result into memory before the size
         guards can act.
 
+        :param result_id: Result ID to poll.
+        :param x_database_id: Database the result belongs to. Results are scoped
+            to a database via the required ``X-Database-Id`` header — pass the
+            same database the query ran in.
         :raises ResultFailedError: the result failed (delivered as HTTP 409).
         :raises ResultTimeoutError: ``ready`` not reached within the poll deadline.
         """
@@ -365,7 +378,7 @@ def wait_for_result(
         last_status = "pending"
         while True:
             try:
-                result = api.get_result(result_id, limit=0)
+                result = api.get_result(result_id, x_database_id=x_database_id, limit=0)
             except ApiException as exc:
                 # A failed result is delivered as HTTP 409: the generated client
                 # raises (response_deserialize raises on any non-2xx) rather than
@@ -449,16 +462,19 @@ def _backoff_delay(self, exc: ApiException, attempt: int) -> float:
     def _materialize_full(
         self,
         preview: QueryResponse,
+        x_database_id: str,
         max_auto_rows: Optional[int],
         max_auto_bytes: Optional[int],
     ) -> QueryResponse:
         if preview.result_id is None:
             raise ResultUnavailableError(warning=preview.warning)
 
         results_api = _ResultsApi(self.api_client)
-        self.wait_for_result(preview.result_id, results_api=results_api)
+        self.wait_for_result(
+            preview.result_id, x_database_id, results_api=results_api
+        )
 
-        total = self._authoritative_total(preview)
+        total = self._authoritative_total(preview, x_database_id)
         # Auto-follow does extra round-trips (poll + paginate) and materializes
         # the full result; log it so the hidden work behind one query() call is
         # observable without being noisy (info, not a warning).
@@ -477,7 +493,12 @@ def _materialize_full(
             )
 
         rows = self._fetch_all_rows(
-            preview.result_id, total, max_auto_rows, max_auto_bytes, results_api
+            preview.result_id,
+            x_database_id,
+            total,
+            max_auto_rows,
+            max_auto_bytes,
+            results_api,
         )
         # Replace the bounded preview with the full row set. truncated /
         # total_row_count stay as the server reported them so the caller can
@@ -487,22 +508,30 @@ def _materialize_full(
             preview.total_row_count = total
         return preview
 
-    def _authoritative_total(self, preview: QueryResponse) -> Optional[int]:
+    def _authoritative_total(
+        self, preview: QueryResponse, x_database_id: str
+    ) -> Optional[int]:
         """The grand total row count. ``total_row_count`` is null while a
         truncated result is still persisting, so fall back to the query-run
         record, which carries the authoritative count once the run succeeds.
+
+        The query run is scoped to a database via the required ``X-Database-Id``
+        header, so ``x_database_id`` is forwarded to the lookup.
         """
         if preview.total_row_count is not None:
             return preview.total_row_count
         try:
-            run = _QueryRunsApi(self.api_client).get_query_run(preview.query_run_id)
+            run = _QueryRunsApi(self.api_client).get_query_run(
+                preview.query_run_id, x_database_id=x_database_id
+            )
         except ApiException:
             return None
         return run.row_count
 
     def _fetch_all_rows(
         self,
         result_id: str,
+        x_database_id: str,
         total: Optional[int],
         max_auto_rows: Optional[int],
         max_auto_bytes: Optional[int],
@@ -515,6 +544,7 @@ def _fetch_all_rows(
         while True:
             page = results_api.get_result(
                 result_id,
+                x_database_id=x_database_id,
                 offset=offset,
                 limit=page_size,
                 format=ResultsFormatQuery.JSON,

tests/integration/test_query_async_polling.py

@@ -56,7 +56,7 @@ def test_query_async_polling(
     deadline = time.monotonic() + POLL_TIMEOUT_S
     run = None
     while time.monotonic() < deadline:
-        run = query_runs_api.get_query_run(query_run_id)
+        run = query_runs_api.get_query_run(query_run_id, x_database_id=database_id)
         if run.status in TERMINAL_STATUSES:
             break
         time.sleep(POLL_INTERVAL_S)
@@ -70,13 +70,13 @@ def test_query_async_polling(
     )
     assert run.row_count == 1
 
-    runs_listing = query_runs_api.list_query_runs(limit=50)
+    runs_listing = query_runs_api.list_query_runs(x_database_id=database_id, limit=50)
     assert any(r.id == query_run_id for r in runs_listing.query_runs), (
         f"query run {query_run_id} not surfaced by list_query_runs"
     )
 
     if run.result_id:
-        result = results_api.get_result(run.result_id)
+        result = results_api.get_result(run.result_id, x_database_id=database_id)
         assert result.result_id == run.result_id
         assert result.status in {"ready", "processing"}
         if result.status == "ready":
@@ -85,7 +85,7 @@ def test_query_async_polling(
 
         # ResultInfo (list_results) exposes the id as `id`, not `result_id` —
         # only GetResultResponse uses `result_id`.
-        results_listing = results_api.list_results(limit=50)
+        results_listing = results_api.list_results(x_database_id=database_id, limit=50)
         assert any(r.id == run.result_id for r in results_listing.results), (
             f"result {run.result_id} not surfaced by list_results"
         )

tests/integration/test_results_arrow.py

@@ -70,7 +70,7 @@ def test_results_arrow(
     deadline = time.monotonic() + POLL_TIMEOUT_S
     run = None
     while time.monotonic() < deadline:
-        run = query_runs_api.get_query_run(query_run_id)
+        run = query_runs_api.get_query_run(query_run_id, x_database_id=database_id)
         if run.status in TERMINAL_STATUSES:
             break
         time.sleep(POLL_INTERVAL_S)
@@ -85,27 +85,27 @@ def test_results_arrow(
     # ResultNotReadyError on 202.
     deadline = time.monotonic() + POLL_TIMEOUT_S
     while time.monotonic() < deadline:
-        result = results_api.get_result(result_id)
+        result = results_api.get_result(result_id, x_database_id=database_id)
         if result.status == "ready":
             break
         time.sleep(POLL_INTERVAL_S)
     else:
         pytest.fail(f"result {result_id} never became ready")
 
     # Buffered: returns a full pyarrow.Table.
-    table = results_api.get_result_arrow(result_id)
+    table = results_api.get_result_arrow(result_id, database_id)
     assert isinstance(table, pa.Table)
     assert table.num_rows == 2
     assert set(table.column_names) == {"x", "msg"}
     assert table.column("x").to_pylist() == [1, 2]
     assert table.column("msg").to_pylist() == ["hello", "world"]
 
     # Streaming: same data via RecordBatchStreamReader.
-    with results_api.stream_result_arrow(result_id) as reader:
+    with results_api.stream_result_arrow(result_id, database_id) as reader:
         streamed = pa.Table.from_batches(list(reader), schema=reader.schema)
     assert streamed.equals(table)
 
     # Pagination forwards correctly.
-    sliced = results_api.get_result_arrow(result_id, offset=1, limit=1)
+    sliced = results_api.get_result_arrow(result_id, database_id, offset=1, limit=1)
     assert sliced.num_rows == 1
     assert sliced.column("x").to_pylist() == [2]

tests/test_arrow.py

@@ -148,7 +148,7 @@ def test_get_result_arrow_returns_table(monkeypatch: pytest.MonkeyPatch) -> None
     _install_fake_response(monkeypatch, fake, captured)
 
     results, _ = _make_results_api()
-    got = results.get_result_arrow("res_123")
+    got = results.get_result_arrow("res_123", "db_x")
 
     assert got.equals(table)
     # Connection is drained (not just released) so a partially-read body can't
@@ -166,6 +166,8 @@ def test_get_result_arrow_returns_table(monkeypatch: pytest.MonkeyPatch) -> None
     assert "format=arrow" in call["url"]
     # Accept header overrides the JSON default.
     assert call["headers"]["Accept"] == ARROW_STREAM_MEDIA_TYPE
+    # Results are database-scoped: the required X-Database-Id header is sent.
+    assert call["headers"]["X-Database-Id"] == "db_x"
 
 
 def test_get_result_arrow_forwards_offset_and_limit(
@@ -180,7 +182,7 @@ def test_get_result_arrow_forwards_offset_and_limit(
     _install_fake_response(monkeypatch, fake, captured)
 
     results, _ = _make_results_api()
-    results.get_result_arrow("res_123", offset=10, limit=100)
+    results.get_result_arrow("res_123", "db_x", offset=10, limit=100)
 
     url = captured[0]["url"]
     assert "offset=10" in url
@@ -199,7 +201,7 @@ def test_stream_result_arrow_yields_reader(
     _install_fake_response(monkeypatch, fake, [])
 
     results, _ = _make_results_api()
-    with results.stream_result_arrow("res_123") as reader:
+    with results.stream_result_arrow("res_123", "db_x") as reader:
         batches = list(reader)
         assert batches, "expected at least one RecordBatch"
         roundtrip = pa.Table.from_batches(batches, schema=reader.schema)
@@ -232,7 +234,7 @@ def test_get_result_arrow_raises_when_not_ready(
 
     results, _ = _make_results_api()
     with pytest.raises(ResultNotReadyError) as ei:
-        results.get_result_arrow("res_pending")
+        results.get_result_arrow("res_pending", "db_x")
 
     assert ei.value.status == "processing"
     assert ei.value.result_id == "res_pending"
@@ -252,7 +254,7 @@ def test_get_result_arrow_raises_api_exception_on_404(
 
     results, _ = _make_results_api()
     with pytest.raises(ApiException) as ei:
-        results.get_result_arrow("res_missing")
+        results.get_result_arrow("res_missing", "db_x")
 
     assert ei.value.status == 404
     assert fake.release_conn_called
@@ -278,7 +280,7 @@ def test_get_result_arrow_raises_api_exception_on_409_failed(
 
     results, _ = _make_results_api()
     with pytest.raises(ApiException) as ei:
-        results.get_result_arrow("res_failed")
+        results.get_result_arrow("res_failed", "db_x")
 
     assert ei.value.status == 409