arrowbricks 0.1.2__tar.gz → 0.3.0__tar.gz

{arrowbricks-0.1.2 → arrowbricks-0.3.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.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 0.1.2
+Version: 0.3.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.3.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.3.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.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "0.1.2"
+version = "0.3.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.3.0/scripts/benchmark_client_reuse.py

@@ -0,0 +1,126 @@
+"""Before/after benchmark for the persistent-client + warehouse-check-cache
+optimizations (see client.py's _get_http_client/_ensure_warehouse_running).
+
+Runs the same small query N times two ways against a REAL Databricks
+warehouse:
+  - "cold": a fresh DatabricksClient per query (simulates the old
+    per-call-httpx.AsyncClient() + per-call warehouse-status-GET behavior).
+  - "warm": one DatabricksClient reused across all N queries (the new
+    default behavior).
+
+Usage:
+    export DATABRICKS_HOST=...           # or falls back to `databricks auth env`
+    export DATABRICKS_WAREHOUSE_ID=...
+    export DATABRICKS_TOKEN=...          # a personal access token
+    uv run python scripts/benchmark_client_reuse.py [--n 10] [--sql "SELECT 1"]
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import shutil
+import statistics
+import subprocess
+import sys
+import time
+
+sys.path.insert(0, "src")
+
+from arrowbricks import DatabricksClient  # noqa: E402
+
+
+def _databricks_auth_env() -> dict[str, str]:
+    """Falls back to `databricks auth env` (the CLI's own resolved
+    credentials) if DATABRICKS_HOST/TOKEN aren't set directly -- lets this
+    script work whether the user set env vars or ran `databricks configure`."""
+    if os.environ.get("DATABRICKS_HOST") and (
+        os.environ.get("DATABRICKS_TOKEN") or os.environ.get("DATABRICKS_CLIENT_ID")
+    ):
+        return dict(os.environ)
+    databricks_cli = shutil.which("databricks")
+    if not databricks_cli:
+        return dict(os.environ)
+    try:
+        out = subprocess.run(  # noqa: S603 -- fixed args, no untrusted input
+            [databricks_cli, "auth", "env"], capture_output=True, text=True, check=True
+        ).stdout
+    except Exception:
+        return dict(os.environ)
+    env = dict(os.environ)
+    for line in out.splitlines():
+        line = line.strip()
+        if line.startswith("export "):
+            line = line[len("export ") :]
+        if "=" in line:
+            k, _, v = line.partition("=")
+            env.setdefault(k.strip(), v.strip().strip('"'))
+    return env
+
+
+async def _run_n_cold(host: str, warehouse_id: str, token: str, sql: str, n: int) -> list[float]:
+    """Old behavior: a fresh client (and fresh http connection pool + a
+    fresh warehouse-status check) per query."""
+    times = []
+    for _ in range(n):
+        client = DatabricksClient(host, warehouse_id, token=token)
+        start = time.perf_counter()
+        await client.execute_json_statement(sql)
+        times.append(time.perf_counter() - start)
+        await client.aclose()
+    return times
+
+
+async def _run_n_warm(host: str, warehouse_id: str, token: str, sql: str, n: int) -> list[float]:
+    """New behavior: one client, one connection pool, warehouse status
+    checked once then cached for the TTL window."""
+    times = []
+    async with DatabricksClient(host, warehouse_id, token=token) as client:
+        for _ in range(n):
+            start = time.perf_counter()
+            await client.execute_json_statement(sql)
+            times.append(time.perf_counter() - start)
+    return times
+
+
+def _report(label: str, times: list[float]) -> None:
+    print(f"{label}: n={len(times)} mean={statistics.mean(times):.3f}s median={statistics.median(times):.3f}s "
+          f"min={min(times):.3f}s max={max(times):.3f}s total={sum(times):.3f}s")
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int, default=10)
+    parser.add_argument("--sql", default="SELECT 1")
+    args = parser.parse_args()
+
+    env = _databricks_auth_env()
+    host = env.get("DATABRICKS_HOST", "adb-8956277663194228.8.azuredatabricks.net")
+    warehouse_id = env.get("DATABRICKS_WAREHOUSE_ID", "c397040753b46093")
+    token = env.get("DATABRICKS_TOKEN")
+    if not token:
+        raise SystemExit(
+            "No DATABRICKS_TOKEN found (env var or `databricks auth env`). "
+            "Set one and re-run -- this script needs a real, working credential."
+        )
+
+    print(f"host={host} warehouse_id={warehouse_id} n={args.n} sql={args.sql!r}\n")
+
+    # Warm the warehouse first (cold-start time shouldn't pollute either measurement).
+    warmup = DatabricksClient(host, warehouse_id, token=token)
+    await warmup.execute_json_statement(args.sql)
+    await warmup.aclose()
+
+    cold = await _run_n_cold(host, warehouse_id, token, args.sql, args.n)
+    _report("cold (fresh client per query)", cold)
+
+    warm = await _run_n_warm(host, warehouse_id, token, args.sql, args.n)
+    _report("warm (one reused client)     ", warm)
+
+    speedup = statistics.mean(cold) / statistics.mean(warm) if statistics.mean(warm) > 0 else float("inf")
+    print(f"\nmean speedup: {speedup:.2f}x")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

arrowbricks-0.3.0/scripts/benchmark_simulated.py

@@ -0,0 +1,157 @@
+"""SIMULATED before/after benchmark for the persistent-client +
+warehouse-check-cache optimizations (client.py's _get_http_client /
+_ensure_warehouse_running) -- runs entirely against a respx-mocked
+transport, no real Databricks warehouse or credentials needed.
+
+This is NOT a measurement against a real network -- respx intercepts at the
+transport layer, so there's no actual socket, TCP handshake, or TLS
+negotiation happening. To make the comparison meaningful anyway, this script
+injects realistic artificial latency for exactly the two things the
+optimizations remove:
+
+  - `_HANDSHAKE_COST_S`: paid once per NEW httpx.AsyncClient (simulates the
+    TCP+TLS handshake a brand-new connection to the Databricks host pays).
+    "cold" mode creates a fresh DatabricksClient (and thus a fresh
+    httpx.AsyncClient) per query, so it pays this every time. "warm" mode
+    reuses one DatabricksClient across all queries, paying it exactly once.
+  - `_WAREHOUSE_CHECK_COST_S`: the warehouse-status GET's round-trip time.
+    This one isn't injected by this script at all -- it falls out of the
+    REAL client.py code: "cold" mode's fresh client has no cached
+    confirmation, so _ensure_warehouse_running does the GET every time;
+    "warm" mode's cache (see warehouse_confirmed_running_ttl_s) skips it
+    after the first call. The respx route itself sleeps
+    _WAREHOUSE_CHECK_COST_S before responding, same for both modes -- the
+    difference in total time comes only from how many times each mode hits
+    that route.
+
+Every other mocked endpoint (statement submission, chunk resolve/fetch) has
+zero added delay in both modes, since that cost doesn't differ between them
+and would only dilute the comparison.
+
+Usage:
+    uv run python scripts/benchmark_simulated.py [--n 20]
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import statistics
+import sys
+import time
+
+sys.path.insert(0, "src")
+sys.path.insert(0, "tests")
+
+import httpx  # noqa: E402
+import respx  # noqa: E402
+from conftest import HOST, WAREHOUSE_ID, build_chunk_bytes  # noqa: E402
+
+from arrowbricks import DatabricksClient  # noqa: E402
+
+_HANDSHAKE_COST_S = 0.06  # a fresh HTTPS connection's TCP+TLS setup, typical for a cross-region call
+_WAREHOUSE_CHECK_COST_S = 0.04  # one small API round-trip
+
+
+def _install_routes(router: respx.Router, statement_id: str) -> None:
+    router.get(f"{HOST}/api/2.0/sql/warehouses/{WAREHOUSE_ID}").mock(
+        side_effect=_slow_response(_WAREHOUSE_CHECK_COST_S, {"state": "RUNNING"})
+    )
+    router.post(f"{HOST}/api/2.0/sql/statements").mock(
+        return_value=httpx.Response(
+            200,
+            json={
+                "statement_id": statement_id,
+                "status": {"state": "SUCCEEDED"},
+                "manifest": {"chunks": [{"chunk_index": 0, "row_count": 1}]},
+            },
+        )
+    )
+    router.get(url__regex=rf"{HOST}/api/2\.0/sql/statements/{statement_id}/result/chunks/\d+").mock(
+        return_value=httpx.Response(200, json={"external_links": [{"external_link": f"{HOST}/_data/chunk-0"}]})
+    )
+    router.get(url__regex=rf"{HOST}/_data/chunk-\d+").mock(
+        return_value=httpx.Response(200, content=build_chunk_bytes(0, 1))
+    )
+
+
+def _slow_response(delay_s: float, body: dict) -> object:
+    async def _handler(request: httpx.Request) -> httpx.Response:
+        await asyncio.sleep(delay_s)
+        return httpx.Response(200, json=body)
+
+    return _handler
+
+
+async def _timed_query(client: DatabricksClient, sql: str) -> float:
+    """Charges _HANDSHAKE_COST_S the first time this client's shared
+    httpx.AsyncClient gets its first real use, simulating a fresh
+    connection's TCP+TLS setup -- a real socket would pay this on its own,
+    respx's in-memory transport doesn't, so it's added explicitly here."""
+    is_first_use = client._http is None
+    start = time.perf_counter()
+    if is_first_use:
+        await asyncio.sleep(_HANDSHAKE_COST_S)
+    await client.execute_json_statement(sql)
+    return time.perf_counter() - start
+
+
+async def _run_cold(n: int, sql: str) -> list[float]:
+    times = []
+    with respx.mock:
+        _install_routes(respx.mock, "stmt-cold")
+        for _ in range(n):
+            token = "fake-token"  # noqa: S105 -- fake token, mocked transport only
+            client = DatabricksClient(HOST, WAREHOUSE_ID, token=token)
+            times.append(await _timed_query(client, sql))
+            await client.aclose()
+    return times
+
+
+async def _run_warm(n: int, sql: str) -> list[float]:
+    times = []
+    with respx.mock:
+        _install_routes(respx.mock, "stmt-warm")
+        token = "fake-token"  # noqa: S105 -- fake token, mocked transport only
+        async with DatabricksClient(HOST, WAREHOUSE_ID, token=token) as client:
+            for _ in range(n):
+                times.append(await _timed_query(client, sql))
+    return times
+
+
+def _report(label: str, times: list[float]) -> None:
+    print(
+        f"{label}: n={len(times)} mean={statistics.mean(times) * 1000:.1f}ms "
+        f"median={statistics.median(times) * 1000:.1f}ms total={sum(times):.3f}s"
+    )
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int, default=20)
+    args = parser.parse_args()
+
+    print("*** SIMULATED benchmark -- respx-mocked transport, no real Databricks call. ***")
+    print(f"Injected costs: fresh-connection handshake={_HANDSHAKE_COST_S * 1000:.0f}ms, "
+          f"warehouse-status round-trip={_WAREHOUSE_CHECK_COST_S * 1000:.0f}ms\n")
+
+    cold = await _run_cold(args.n, "SELECT 1")
+    _report("cold (old behavior: fresh client + warehouse-check per query)", cold)
+
+    warm = await _run_warm(args.n, "SELECT 1")
+    _report("warm (new behavior: one reused client)                      ", warm)
+
+    saved_per_query_ms = (statistics.mean(cold) - statistics.mean(warm)) * 1000
+    speedup = statistics.mean(cold) / statistics.mean(warm) if statistics.mean(warm) > 0 else float("inf")
+    print(f"\nmean saved per query (after the first): ~{saved_per_query_ms:.1f}ms")
+    print(f"mean speedup: {speedup:.2f}x")
+    print(
+        "\nNote: this isolates exactly the two mechanisms changed (connection reuse, "
+        "warehouse-check caching) with realistic but assumed latency values -- it does "
+        "not measure real Databricks/network round-trip times, statement execution time, "
+        "or chunk-fetch time, none of which changed."
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

{arrowbricks-0.1.2 → arrowbricks-0.3.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.3.0}/src/arrowbricks/client.py

@@ -82,6 +82,7 @@ class DatabricksClient:
         wait_timeout: str = "30s",
         chunk_fetch_concurrency: int = 6,
         warehouse_start_timeout: float = 300.0,
+        warehouse_confirmed_running_ttl_s: float = 30.0,
     ) -> None:
         if not token and not token_provider:
             raise ValueError("DatabricksClient needs either `token` or `token_provider`")
