arrowbricks 0.1.2__tar.gz → 0.2.0__tar.gz

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/AGENTS.md

@@ -6,9 +6,9 @@ Runs SQL against a Databricks SQL warehouse via the Statement Execution API and
 
 - `src/arrowbricks/client.py` -- pure REST client (auth, statement submission/polling, backpressure-bounded concurrent chunk download). No Arrow dependency at all, intentionally: someone who only wants `execute_json_statement` or `upload_volume_file`/`delete_volume_file` shouldn't need arro3 pulled in conceptually either, even though it's a hard dependency of the package as a whole. Retries are a small hand-rolled `_retry_call` loop, not a dependency (see "Design invariants" below).
 - `src/arrowbricks/_streaming.py` -- Arrow (de)serialization via arro3, always: `ReplayableArrowChunk`, `write_ipc_stream` (always uncompressed, see below), heartbeat helpers, chunk fetching, and `stream_query_json` (arro3's `write_ndjson`). One Arrow engine, no pluggable backend -- unlike duckbricks, there's nothing here to make pluggable; arro3 *is* the whole point.
-- `src/arrowbricks/cursor.py` -- `Connection`/`Cursor`, the DB-API-ish surface (`execute`/`execute_streamed`, `fetchone`/`fetchmany`/`fetchall`, `fetchall_arrow`/`fetchmany_arrow`). `_ResultSet` buffers at the Arrow-Table level (not materialized Python rows) so the Arrow-native fetch methods stay zero-copy; row-based fetches materialize lazily off that buffer.
+- `src/arrowbricks/cursor.py` -- `Connection`/`Cursor`, the DB-API-ish surface (`execute`/`execute_streamed`, `fetchone`/`fetchmany`/`fetchall`, `fetchall_arrow`/`fetchmany_arrow`, `fetchall_streamed`/`fetchall_arrow_streamed`). `_ResultSet` buffers at the Arrow-Table level (not materialized Python rows) so the Arrow-native fetch methods stay zero-copy; row-based fetches materialize lazily off that buffer.
 - `tests/` -- respx mocks the Databricks REST endpoints (warehouse status, statement submit, chunk-link resolution, external-link byte download); no real warehouse or credentials needed to run the suite.
-- `examples/` -- `basic.py` (static token), `cursor_paging.py` (fetchmany/fetchmany_arrow over a large result), `fastapi_sse.py` (streaming NDJSON as SSE), `azure_auth.py` (Azure AD `token_provider` via `azure-identity`, kept out of core deps on purpose -- don't widen `ty check`'s scope to include it).
+- `examples/` -- `basic.py` (static token), `cursor_paging.py` (fetchmany/fetchmany_arrow over a large result), `fastapi_sse.py` (streaming NDJSON as SSE), `fastapi_sse_pivot.py` (buffered fetchall_streamed with one combined heartbeat/timeout budget across the wait+download phases), `azure_auth.py` (Azure AD `token_provider` via `azure-identity`, kept out of core deps on purpose -- don't widen `ty check`'s scope to include it).
 
 ## Commands
 
@@ -25,8 +25,9 @@ One-time setup per clone: `prek install` (needs `uv tool install prek` first if
 
 - **No cloud-SDK dependency.** Auth is `token: str` or `token_provider: Callable[[], str | Awaitable[str]]`. Do not add `azure-identity`/`boto3`/etc. as a real dependency -- that belongs in the caller's app.
 - **No hardcoded catalog/schema.** `catalog`/`schema` default to `None` everywhere. This package has zero knowledge of any specific Databricks workspace's naming.
-- **Chunk order is not fetch order.** `DatabricksClient` fetches chunks concurrently (bounded, with backpressure) and they can complete out of order. `_ResultSet`/`stream_query_json` both hold a `pending: dict[int, chunk]` reorder buffer keyed by `chunk_index`, releasing in order as the next expected index shows up. If you touch either, keep a test proving order survives out-of-order arrival (see `test_fetchall_preserves_order_despite_out_of_order_chunks`, `test_stream_query_json_preserves_order_despite_out_of_order_chunks`).
+- **Chunk order is not fetch order, and a chunk_index is not guaranteed unique or contiguous.** `DatabricksClient` fetches chunks concurrently (bounded, with backpressure) and they can complete out of order. `_ResultSet`/`stream_query_json` both hold a `pending: dict[int, list[chunk]]` reorder buffer keyed by `chunk_index` -- a **list**, not a single chunk, because `_fetch_chunk_index` can yield more than one blob for the same index (multiple `external_links` per chunk, "usually exactly one" but not guaranteed). Releases in order as the next expected index empties out; once the source is exhausted, a genuine gap (an index that never showed up at all) drains the lowest remaining index instead of stranding everything buffered past it. If you touch either, keep a test proving order survives out-of-order arrival AND that duplicate/missing indices never lose rows (see `test_fetchall_preserves_order_despite_out_of_order_chunks`, `test_stream_query_json_preserves_order_despite_out_of_order_chunks`, `tests/test_result_set.py`).
 - **Chunks are fetched lazily, not all upfront.** `_ResultSet` only pulls the next chunk from `_chunk_aiter` when the caller's `fetchone`/`fetchmany`/`fetchall` actually needs more rows than are already buffered. Don't "simplify" this into draining the whole chunk iterator inside `execute()` -- that defeats the point of a paginated cursor.
+- **`execute_streamed`'s heartbeat/timeout only covers the wait for the statement to become ready -- not downloading any chunk.** That was a real bug in this package's first release: a caller doing `execute_streamed()` then `fetchall()` got zero timeout enforcement and zero heartbeats during a slow multi-chunk download, exactly the case heartbeats exist for. `fetchall_streamed`/`fetchall_arrow_streamed` cover that second phase -- a caller wanting one combined budget across both phases must track its own deadline and pass the *remaining* time into the second call (see `examples/fastapi_sse_pivot.py`), since two independently-clocked `total_timeout_s`s would let a pathological case run up to 2x the intended ceiling.
 - **No silent row caps.** There's no `ABSOLUTE_ROW_LIMIT`-style ceiling baked in. If a caller wants one, that's `row_limit`, which they pass explicitly.
 - **No retry dependency.** `client.py`'s `_retry_call` is a ~10-line hand-rolled exponential-backoff loop, replacing tenacity on purpose -- it's the only retry pattern in the whole client, so a dependency for it wasn't worth it.
 - **One Arrow engine, no pluggable backend.** Unlike duckbricks (which supports nanoarrow *or* arro3 *or* bring-your-own), arrowbricks is arro3-only by design -- that's the entire "single responsibility" pitch. Don't add a backend-abstraction layer back in; if a caller needs a different Arrow engine, that's duckbricks' `set_arrow_backend()`, not this package.

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 0.1.2
+Version: 0.2.0
 Summary: Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a DB-API-ish Cursor (fetchone/fetchmany/fetchall/fetchall_arrow) or NDJSON streaming. Single Arrow engine (arro3), no DuckDB.
 Project-URL: Repository, https://github.com/bmsuisse/arrowbricks
 License-Expression: MIT
@@ -72,8 +72,11 @@ See [`examples/basic.py`](examples/basic.py) for a runnable version,
 [`examples/cursor_paging.py`](examples/cursor_paging.py) for paging a large
 result with `fetchmany`/`fetchmany_arrow` without buffering it all upfront,
 [`examples/fastapi_sse.py`](examples/fastapi_sse.py) for streaming a query to
-a client as Server-Sent Events, or [`examples/azure_auth.py`](examples/azure_auth.py)
-for a caching `token_provider` built on Azure AD (`DefaultAzureCredential`).
+a client as Server-Sent Events, [`examples/fastapi_sse_pivot.py`](examples/fastapi_sse_pivot.py)
+for the same over a buffered `Cursor.fetchall_streamed` result with one
+combined heartbeat/timeout budget across both the wait and the download, or
+[`examples/azure_auth.py`](examples/azure_auth.py) for a caching
+`token_provider` built on Azure AD (`DefaultAzureCredential`).
 
 ## Why not `databricks-sql-connector`?
 
@@ -102,9 +105,10 @@ conn = connect(host=..., warehouse_id=..., token_provider=my_token_provider)
 - `Connection.cursor() -> Cursor`
 - `Connection.client -> DatabricksClient` -- the same client `cursor()` uses, for lower-level access (e.g. `stream_query_json`, `execute_json_statement`, `upload_volume_file`).
 - `Cursor.execute(sql, parameters=None, *, row_limit=None, offset=None, catalog=None, schema=None, total_timeout_s=None) -> Cursor` -- submits and waits for the statement, like a real DB-API cursor. `parameters`, if given, is Databricks' own named-parameter format -- `[{"name": ..., "value": ..., "type": ...}]` bound against `:name` markers in `sql`.
-- `Cursor.execute_streamed(...)` -- same args, but an async generator yielding `HEARTBEAT` while waiting on a slow cold start, then the ready `Cursor` -- for bridging e.g. an SSE connection.
+- `Cursor.execute_streamed(...)` -- same args, but an async generator yielding `HEARTBEAT` while waiting on a slow cold start, then the ready `Cursor` -- for bridging e.g. an SSE connection. Its timeout/heartbeats stop the moment the statement is ready, *before* any chunk has been downloaded -- see `fetchall_streamed` below for the download phase itself.
 - `Cursor.fetchone() -> tuple | None`, `Cursor.fetchmany(size) -> list[tuple]`, `Cursor.fetchall() -> list[tuple]`
 - `Cursor.fetchmany_arrow(size) -> arro3.core.Table`, `Cursor.fetchall_arrow() -> arro3.core.Table`
+- `Cursor.fetchall_streamed(*, total_timeout_s=None)` / `Cursor.fetchall_arrow_streamed(*, total_timeout_s=None)` -- like `fetchall()`/`fetchall_arrow()`, but yield `HEARTBEAT` while pulling chunks instead of blocking silently, then the final rows/Table -- for a caller downloading a large result over SSE who needs heartbeats (and a timeout) through the *download*, not just the initial wait. Compose with `execute_streamed` and a shared deadline if you want one combined budget across both phases (see `examples/fastapi_sse_pivot.py`).
 - `Cursor` is an async iterator, yielding one row (tuple) at a time.
 - `Cursor.description` -- DB-API-style `[(name, type_name, None, None, None, None, None), ...]` after `execute()`.
 - `stream_query_json(client, sql, **kwargs)` -- yields `HEARTBEAT`, then each row as a JSON string, as soon as its chunk arrives. Timestamps come out as full ISO-8601, every column key is always present (`"col":null` for a null value, never an omitted key).

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/README.md

@@ -59,8 +59,11 @@ See [`examples/basic.py`](examples/basic.py) for a runnable version,
 [`examples/cursor_paging.py`](examples/cursor_paging.py) for paging a large
 result with `fetchmany`/`fetchmany_arrow` without buffering it all upfront,
 [`examples/fastapi_sse.py`](examples/fastapi_sse.py) for streaming a query to
-a client as Server-Sent Events, or [`examples/azure_auth.py`](examples/azure_auth.py)
-for a caching `token_provider` built on Azure AD (`DefaultAzureCredential`).
+a client as Server-Sent Events, [`examples/fastapi_sse_pivot.py`](examples/fastapi_sse_pivot.py)
+for the same over a buffered `Cursor.fetchall_streamed` result with one
+combined heartbeat/timeout budget across both the wait and the download, or
+[`examples/azure_auth.py`](examples/azure_auth.py) for a caching
+`token_provider` built on Azure AD (`DefaultAzureCredential`).
 
 ## Why not `databricks-sql-connector`?
 
@@ -89,9 +92,10 @@ conn = connect(host=..., warehouse_id=..., token_provider=my_token_provider)
 - `Connection.cursor() -> Cursor`
 - `Connection.client -> DatabricksClient` -- the same client `cursor()` uses, for lower-level access (e.g. `stream_query_json`, `execute_json_statement`, `upload_volume_file`).
 - `Cursor.execute(sql, parameters=None, *, row_limit=None, offset=None, catalog=None, schema=None, total_timeout_s=None) -> Cursor` -- submits and waits for the statement, like a real DB-API cursor. `parameters`, if given, is Databricks' own named-parameter format -- `[{"name": ..., "value": ..., "type": ...}]` bound against `:name` markers in `sql`.
-- `Cursor.execute_streamed(...)` -- same args, but an async generator yielding `HEARTBEAT` while waiting on a slow cold start, then the ready `Cursor` -- for bridging e.g. an SSE connection.
+- `Cursor.execute_streamed(...)` -- same args, but an async generator yielding `HEARTBEAT` while waiting on a slow cold start, then the ready `Cursor` -- for bridging e.g. an SSE connection. Its timeout/heartbeats stop the moment the statement is ready, *before* any chunk has been downloaded -- see `fetchall_streamed` below for the download phase itself.
 - `Cursor.fetchone() -> tuple | None`, `Cursor.fetchmany(size) -> list[tuple]`, `Cursor.fetchall() -> list[tuple]`
 - `Cursor.fetchmany_arrow(size) -> arro3.core.Table`, `Cursor.fetchall_arrow() -> arro3.core.Table`
+- `Cursor.fetchall_streamed(*, total_timeout_s=None)` / `Cursor.fetchall_arrow_streamed(*, total_timeout_s=None)` -- like `fetchall()`/`fetchall_arrow()`, but yield `HEARTBEAT` while pulling chunks instead of blocking silently, then the final rows/Table -- for a caller downloading a large result over SSE who needs heartbeats (and a timeout) through the *download*, not just the initial wait. Compose with `execute_streamed` and a shared deadline if you want one combined budget across both phases (see `examples/fastapi_sse_pivot.py`).
 - `Cursor` is an async iterator, yielding one row (tuple) at a time.
 - `Cursor.description` -- DB-API-style `[(name, type_name, None, None, None, None, None), ...]` after `execute()`.
 - `stream_query_json(client, sql, **kwargs)` -- yields `HEARTBEAT`, then each row as a JSON string, as soon as its chunk arrives. Timestamps come out as full ISO-8601, every column key is always present (`"col":null` for a null value, never an omitted key).

arrowbricks-0.2.0/examples/fastapi_sse_pivot.py

@@ -0,0 +1,78 @@
+"""Stream a large result to a client over SSE with ONE combined heartbeat/
+timeout budget spanning both phases: waiting for the statement to complete
+(execute_streamed) AND downloading every chunk afterwards (fetchall_streamed)
+-- the two-phase composition `fastapi_sse.py`'s simpler stream_query_json
+example doesn't need, since that one streams rows as chunks arrive rather
+than buffering a full result first.
+
+Splitting the timeout into two separately-clocked halves (each given the
+same total_timeout_s) would let a pathological case run up to 2x the
+intended ceiling -- tracking one deadline and passing the *remaining* budget
+into the second phase keeps it as one real ceiling across both.
+
+Requires `pip install arrowbricks fastapi uvicorn`.
+
+    DATABRICKS_HOST=adb-1234567890.1.azuredatabricks.net \\
+    DATABRICKS_WAREHOUSE_ID=abcd1234efgh5678 \\
+    DATABRICKS_TOKEN=dapi... \\
+    uvicorn examples.fastapi_sse_pivot:app --reload
+
+Then, in another terminal:
+
+    curl -N "http://localhost:8000/pivot?sql=SELECT+*+FROM+range(1000000)"
+
+Same caveat as fastapi_sse.py: `sql` comes straight from the request for
+brevity -- validate/allowlist it in a real deployment.
+"""
+
+import asyncio
+import os
+from collections.abc import AsyncIterator
+
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+
+from arrowbricks import HEARTBEAT, connect
+from arrowbricks.cursor import Cursor
+
+app = FastAPI()
+
+conn = connect(
+    host=os.environ["DATABRICKS_HOST"],
+    warehouse_id=os.environ["DATABRICKS_WAREHOUSE_ID"],
+    token=os.environ["DATABRICKS_TOKEN"],
+)
+
+_TOTAL_TIMEOUT_S = 300.0
+
+
+async def _fetch_all_rows_with_heartbeat(cursor: Cursor, deadline: float) -> AsyncIterator[object]:
+    """Yields HEARTBEAT while downloading, then the final list[Row] -- bounded
+    by whatever's left of the shared deadline, not a fresh full timeout."""
+    loop = asyncio.get_running_loop()
+    remaining = max(deadline - loop.time(), 0.0)
+    async for item in cursor.fetchall_streamed(total_timeout_s=remaining):
+        yield item
+
+
+async def _sse(sql: str) -> AsyncIterator[str]:
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + _TOTAL_TIMEOUT_S
+    cursor = conn.cursor()
+
+    async for item in cursor.execute_streamed(sql, total_timeout_s=_TOTAL_TIMEOUT_S):
+        if item is HEARTBEAT:
+            yield ": keep-alive\n\n"
+            continue
+        async for fetch_item in _fetch_all_rows_with_heartbeat(cursor, deadline):
+            if fetch_item is HEARTBEAT:
+                yield ": keep-alive\n\n"
+            else:
+                for row in fetch_item:
+                    yield f"data: {row}\n\n"
+    yield "event: end\ndata: {}\n\n"
+
+
+@app.get("/pivot")
+async def pivot(sql: str) -> StreamingResponse:
+    return StreamingResponse(_sse(sql), media_type="text/event-stream")

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "0.1.2"
+version = "0.2.0"
 description = "Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a DB-API-ish Cursor (fetchone/fetchmany/fetchall/fetchall_arrow) or NDJSON streaming. Single Arrow engine (arro3), no DuckDB."
 readme = "README.md"
 license = "MIT"

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/src/arrowbricks/_streaming.py

@@ -236,7 +236,11 @@ async def stream_query_json(
     O(whole result). Chunks can arrive out of order, so out-of-order arrivals
     sit in a small `pending` buffer until the next expected chunk_index shows
     up -- that buffer stays bounded by concurrency, it never grows to the full
-    result.
+    result. `pending` holds a *list* per index, not a single chunk, since a
+    chunk_index can have more than one blob (DatabricksClient._fetch_chunk_index
+    gathers over possibly-multiple `external_links` per chunk); once the
+    source is exhausted, any index that never showed up at all is skipped
+    rather than stranding everything buffered past it forever.
 
     Note this yields a whole chunk's rows at once (write_ndjson has no
     incremental/row-at-a-time mode) -- Databricks' own chunk sizing already
@@ -247,17 +251,31 @@ async def stream_query_json(
         client, sql, catalog=catalog, schema=schema, parameters=params
     )
 
-    pending: dict[int, ReplayableArrowChunk] = {}
+    pending: dict[int, list[ReplayableArrowChunk]] = {}
     next_idx = 0
     loop = asyncio.get_running_loop()
+
+    async def _emit_lines(chunk: ReplayableArrowChunk) -> AsyncIterator[str]:
+        blob = await loop.run_in_executor(None, _write_ndjson, chunk)
+        for line in blob.splitlines():
+            yield line.decode()
+
     async for item in heartbeat_over_stream(chunk_iter, total_timeout_s=total_timeout_s):
         if item is HEARTBEAT:
             yield HEARTBEAT
             continue
-        pending[item.chunk_index] = item
+        pending.setdefault(item.chunk_index, []).append(item)
         while next_idx in pending:
-            chunk = pending.pop(next_idx)
-            blob = await loop.run_in_executor(None, _write_ndjson, chunk)
-            for line in blob.splitlines():
-                yield line.decode()
-            next_idx += 1
+            chunk = pending[next_idx].pop(0)
+            if not pending[next_idx]:
+                del pending[next_idx]
+                next_idx += 1
+            async for line in _emit_lines(chunk):
+                yield line
+
+    # A genuine gap (an index that never arrived) must not strand chunks
+    # buffered past it -- drain whatever's left, in ascending index order.
+    for idx in sorted(pending):
+        for chunk in pending[idx]:
+            async for line in _emit_lines(chunk):
+                yield line

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/src/arrowbricks/cursor.py

@@ -54,32 +54,57 @@ class _ResultSet:
     def __init__(self, schema: core.Schema | None, chunk_aiter: AsyncIterator[ReplayableArrowChunk]) -> None:
         self.schema = schema
         self._chunk_aiter = chunk_aiter
-        self._pending: dict[int, ReplayableArrowChunk] = {}
+        # A list per index, not a single chunk -- a chunk_index can have more
+        # than one blob (DatabricksClient._fetch_chunk_index gathers over
+        # possibly-multiple `external_links` per chunk, "usually exactly one"
+        # but not guaranteed to be), and a dict keyed by index alone would
+        # silently overwrite/lose the first blob when a second one for the
+        # same index arrives.
+        self._pending: dict[int, list[ReplayableArrowChunk]] = {}
         self._next_idx = 0
         self._exhausted = False
         self._buffer: core.Table | None = None
         self.rownumber = 0
 
+    def _pop_pending(self, idx: int) -> ReplayableArrowChunk:
+        blobs = self._pending[idx]
+        chunk = blobs.pop(0)
+        if not blobs:
+            del self._pending[idx]
+        return chunk
+
     async def _pull_one_chunk_table(self) -> core.Table | None:
-        """Returns the next expected chunk's Table, in order -- checking
-        `_pending` FIRST, since a single earlier call can have pulled several
-        chunks off `_chunk_aiter` before the one it actually needed showed up
-        (arrival order is completion order, not chunk_index order), leaving
-        the rest already-fetched-and-buffered here. Only touches the network
-        (`_chunk_aiter.__anext__()`) once `_pending` has nothing more to give."""
+        """Returns one chunk's Table, preferring the next expected index --
+        checking `_pending` FIRST, since a single earlier call can have pulled
+        several chunks off `_chunk_aiter` before the one it actually needed
+        showed up (arrival order is completion order, not chunk_index order),
+        leaving the rest already-fetched-and-buffered here. Only touches the
+        network (`_chunk_aiter.__anext__()`) once `_pending` has nothing more
+        to give for the current index.
+
+        Once the source is exhausted, `_next_idx` reaching that exact value is
+        no longer required: if some index was skipped entirely (e.g. a chunk
+        whose bytes came back empty, see fetch_arrow_chunks_for_statement),
+        waiting for it forever would silently strand every higher-indexed
+        chunk already sitting in `_pending`. Draining the lowest remaining
+        index instead means a genuine gap costs row *order* past that point,
+        never lost rows."""
         while True:
             if self._next_idx in self._pending:
-                ready = self._pending.pop(self._next_idx)
-                self._next_idx += 1
-                return ready.to_table()
+                chunk = self._pop_pending(self._next_idx)
+                if self._next_idx not in self._pending:
+                    self._next_idx += 1
+                return chunk.to_table()
             if self._exhausted:
+                if self._pending:
+                    return self._pop_pending(min(self._pending)).to_table()
                 return None
             try:
                 chunk = await self._chunk_aiter.__anext__()
             except StopAsyncIteration:
                 self._exhausted = True
                 continue
-            self._pending[chunk.chunk_index] = chunk
+            self._pending.setdefault(chunk.chunk_index, []).append(chunk)
 
     async def _ensure_buffer(self, want: int) -> None:
         while (self._buffer is None or self._buffer.num_rows < want) and not self._exhausted:
@@ -233,6 +258,31 @@ class Cursor:
     async def fetchall_arrow(self) -> core.Table:
         return await self._require_result().fetchall_arrow()
 
+    def fetchall_streamed(self, *, total_timeout_s: float | None = None) -> AsyncIterator[Any]:
+        """Like fetchall(), but yields HEARTBEAT while pulling chunks instead
+        of blocking silently -- for a caller bridging e.g. an SSE connection
+        through the full download, not just the initial `execute_streamed`
+        wait for the statement to complete. Downloading many chunks for a
+        large result can itself take a while; `execute_streamed`'s own
+        heartbeats stop the moment the statement is ready, before any chunk
+        has actually been fetched. Yields HEARTBEAT zero or more times, then
+        the final `list[Row]`."""
+
+        async def _gen() -> AsyncIterator[Any]:
+            async for item in await_with_heartbeat(self.fetchall(), total_timeout_s=total_timeout_s):
+                yield item
+
+        return _gen()
+
+    def fetchall_arrow_streamed(self, *, total_timeout_s: float | None = None) -> AsyncIterator[Any]:
+        """Arrow-`Table` counterpart to fetchall_streamed -- see its docstring."""
+
+        async def _gen() -> AsyncIterator[Any]:
+            async for item in await_with_heartbeat(self.fetchall_arrow(), total_timeout_s=total_timeout_s):
+                yield item
+
+        return _gen()
+
     def __aiter__(self) -> Cursor:
         return self
 

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/tests/test_cursor.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import asyncio
+
 import httpx
 import pytest
 import respx
 
-from arrowbricks import HEARTBEAT, DatabricksClient
+from arrowbricks import HEARTBEAT, DatabricksClient, QueryTimeout
 from arrowbricks.cursor import Cursor
 
 
@@ -153,6 +155,60 @@ async def test_execute_streamed_emits_heartbeat_before_cursor_ready(mock_warehou
     assert len(await cursor.fetchall()) == 3
 
 
+@pytest.mark.asyncio
+@respx.mock
+async def test_fetchall_streamed_times_out_on_a_slow_chunk_download(warehouse_host_id):
+    """execute_streamed's own timeout only covers the wait for the statement
+    to become ready -- it stops the moment chunks are available to fetch,
+    before any chunk has actually been downloaded. A slow chunk download must
+    still be bounded, which is exactly what fetchall_streamed adds."""
+    host, warehouse_id = warehouse_host_id
+    respx.mock.get(f"{host}/api/2.0/sql/warehouses/{warehouse_id}").mock(
+        return_value=httpx.Response(200, json={"state": "RUNNING"})
+    )
+    respx.mock.post(f"{host}/api/2.0/sql/statements").mock(
+        return_value=httpx.Response(
+            200,
+            json={
+                "statement_id": "stmt-slow",
+                "status": {"state": "SUCCEEDED"},
+                "manifest": {"chunks": [{"chunk_index": 0, "row_count": 3}]},
+            },
+        )
+    )
+    respx.mock.get(f"{host}/api/2.0/sql/statements/stmt-slow/result/chunks/0").mock(
+        return_value=httpx.Response(200, json={"external_links": [{"external_link": f"{host}/_data/slow-chunk"}]})
+    )
+
+    async def _slow_chunk(request: httpx.Request) -> httpx.Response:
+        await asyncio.sleep(10)
+        raise AssertionError("unreachable -- test should time out first")
+
+    respx.mock.get(f"{host}/_data/slow-chunk").mock(side_effect=_slow_chunk)
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+    cursor = Cursor(client)
+    await cursor.execute("SELECT * FROM whatever")
+
+    with pytest.raises(QueryTimeout):
+        async for _ in cursor.fetchall_streamed(total_timeout_s=0.05):
+            pass
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_fetchall_streamed_yields_result_after_zero_or_more_heartbeats(mock_warehouse, warehouse_host_id):
+    host, warehouse_id = warehouse_host_id
+    mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=3)
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+    cursor = Cursor(client)
+    await cursor.execute("SELECT * FROM whatever")
+
+    items = [item async for item in cursor.fetchall_streamed(total_timeout_s=5)]
+
+    assert all(item is HEARTBEAT for item in items[:-1])
+    assert [r[0] for r in items[-1]] == [0, 1, 2]
+
+
 @pytest.mark.asyncio
 async def test_fetch_before_execute_raises():
     client = DatabricksClient("https://fake", "wh", token="test-token")

arrowbricks-0.2.0/tests/test_result_set.py

@@ -0,0 +1,74 @@
+"""Regression tests for _ResultSet's chunk reorder buffer -- specifically the
+two failure modes a naive dict[int, chunk] keyed-by-index buffer gets wrong:
+a chunk_index that shows up more than once (DatabricksClient._fetch_chunk_index
+gathers over possibly-multiple external_links per chunk), and a chunk_index
+that never shows up at all (fetch_arrow_chunks_for_statement skips falsy
+blobs). Both must never lose rows -- see cursor.py's _pull_one_chunk_table
+docstring."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+
+import pytest
+
+from arrowbricks._streaming import ReplayableArrowChunk
+from arrowbricks.cursor import _ResultSet
+
+
+async def _aiter(chunks: list[ReplayableArrowChunk]) -> AsyncIterator[ReplayableArrowChunk]:
+    for c in chunks:
+        yield c
+
+
+def _ids(chunk_bytes_builder, lo: int, hi: int, chunk_index: int) -> ReplayableArrowChunk:
+    return ReplayableArrowChunk(chunk_bytes_builder(lo, hi), chunk_index=chunk_index)
+
+
+@pytest.mark.asyncio
+async def test_duplicate_chunk_index_keeps_both_blobs(chunk_bytes_builder):
+    """Two separate blobs arriving under the SAME chunk_index (multiple
+    external_links for one chunk) must both survive, not overwrite each
+    other."""
+    chunks = [
+        _ids(chunk_bytes_builder, 0, 2, chunk_index=0),
+        _ids(chunk_bytes_builder, 5, 7, chunk_index=0),  # same index, second blob
+        _ids(chunk_bytes_builder, 2, 4, chunk_index=1),
+    ]
+    result = _ResultSet(schema=None, chunk_aiter=_aiter(chunks))
+
+    table = await result.fetchall_arrow()
+
+    assert sorted(table.column(0).combine_chunks().to_pylist()) == [0, 1, 2, 3, 5, 6]
+
+
+@pytest.mark.asyncio
+async def test_gap_in_chunk_index_does_not_strand_later_chunks(chunk_bytes_builder):
+    """chunk_index 1 never arrives at all (e.g. its bytes came back empty and
+    got filtered upstream) -- chunk 2's rows must still come out, not get
+    stranded behind a hole that never fills in."""
+    chunks = [
+        _ids(chunk_bytes_builder, 0, 2, chunk_index=0),
+        _ids(chunk_bytes_builder, 2, 4, chunk_index=2),  # index 1 is missing
+    ]
+    result = _ResultSet(schema=None, chunk_aiter=_aiter(chunks))
+
+    table = await result.fetchall_arrow()
+
+    assert sorted(table.column(0).combine_chunks().to_pylist()) == [0, 1, 2, 3]
+
+
+@pytest.mark.asyncio
+async def test_out_of_order_arrival_still_preserves_row_order(chunk_bytes_builder):
+    """The common case (no gaps, no duplicates, just network completion
+    order != chunk_index order) must still come out in chunk_index order."""
+    chunks = [
+        _ids(chunk_bytes_builder, 4, 6, chunk_index=2),
+        _ids(chunk_bytes_builder, 0, 2, chunk_index=0),
+        _ids(chunk_bytes_builder, 2, 4, chunk_index=1),
+    ]
+    result = _ResultSet(schema=None, chunk_aiter=_aiter(chunks))
+
+    table = await result.fetchall_arrow()
+
+    assert table.column(0).combine_chunks().to_pylist() == [0, 1, 2, 3, 4, 5]

arrowbricks-0.2.0/tests/test_streaming.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import json
+
+import httpx
+import pytest
+import respx
+
+from arrowbricks import DatabricksClient, stream_query_json
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_stream_query_json_preserves_order_despite_out_of_order_chunks(mock_warehouse, warehouse_host_id):
+    host, warehouse_id = warehouse_host_id
+    # Chunks resolve in REVERSE completion order (see conftest's
+    # reverse_arrival) -- this is the scenario that silently breaks ORDER BY
+    # if a caller doesn't reorder before emitting.
+    mock_warehouse(respx.mock, n_chunks=4, rows_per_chunk=5, reverse_arrival=True)
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+
+    rows = [json.loads(row) async for row in stream_query_json(client, "SELECT * FROM whatever ORDER BY id")]
+
+    assert [r["id"] for r in rows] == list(range(20))
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_stream_query_json_row_shape(mock_warehouse, warehouse_host_id):
+    host, warehouse_id = warehouse_host_id
+    mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=3)
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+
+    rows = [json.loads(row) async for row in stream_query_json(client, "SELECT * FROM whatever")]
+
+    assert rows == [
+        {"id": 0, "label": "row_0"},
+        {"id": 1, "label": "row_1"},
+        {"id": 2, "label": "row_2"},
+    ]
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_stream_query_json_survives_duplicate_index_and_a_gap(warehouse_host_id, chunk_bytes_builder):
+    """chunk_index 0 has TWO external_links (a real, if uncommon, shape --
+    DatabricksClient._fetch_chunk_index gathers over all of them), and index 1
+    never shows up at all. Neither must lose rows -- see cursor.py's
+    _pull_one_chunk_table / this module's stream_query_json docstrings."""
+    host, warehouse_id = warehouse_host_id
+    statement_id = "stmt-dup-gap"
+    respx.mock.get(f"{host}/api/2.0/sql/warehouses/{warehouse_id}").mock(
+        return_value=httpx.Response(200, json={"state": "RUNNING"})
+    )
+    respx.mock.post(f"{host}/api/2.0/sql/statements").mock(
+        return_value=httpx.Response(
+            200,
+            json={
+                "statement_id": statement_id,
+                "status": {"state": "SUCCEEDED"},
+                # index 1 deliberately absent -- a genuine gap.
+                "manifest": {"chunks": [{"chunk_index": 0, "row_count": 4}, {"chunk_index": 2, "row_count": 2}]},
+            },
+        )
+    )
+    respx.mock.get(f"{host}/api/2.0/sql/statements/{statement_id}/result/chunks/0").mock(
+        return_value=httpx.Response(
+            200,
+            json={"external_links": [{"external_link": f"{host}/_data/0a"}, {"external_link": f"{host}/_data/0b"}]},
+        )
+    )
+    respx.mock.get(f"{host}/api/2.0/sql/statements/{statement_id}/result/chunks/2").mock(
+        return_value=httpx.Response(200, json={"external_links": [{"external_link": f"{host}/_data/2"}]})
+    )
+    respx.mock.get(f"{host}/_data/0a").mock(return_value=httpx.Response(200, content=chunk_bytes_builder(0, 2)))
+    respx.mock.get(f"{host}/_data/0b").mock(return_value=httpx.Response(200, content=chunk_bytes_builder(10, 12)))
+    respx.mock.get(f"{host}/_data/2").mock(return_value=httpx.Response(200, content=chunk_bytes_builder(20, 22)))
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+
+    rows = [json.loads(row) async for row in stream_query_json(client, "SELECT * FROM whatever")]
+
+    assert sorted(r["id"] for r in rows) == [0, 1, 10, 11, 20, 21]

{arrowbricks-0.1.2 → arrowbricks-0.2.0}/uv.lock

@@ -145,7 +145,7 @@ wheels = [
 
 [[package]]
 name = "arrowbricks"
-version = "0.1.2"
+version = "0.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "arro3-core" },

arrowbricks-0.1.2/tests/test_streaming.py

@@ -1,39 +0,0 @@
-from __future__ import annotations
-
-import json
-
-import pytest
-import respx
-
-from arrowbricks import DatabricksClient, stream_query_json
-
-
-@pytest.mark.asyncio
-@respx.mock
-async def test_stream_query_json_preserves_order_despite_out_of_order_chunks(mock_warehouse, warehouse_host_id):
-    host, warehouse_id = warehouse_host_id
-    # Chunks resolve in REVERSE completion order (see conftest's
-    # reverse_arrival) -- this is the scenario that silently breaks ORDER BY
-    # if a caller doesn't reorder before emitting.
-    mock_warehouse(respx.mock, n_chunks=4, rows_per_chunk=5, reverse_arrival=True)
-    client = DatabricksClient(host, warehouse_id, token="test-token")
-
-    rows = [json.loads(row) async for row in stream_query_json(client, "SELECT * FROM whatever ORDER BY id")]
-
-    assert [r["id"] for r in rows] == list(range(20))
-
-
-@pytest.mark.asyncio
-@respx.mock
-async def test_stream_query_json_row_shape(mock_warehouse, warehouse_host_id):
-    host, warehouse_id = warehouse_host_id
-    mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=3)
-    client = DatabricksClient(host, warehouse_id, token="test-token")
-
-    rows = [json.loads(row) async for row in stream_query_json(client, "SELECT * FROM whatever")]
-
-    assert rows == [
-        {"id": 0, "label": "row_0"},
-        {"id": 1, "label": "row_1"},
-        {"id": 2, "label": "row_2"},
-    ]