@@ -101,6 +102,40 @@ class DatabricksClient:
         # slow consumer.
         self.chunk_fetch_concurrency = chunk_fetch_concurrency
         self.warehouse_start_timeout = warehouse_start_timeout
+        # See _ensure_warehouse_running: a warm warehouse doesn't need its
+        # RUNNING state re-verified on every single statement -- one round
+        # trip saved per query once confirmed within this window.
+        self._warehouse_confirmed_running_ttl_s = warehouse_confirmed_running_ttl_s
+        self._warehouse_confirmed_running_at: float | None = None
+        # One shared connection pool for this client's whole lifetime instead
+        # of a fresh httpx.AsyncClient() (and its own TCP+TLS handshake) per
+        # call -- every method below hits the same Databricks host repeatedly,
+        # so keep-alive/pooling actually pays off across calls. Presigned
+        # external-link downloads (a different host: blob storage) still each
+        # get their own connection from this same pool as needed; httpx pools
+        # per-host internally, so sharing one client across hosts is safe.
+        self._http: httpx.AsyncClient | None = None
+        self._http_lock = asyncio.Lock()
+
+    async def _get_http_client(self) -> httpx.AsyncClient:
+        if self._http is None:
+            async with self._http_lock:
+                if self._http is None:
+                    self._http = httpx.AsyncClient()
+        return self._http
+
+    async def aclose(self) -> None:
+        """Closes the shared connection pool. Safe to call even if no request
+        was ever made (no-op) or more than once."""
+        if self._http is not None:
+            await self._http.aclose()
+            self._http = None
+
+    async def __aenter__(self) -> DatabricksClient:
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.aclose()
 
     async def _bearer_token(self) -> str:
         if self._token is not None:
@@ -135,12 +170,23 @@ class DatabricksClient:
         implicit auto-start. A cold warehouse's catalog credential cache needs
         a moment to catch up right after startup -- submitting straight into
         that window is a common source of transient, identity-scoped 403s.
-        Fast path: a single GET when already RUNNING, so this adds no
-        meaningful overhead once warm."""
+
+        Skips the check entirely if RUNNING was already confirmed within
+        `_warehouse_confirmed_running_ttl_s` -- a warm, always-on warehouse
+        doesn't need re-verifying on every single statement; that GET is a
+        full round trip that buys nothing once already known-good."""
+        now = time.monotonic()
+        if (
+            self._warehouse_confirmed_running_at is not None
+            and now - self._warehouse_confirmed_running_at < self._warehouse_confirmed_running_ttl_s
+        ):
+            return
+
         url = f"{self._host}/api/2.0/sql/warehouses/{self.warehouse_id}"
         resp = await self._authed_request(client, "GET", url)
         state = resp.json().get("state")
         if state == "RUNNING":
+            self._warehouse_confirmed_running_at = time.monotonic()
             return
 
         if state == "STOPPED":
@@ -152,6 +198,7 @@ class DatabricksClient:
             resp = await self._authed_request(client, "GET", url)
             state = resp.json().get("state")
             if state == "RUNNING":
+                self._warehouse_confirmed_running_at = time.monotonic()
                 return
         # Falls through and lets the statement submission itself surface
         # whatever's actually wrong -- proceeding anyway rather than raising
@@ -187,21 +234,21 @@ class DatabricksClient:
         if parameters:
             body["parameters"] = parameters
 
-        async with httpx.AsyncClient() as client:
-            await self._ensure_warehouse_running(client)
-            resp = await self._authed_request(client, "POST", f"{self._host}/api/2.0/sql/statements", json=body)
-            data = resp.json()
+        client = await self._get_http_client()
+        await self._ensure_warehouse_running(client)
+        resp = await self._authed_request(client, "POST", f"{self._host}/api/2.0/sql/statements", json=body)
+        data = resp.json()
 
+        status = data.get("status", {})
+        while status.get("state") not in _TERMINAL_STATES:
+            statement_id = data["statement_id"]
+            await asyncio.sleep(_POLL_INTERVAL_S)
+            resp = await self._authed_request(client, "GET", f"{self._host}/api/2.0/sql/statements/{statement_id}")
+            data = resp.json()
             status = data.get("status", {})
-            while status.get("state") not in _TERMINAL_STATES:
-                statement_id = data["statement_id"]
-                await asyncio.sleep(_POLL_INTERVAL_S)
-                resp = await self._authed_request(client, "GET", f"{self._host}/api/2.0/sql/statements/{statement_id}")
-                data = resp.json()
-                status = data.get("status", {})
 
-            _raise_for_failed(status)
-            return data["statement_id"], data.get("manifest") or {}
+        _raise_for_failed(status)
+        return data["statement_id"], data.get("manifest") or {}
 
     async def execute_arrow_statement(
         self,
@@ -249,26 +296,26 @@ class DatabricksClient:
         overwriting anything already there. `volume_path` is caller-supplied
         in full (e.g. `/Volumes/my_catalog/my_schema/my_volume/some/file.parquet`)
         -- this package has no knowledge of any specific catalog/schema/volume."""
-        async with httpx.AsyncClient() as client:
-            await self._authed_request(
-                client,
-                "PUT",
-                f"{self._host}/api/2.0/fs/files{volume_path}",
-                params={"overwrite": "true"},
-                content_type="application/octet-stream",
-                content=data,
-            )
+        client = await self._get_http_client()
+        await self._authed_request(
+            client,
+            "PUT",
+            f"{self._host}/api/2.0/fs/files{volume_path}",
+            params={"overwrite": "true"},
+            content_type="application/octet-stream",
+            content=data,
+        )
 
     async def delete_volume_file(self, volume_path: str) -> None:
         """Deletes a file at `volume_path` (see upload_volume_file). A 404 is
         treated as success -- the file is already gone, which is fine for
         idempotent staging cleanup."""
-        async with httpx.AsyncClient() as client:
-            try:
-                await self._authed_request(client, "DELETE", f"{self._host}/api/2.0/fs/files{volume_path}")
-            except httpx.HTTPStatusError as exc:
-                if exc.response.status_code != 404:
-                    raise
+        client = await self._get_http_client()
+        try:
+            await self._authed_request(client, "DELETE", f"{self._host}/api/2.0/fs/files{volume_path}")
+        except httpx.HTTPStatusError as exc:
+            if exc.response.status_code != 404:
+                raise
 
     async def stream_chunks_by_index(
         self, statement_id: str, chunk_metas: list[dict[str, Any]]
@@ -281,11 +328,11 @@ class DatabricksClient:
         manifest didn't carry one) and its own chunk_index, so a caller that
         cares about the original row order (e.g. a query with ORDER BY) can
         restore it even though chunks can complete out of order."""
-        async with httpx.AsyncClient() as client:
-            async for blob, row_count, chunk_index in self._fetch_chunks_with_backpressure(
-                client, statement_id, chunk_metas
-            ):
-                yield blob, row_count, chunk_index
+        client = await self._get_http_client()
+        async for blob, row_count, chunk_index in self._fetch_chunks_with_backpressure(
+            client, statement_id, chunk_metas
+        ):
+            yield blob, row_count, chunk_index
 
     async def _fetch_link_bytes(self, client: httpx.AsyncClient, url: str) -> bytes:
         async def _do() -> bytes:

{arrowbricks-0.1.2 → arrowbricks-0.3.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.3.0}/tests/conftest.py

@@ -47,12 +47,14 @@ def mock_warehouse():
     Callers configure it via mock_warehouse(...) inside a `with respx.mock:`
     block (or use the `respx_router` param)."""
 
-    def _install(router: respx.Router, n_chunks: int, rows_per_chunk: int, *, reverse_arrival: bool = False) -> None:
+    def _install(
+        router: respx.Router, n_chunks: int, rows_per_chunk: int, *, reverse_arrival: bool = False
+    ) -> respx.Route:
         statement_id = "stmt-abc"
         chunks = [{"chunk_index": i, "row_count": rows_per_chunk} for i in range(n_chunks)]
         chunk_bytes = {i: build_chunk_bytes(i * rows_per_chunk, (i + 1) * rows_per_chunk) for i in range(n_chunks)}
 
-        router.get(f"{HOST}/api/2.0/sql/warehouses/{WAREHOUSE_ID}").mock(
+        warehouse_route = router.get(f"{HOST}/api/2.0/sql/warehouses/{WAREHOUSE_ID}").mock(
             return_value=httpx.Response(200, json={"state": "RUNNING"})
         )
         router.post(f"{HOST}/api/2.0/sql/statements").mock(
@@ -89,6 +91,7 @@ def mock_warehouse():
             return httpx.Response(200, content=chunk_bytes[idx])
 
         router.get(url__regex=rf"{HOST}/_data/chunk-\d+").mock(side_effect=serve_chunk_bytes)
+        return warehouse_route
 
     return _install
 

{arrowbricks-0.1.2 → arrowbricks-0.3.0}/tests/test_client.py

@@ -163,3 +163,84 @@ def test_requires_token_or_provider(warehouse_host_id):
     host, warehouse_id = warehouse_host_id
     with pytest.raises(ValueError, match="token"):
         DatabricksClient(host, warehouse_id)
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_shares_one_http_client_across_calls(mock_warehouse, warehouse_host_id):
+    """A DatabricksClient should reuse one httpx.AsyncClient (and its
+    connection pool) across statements and chunk fetches, not open a fresh
+    one per call -- see client.py's _get_http_client."""
+    host, warehouse_id = warehouse_host_id
+    mock_warehouse(respx.mock, n_chunks=2, rows_per_chunk=2)
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+
+    assert client._http is None
+    await client.execute_json_statement("SELECT 1")
+    first = client._http
+    assert first is not None
+    await client.execute_json_statement("SELECT 2")
+    assert client._http is first  # same instance reused, not recreated
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_aclose_closes_and_allows_reopening(mock_warehouse, warehouse_host_id):
+    host, warehouse_id = warehouse_host_id
+    mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=1)
+    client = DatabricksClient(host, warehouse_id, token="test-token")
+
+    await client.execute_json_statement("SELECT 1")
+    first = client._http
+    assert first is not None
+    await client.aclose()
+    assert client._http is None
+    assert first.is_closed
+
+    await client.execute_json_statement("SELECT 1")  # still usable after close
+    assert client._http is not None
+    assert client._http is not first
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_async_context_manager_closes_on_exit(mock_warehouse, warehouse_host_id):
+    host, warehouse_id = warehouse_host_id
+    mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=1)
+
+    async with DatabricksClient(host, warehouse_id, token="test-token") as client:
+        await client.execute_json_statement("SELECT 1")
+        http_client = client._http
+
+    assert http_client is not None
+    assert http_client.is_closed
+    assert client._http is None
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_warehouse_running_check_is_cached_across_statements(mock_warehouse, warehouse_host_id):
+    """Once RUNNING is confirmed, a second statement within the TTL window
+    shouldn't re-GET the warehouse status -- that round trip buys nothing on
+    an already-known-warm warehouse (see _ensure_warehouse_running)."""
+    host, warehouse_id = warehouse_host_id
+    warehouse_route = mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=1)
+    client = DatabricksClient(host, warehouse_id, token="test-token", warehouse_confirmed_running_ttl_s=60.0)
+
+    await client.execute_json_statement("SELECT 1")
+    assert warehouse_route.call_count == 1
+    await client.execute_json_statement("SELECT 2")
+    assert warehouse_route.call_count == 1  # still 1 -- cached, no second GET
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_warehouse_running_check_re_verifies_after_ttl_expires(mock_warehouse, warehouse_host_id):
+    host, warehouse_id = warehouse_host_id
+    warehouse_route = mock_warehouse(respx.mock, n_chunks=1, rows_per_chunk=1)
+    client = DatabricksClient(host, warehouse_id, token="test-token", warehouse_confirmed_running_ttl_s=0.0)
+
+    await client.execute_json_statement("SELECT 1")
+    assert warehouse_route.call_count == 1
+    await client.execute_json_statement("SELECT 2")
+    assert warehouse_route.call_count == 2  # TTL is 0 -- re-verified every time

{arrowbricks-0.1.2 → arrowbricks-0.3.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.3.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.3.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.3.0}/uv.lock

@@ -145,7 +145,7 @@ wheels = [
 
 [[package]]
 name = "arrowbricks"
-version = "0.1.2"
+version = "0.3.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"},
-    ]