arrowbricks 0.1.0__py3-none-any.whl

arrowbricks/__init__.py

@@ -0,0 +1,23 @@
+from ._streaming import (
+    HEARTBEAT,
+    QueryTimeout,
+    ReplayableArrowChunk,
+    await_with_heartbeat,
+    stream_query_json,
+    write_ipc_stream,
+)
+from .client import DatabricksClient
+from .cursor import Connection, Cursor, connect
+
+__all__ = [
+    "HEARTBEAT",
+    "Connection",
+    "Cursor",
+    "DatabricksClient",
+    "QueryTimeout",
+    "ReplayableArrowChunk",
+    "await_with_heartbeat",
+    "connect",
+    "stream_query_json",
+    "write_ipc_stream",
+]

arrowbricks/_streaming.py

@@ -0,0 +1,263 @@
+"""Chunk fetching, heartbeats, and Arrow (de)serialization -- everything that
+turns a DatabricksClient's raw chunk bytes into Arrow, via arro3 directly.
+Single Arrow engine, no pluggable backend: arrowbricks' whole reason to exist
+is "Databricks to Arrow via arro3", so there's nothing to make pluggable.
+
+`write_ipc_stream` always passes `compression=None` -- arro3's own default
+(`compression="LZ4"`) body-compresses every record batch, which DuckDB's own
+Arrow C Data Interface reader decompresses transparently but which other
+Arrow IPC readers may not support at all (observed: duckdb-wasm's browser-side
+decoder silently fails to parse it). Plain, uncompressed bodies are the safe
+default for a library whose bytes might end up read by anything."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import io
+from collections.abc import AsyncIterator, Awaitable
+from typing import Any, BinaryIO, TypeVar
+
+import arro3.core as core
+import arro3.io as aio
+
+from .client import DatabricksClient
+
+__all__ = [
+    "HEARTBEAT",
+    "QueryTimeout",
+    "ReplayableArrowChunk",
+    "await_with_heartbeat",
+    "fetch_arrow_chunks_for_statement",
+    "fetch_arrow_chunks_with_manifest",
+    "stream_query_json",
+    "write_ipc_stream",
+]
+
+# How often a caller waiting on a slow Databricks round-trip (warehouse cold
+# start, a long-running statement) gets a HEARTBEAT -- pick something well
+# under whatever idle-connection ceiling sits between your server and its
+# client (e.g. many PaaS load balancers cut an idle SSE connection around
+# ~230s) if you're forwarding these as keep-alive pings.
+_HEARTBEAT_INTERVAL_S = 15.0
+
+
+class QueryTimeout(RuntimeError):
+    """Raised when a query exceeds its `total_timeout_s`."""
+
+
+class _Heartbeat:
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "HEARTBEAT"
+
+
+HEARTBEAT = _Heartbeat()
+
+T = TypeVar("T")
+
+
+def write_ipc_stream(stream: Any, buf: BinaryIO) -> None:
+    """Writes `stream` (anything implementing `__arrow_c_stream__`, e.g. an
+    arro3 Table/RecordBatchReader or a ReplayableArrowChunk) as Arrow-IPC
+    stream bytes -- always uncompressed, see module docstring."""
+    aio.write_ipc_stream(stream, buf, compression=None)
+
+
+class ReplayableArrowChunk:
+    """Wraps one Arrow IPC-stream byte chunk so it can be handed to something
+    that calls `__arrow_c_stream__` more than once per relation (a schema
+    peek, then the actual scan -- DuckDB's registration path does this, for
+    one). A plain parsed stream is single-use and raises on the second call,
+    so this re-parses from the cached bytes every call instead. The bytes are
+    already fully in memory (just downloaded), so re-parsing costs a cheap
+    second pass, not a second network fetch."""
+
+    __slots__ = ("_data", "chunk_index", "declared_row_count")
+
+    def __init__(self, data: bytes, chunk_index: int, declared_row_count: int | None = None) -> None:
+        self._data = data
+        self.chunk_index = chunk_index
+        self.declared_row_count = declared_row_count
+
+    def __arrow_c_stream__(self, requested_schema: object = None) -> object:
+        reader = aio.read_ipc_stream(io.BytesIO(self._data))
+        return reader.__arrow_c_stream__(requested_schema)
+
+    def nbytes(self) -> int:
+        return len(self._data)
+
+    def to_table(self) -> core.Table:
+        """Parses this chunk's bytes into an arro3 Table -- the entry point
+        `cursor.py`'s result-set buffering uses to work with chunks at the
+        Arrow level (concat/slice) rather than re-parsing bytes per access."""
+        return core.Table.from_arrow(self)
+
+
+async def fetch_arrow_chunks_with_manifest(
+    client: DatabricksClient,
+    sql: str,
+    *,
+    catalog: str | None = None,
+    schema: str | None = None,
+    parameters: list[dict[str, Any]] | None = None,
+) -> tuple[str, dict[str, Any], AsyncIterator[ReplayableArrowChunk]]:
+    """Submits `sql` and returns (statement_id, manifest, chunk_iterator).
+    total_row_count/total_chunk_count are already in the manifest once the
+    statement succeeds, so callers needing a progress-bar target don't need a
+    separate preflight COUNT(*)."""
+    statement_id, manifest = await client.execute_arrow_statement(
+        sql, catalog=catalog, schema=schema, parameters=parameters
+    )
+    chunk_metas = manifest.get("chunks") or []
+    return statement_id, manifest, fetch_arrow_chunks_for_statement(client, statement_id, chunk_metas)
+
+
+async def fetch_arrow_chunks_for_statement(
+    client: DatabricksClient, statement_id: str, chunk_metas: list[dict[str, Any]]
+) -> AsyncIterator[ReplayableArrowChunk]:
+    async for chunk_bytes, row_count, chunk_index in client.stream_chunks_by_index(statement_id, chunk_metas):
+        if chunk_bytes:
+            yield ReplayableArrowChunk(chunk_bytes, chunk_index, declared_row_count=row_count)
+
+
+async def await_with_heartbeat(
+    aw: Awaitable[T], *, interval_s: float = _HEARTBEAT_INTERVAL_S, total_timeout_s: float | None = None
+) -> AsyncIterator[Any]:
+    """Wraps a single slow awaitable with periodic HEARTBEAT yields, so a
+    caller streaming this over e.g. SSE never goes silent for the whole wait.
+    Yields HEARTBEAT zero or more times, then yields the awaitable's real
+    result exactly once. Re-raises whatever `aw` raised, or QueryTimeout if
+    `total_timeout_s` elapses first."""
+    task: asyncio.Task[T] = asyncio.ensure_future(aw)
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + total_timeout_s if total_timeout_s is not None else None
+    try:
+        while not task.done():
+            wait_for = interval_s if deadline is None else min(interval_s, max(deadline - loop.time(), 0.0))
+            done, _ = await asyncio.wait({task}, timeout=wait_for)
+            if not done:
+                if deadline is not None and loop.time() >= deadline:
+                    task.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await task
+                    raise QueryTimeout(f"Query exceeded {total_timeout_s}s timeout")
+                yield HEARTBEAT
+        yield await task
+    finally:
+        if not task.done():
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+
+async def heartbeat_over_stream(
+    aiter: AsyncIterator[T], *, interval_s: float = _HEARTBEAT_INTERVAL_S, total_timeout_s: float | None = None
+) -> AsyncIterator[Any]:
+    """Like await_with_heartbeat, but for a stream of items rather than one
+    final result: yields HEARTBEAT whenever the wait for the *next* item
+    exceeds interval_s, and otherwise passes each item through as it arrives.
+    Used by stream_query_json (and cursor.py's execute_streamed) so a slow
+    chunk mid-stream -- not just a cold warehouse start before the first one
+    -- never lets a downstream SSE connection go silent for the whole wait."""
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + total_timeout_s if total_timeout_s is not None else None
+    it = aiter.__aiter__()
+    task: asyncio.Task[Any] | None = None
+    try:
+        while True:
+            if task is None:
+                task = asyncio.ensure_future(it.__anext__())
+            wait_for = interval_s if deadline is None else min(interval_s, max(deadline - loop.time(), 0.0))
+            done, _ = await asyncio.wait({task}, timeout=wait_for)
+            if not done:
+                if deadline is not None and loop.time() >= deadline:
+                    task.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await task
+                    raise QueryTimeout(f"Query exceeded {total_timeout_s}s timeout")
+                yield HEARTBEAT
+                continue
+            try:
+                yield task.result()
+            except StopAsyncIteration:
+                return
+            finally:
+                task = None
+    finally:
+        if task is not None and not task.done():
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+
+def windowed_sql(sql: str, *, row_limit: int | None, offset: int | None) -> str:
+    """Pushes LIMIT/OFFSET into the SQL submitted to Databricks -- a query
+    should never fetch more rows from the warehouse than the caller wants."""
+    if row_limit is None and not offset:
+        return sql
+    if row_limit is None:
+        return f"SELECT * FROM ({sql}) _q OFFSET {offset}"  # noqa: S608
+    if offset:
+        return f"SELECT * FROM ({sql}) _q LIMIT {row_limit} OFFSET {offset}"  # noqa: S608
+    return f"SELECT * FROM ({sql}) _q LIMIT {row_limit}"  # noqa: S608
+
+
+def _write_ndjson(chunk: ReplayableArrowChunk) -> bytes:
+    buf = io.BytesIO()
+    # explicit_nulls=True: arro3 omits null-valued keys by default -- without
+    # this a row's JSON shape would vary by which columns happen to be null
+    # in it.
+    aio.write_ndjson(chunk, buf, explicit_nulls=True)
+    return buf.getvalue()
+
+
+async def stream_query_json(
+    client: DatabricksClient,
+    sql: str,
+    *,
+    params: list[dict[str, Any]] | None = None,
+    row_limit: int | None = None,
+    offset: int | None = None,
+    catalog: str | None = None,
+    schema: str | None = None,
+    total_timeout_s: float | None = None,
+) -> AsyncIterator[Any]:
+    """Yields HEARTBEAT while waiting on Databricks, then each result row as a
+    ready-to-send JSON string (arro3's native `write_ndjson`) -- one per SSE
+    frame, say.
+
+    Unlike cursor.py's fetch methods, this registers and emits each Databricks
+    chunk AS IT ARRIVES rather than buffering the full result first -- the
+    first row reaches the caller after ~one chunk's fetch time, not the whole
+    statement's, and at most a handful of chunks' bytes (bounded by the
+    client's own chunk_fetch_concurrency) are ever held in memory at once, not
+    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.
+
+    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
+    bounds how much that is, so the overall memory-bounded-across-chunks
+    guarantee above still holds, just at chunk granularity."""
+    sql = windowed_sql(sql, row_limit=row_limit, offset=offset)
+    _statement_id, _manifest, chunk_iter = await fetch_arrow_chunks_with_manifest(
+        client, sql, catalog=catalog, schema=schema, parameters=params
+    )
+
+    pending: dict[int, ReplayableArrowChunk] = {}
+    next_idx = 0
+    loop = asyncio.get_running_loop()
+    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
+        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

arrowbricks/client.py

@@ -0,0 +1,374 @@
+"""Async client for the Databricks SQL Statement Execution API, tuned for
+pulling large results out as Arrow-IPC chunks rather than running a single
+small query. No dependency on Arrow at all -- that lives in `_streaming.py`
+and `cursor.py`, so this client is reusable by any caller regardless of what
+they do with the raw chunk bytes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import Any, TypeVar, cast
+
+import httpx
+
+_POLL_INTERVAL_S = 2.0
+_TERMINAL_STATES = {"SUCCEEDED", "FAILED", "CANCELED", "CLOSED"}
+_RETRY_ATTEMPTS = 6
+_RETRY_MAX_WAIT_S = 20.0
+
+TokenProvider = Callable[[], "str | Awaitable[str]"]
+
+_T = TypeVar("_T")
+
+
+def _is_transient_error(exc: BaseException) -> bool:
+    """401/403/429/5xx are treated as transient and retried -- these have been
+    observed in practice to resolve themselves (e.g. a just-granted permission
+    or a just-refreshed token not yet visible to the request that raced it)
+    rather than indicating a persistent problem. httpx.TimeoutException/
+    NetworkError are deliberately excluded: a genuinely stalled connection
+    should fail fast on the caller's own timeout, not be retried here."""
+    return isinstance(exc, httpx.HTTPStatusError) and (
+        exc.response.status_code in (401, 403, 408, 429) or exc.response.status_code >= 500
+    )
+
+
+async def _retry_call(fn: Callable[[], Awaitable[_T]]) -> _T:
+    """Retry fn() with exponential backoff (1s, 2s, 4s, ... capped at
+    _RETRY_MAX_WAIT_S) on transient errors, up to _RETRY_ATTEMPTS total tries.
+    A hand-rolled loop instead of tenacity -- this is the only retry pattern
+    in the whole client, so a dependency for it isn't worth it."""
+    attempt = 0
+    while True:
+        try:
+            return await fn()
+        except Exception as exc:
+            if attempt == _RETRY_ATTEMPTS - 1 or not _is_transient_error(exc):
+                raise
+            await asyncio.sleep(min(_RETRY_MAX_WAIT_S, 2.0**attempt))
+            attempt += 1
+
+
+def _raise_for_failed(status: dict[str, Any]) -> None:
+    state = status.get("state")
+    if state == "FAILED":
+        err = status.get("error", {})
+        raise RuntimeError(f"Databricks statement failed [{err.get('error_code')}]: {err.get('message')}")
+    if state == "CANCELED":
+        raise RuntimeError("Databricks statement was canceled")
+
+
+class DatabricksClient:
+    """One Databricks SQL warehouse endpoint. Auth is entirely bring-your-own:
+    pass a static `token`, or a `token_provider` callable (sync or async) that
+    returns one -- this client has no opinion on *how* you get a token (a
+    personal access token, an OAuth M2M flow, a cloud-provider credential
+    chain) and no cloud-SDK dependency baked in. If your provider is expensive
+    to call (e.g. shells out to a CLI), cache/refresh inside it -- this client
+    calls it on every request and does no caching of its own."""
+
+    def __init__(
+        self,
+        host: str,
+        warehouse_id: str,
+        *,
+        token: str | None = None,
+        token_provider: TokenProvider | None = None,
+        http_timeout: float = 60.0,
+        wait_timeout: str = "30s",
+        chunk_fetch_concurrency: int = 6,
+        warehouse_start_timeout: float = 300.0,
+    ) -> None:
+        if not token and not token_provider:
+            raise ValueError("DatabricksClient needs either `token` or `token_provider`")
+        self._host = host.rstrip("/")
+        if not self._host.startswith("https://"):
+            self._host = f"https://{self._host}"
+        self.warehouse_id = warehouse_id
+        self._token = token
+        self._token_provider = token_provider
+        self.http_timeout = http_timeout
+        self.wait_timeout = wait_timeout
+        # Each concurrent fetch slot (plus its matching queued-but-unconsumed
+        # slot, see _fetch_chunks_with_backpressure) holds a whole chunk's raw
+        # bytes in plain Python memory. Keep this modest by default -- it's
+        # I/O-bound network fetching, so parallelism past single digits buys
+        # little, while every extra slot is more memory held hostage behind a
+        # slow consumer.
+        self.chunk_fetch_concurrency = chunk_fetch_concurrency
+        self.warehouse_start_timeout = warehouse_start_timeout
+
+    async def _bearer_token(self) -> str:
+        if self._token is not None:
+            return self._token
+        assert self._token_provider is not None  # noqa: S101 -- enforced in __init__, not a test assertion
+        result = self._token_provider()
+        if inspect.isawaitable(result):
+            return await cast("Awaitable[str]", result)
+        return cast(str, result)
+
+    async def _headers(self, content_type: str = "application/json") -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {await self._bearer_token()}",
+            "Content-Type": content_type,
+        }
+
+    async def _authed_request(
+        self, client: httpx.AsyncClient, method: str, url: str, *, content_type: str = "application/json", **kwargs: Any
+    ) -> httpx.Response:
+        async def _do() -> httpx.Response:
+            resp = await client.request(
+                method, url, headers=await self._headers(content_type), timeout=self.http_timeout, **kwargs
+            )
+            resp.raise_for_status()
+            return resp
+
+        return await _retry_call(_do)
+
+    async def _ensure_warehouse_running(self, client: httpx.AsyncClient) -> None:
+        """Explicitly wait for the warehouse to reach RUNNING before submitting
+        a statement, rather than relying on the Statement Execution API's own
+        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."""
+        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":
+            return
+
+        if state == "STOPPED":
+            await self._authed_request(client, "POST", f"{url}/start")
+
+        deadline = time.monotonic() + self.warehouse_start_timeout
+        while time.monotonic() < deadline:
+            await asyncio.sleep(_POLL_INTERVAL_S)
+            resp = await self._authed_request(client, "GET", url)
+            state = resp.json().get("state")
+            if state == "RUNNING":
+                return
+        # Falls through and lets the statement submission itself surface
+        # whatever's actually wrong -- proceeding anyway rather than raising
+        # here avoids a false negative if the warehouse comes up moments later.
+
+    async def _execute_statement(
+        self,
+        statement: str,
+        *,
+        format: str,
+        catalog: str | None,
+        schema: str | None,
+        wait_timeout: str | None,
+        parameters: list[dict[str, Any]] | None,
+    ) -> tuple[str, dict[str, Any]]:
+        """Submit an EXTERNAL_LINKS statement in the given result format and
+        poll until terminal. Returns (statement_id, manifest). `parameters`,
+        if given, is Databricks' own named-parameter format --
+        [{"name": ..., "value": ..., "type": ...}] bound against `:name`
+        markers in `statement`."""
+        body: dict[str, Any] = {
+            "warehouse_id": self.warehouse_id,
+            "statement": statement,
+            "disposition": "EXTERNAL_LINKS",
+            "format": format,
+            "wait_timeout": wait_timeout or self.wait_timeout,
+            "on_wait_timeout": "CONTINUE",
+        }
+        if catalog:
+            body["catalog"] = catalog
+        if schema:
+            body["schema"] = schema
+        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()
+
+            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 {}
+
+    async def execute_arrow_statement(
+        self,
+        statement: str,
+        *,
+        catalog: str | None = None,
+        schema: str | None = None,
+        wait_timeout: str | None = None,
+        parameters: list[dict[str, Any]] | None = None,
+    ) -> tuple[str, dict[str, Any]]:
+        """Like _execute_statement, fixed to ARROW_STREAM -- for callers that
+        ingest the result as Arrow (see _streaming.py/cursor.py)."""
+        return await self._execute_statement(
+            statement,
+            format="ARROW_STREAM",
+            catalog=catalog,
+            schema=schema,
+            wait_timeout=wait_timeout,
+            parameters=parameters,
+        )
+
+    async def execute_json_statement(
+        self,
+        statement: str,
+        *,
+        catalog: str | None = None,
+        schema: str | None = None,
+        wait_timeout: str | None = None,
+        parameters: list[dict[str, Any]] | None = None,
+    ) -> tuple[str, dict[str, Any]]:
+        """Like _execute_statement, fixed to JSON_ARRAY -- each fetched
+        chunk's bytes (see stream_chunks_by_index) are then a plain JSON array
+        of rows, for callers that want Python values without an Arrow parse."""
+        return await self._execute_statement(
+            statement,
+            format="JSON_ARRAY",
+            catalog=catalog,
+            schema=schema,
+            wait_timeout=wait_timeout,
+            parameters=parameters,
+        )
+
+    async def upload_volume_file(self, volume_path: str, data: bytes) -> None:
+        """Uploads `data` to a Unity Catalog volume path via the Files API,
+        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,
+            )
+
+    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
+
+    async def stream_chunks_by_index(
+        self, statement_id: str, chunk_metas: list[dict[str, Any]]
+    ) -> AsyncIterator[tuple[bytes, int | None, int]]:
+        """Fetch already-known chunks for a completed statement (see
+        execute_arrow_statement/execute_json_statement), with real
+        backpressure -- see _fetch_chunks_with_backpressure. Each yielded
+        (blob, row_count, chunk_index) is an independent, self-contained
+        result blob plus its row count from the manifest (`None` if the
+        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
+
+    async def _fetch_link_bytes(self, client: httpx.AsyncClient, url: str) -> bytes:
+        async def _do() -> bytes:
+            resp = await client.get(url, timeout=self.http_timeout)
+            resp.raise_for_status()
+            return resp.content
+
+        return await _retry_call(_do)
+
+    async def _fetch_chunk_index(self, client: httpx.AsyncClient, statement_id: str, chunk_index: int) -> list[bytes]:
+        """Resolve one chunk index to its external link(s) (usually exactly
+        one) and download the bytes. Each chunk index can be requested
+        independently -- that's what makes concurrent fetching possible."""
+        resp = await self._authed_request(
+            client, "GET", f"{self._host}/api/2.0/sql/statements/{statement_id}/result/chunks/{chunk_index}"
+        )
+        links = resp.json().get("external_links") or []
+        return await asyncio.gather(*(self._fetch_link_bytes(client, link["external_link"]) for link in links))
+
+    async def _fetch_chunks_with_backpressure(
+        self,
+        client: httpx.AsyncClient,
+        statement_id: str,
+        chunk_metas: list[dict[str, Any]],
+    ) -> AsyncIterator[tuple[bytes, int | None, int]]:
+        """Fetch every chunk concurrently (bounded), instead of one round trip
+        at a time -- for a result with hundreds of chunks, that sequential
+        chain dominates total latency. Chunks complete out of order; ordering
+        (if the caller's SQL needs it) is restored downstream, not here.
+
+        This is a bounded *worker pool*, not `asyncio.as_completed` over every
+        chunk task scheduled up front -- that shape bounds how many downloads
+        run *simultaneously* but not how many completed-and-downloaded chunks
+        pile up waiting for a slow consumer, since a slot frees the instant
+        its own download finishes, whether or not anything downstream has
+        consumed that chunk yet. If the consumer is slower than the network --
+        plausible, one async task processing vs. several concurrent
+        downloads -- unconsumed bytes still pile up to O(whole result)
+        regardless of the concurrency bound. Here, a worker blocks on
+        `out.put()` (a bounded queue) until the consumer actually takes the
+        previous chunk, so a fetch slot is only freed once its chunk has been
+        handed off -- backpressure reaches all the way to the download, and
+        peak resident chunks stay at ~concurrency + queue size, not
+        O(whole result)."""
+        concurrency = self.chunk_fetch_concurrency
+        work: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        for meta in chunk_metas:
+            work.put_nowait(meta)
+        out: asyncio.Queue[tuple[bytes, int | None, int] | None] = asyncio.Queue(maxsize=concurrency)
+
+        async def worker() -> None:
+            while True:
+                try:
+                    meta = work.get_nowait()
+                except asyncio.QueueEmpty:
+                    return
+                for blob in await self._fetch_chunk_index(client, statement_id, meta["chunk_index"]):
+                    await out.put((blob, meta.get("row_count"), meta["chunk_index"]))  # blocks here = backpressure
+
+        workers = [asyncio.create_task(worker()) for _ in range(min(concurrency, len(chunk_metas) or 1))]
+
+        async def close_when_done() -> None:
+            # return_exceptions=True, not the default False: with the default,
+            # a failing worker makes gather raise *immediately* and this
+            # coroutine never reaches `out.put(None)` -- the consumer's
+            # `await out.get()` below then blocks forever. Collect every
+            # worker's outcome first, unblock the consumer unconditionally,
+            # then re-raise whichever failed.
+            results = await asyncio.gather(*workers, return_exceptions=True)
+            await out.put(None)
+            for result in results:
+                if isinstance(result, BaseException):
+                    raise result
+
+        closer = asyncio.create_task(close_when_done())
+        try:
+            while True:
+                item = await out.get()
+                if item is None:
+                    break
+                yield item
+            await closer  # propagate a worker failure that only surfaced after the last successful yield
+        finally:
+            closer.cancel()
+            for w in workers:
+                w.cancel()

arrowbricks/cursor.py

@@ -0,0 +1,261 @@
+"""A Cursor/Connection surface shaped like databricks-sql-python's (execute,
+fetchone/fetchmany/fetchall, fetchall_arrow/fetchmany_arrow) -- familiar DB-API
+ergonomics on top of the Statement Execution API + arro3, for callers who want
+to pull rows/Arrow incrementally rather than get everything back at once.
+
+Unlike a real DB-API cursor, this is async throughout (`execute`, `fetchone`,
+etc. are all coroutines) since the underlying client is -- there's no
+synchronous escape hatch, matching every other function in this package.
+
+Chunks are fetched from Databricks lazily, only as fetchone/fetchmany/fetchall
+actually need more rows -- a `fetchmany(100)` loop over a 700k-row result
+never pulls more chunks than it's consumed. Chunks can arrive out of order
+over the network (see client.py); a `pending` buffer holds early arrivals
+until the next expected chunk_index shows up, same reasoning as
+_streaming.py's stream_query_json."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+import arro3.core as core
+
+from ._streaming import (
+    HEARTBEAT,
+    ReplayableArrowChunk,
+    await_with_heartbeat,
+    fetch_arrow_chunks_with_manifest,
+    windowed_sql,
+)
+from .client import DatabricksClient
+
+__all__ = ["Connection", "Cursor", "connect"]
+
+Row = tuple[Any, ...]
+Description = tuple[str, str | None, None, None, None, None, None]
+
+
+def _empty_table(schema: core.Schema | None) -> core.Table:
+    return core.Table.from_batches([], schema=schema) if schema is not None else core.Table.from_pydict({})
+
+
+def _description_from_manifest(manifest: dict[str, Any]) -> list[Description]:
+    columns = manifest.get("schema", {}).get("columns") or []
+    return [(c.get("name"), c.get("type_name"), None, None, None, None, None) for c in columns]
+
+
+class _ResultSet:
+    """Buffers already-fetched-but-not-yet-returned rows at the Arrow level
+    (an arro3 Table), so fetchall_arrow/fetchmany_arrow stay zero-copy and
+    fetchone/fetchmany/fetchall just materialize whatever slice they need.
+    Not part of the public API -- reached only via Cursor."""
+
+    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] = {}
+        self._next_idx = 0
+        self._exhausted = False
+        self._buffer: core.Table | None = None
+        self.rownumber = 0
+
+    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."""
+        while True:
+            if self._next_idx in self._pending:
+                ready = self._pending.pop(self._next_idx)
+                self._next_idx += 1
+                return ready.to_table()
+            if self._exhausted:
+                return None
+            try:
+                chunk = await self._chunk_aiter.__anext__()
+            except StopAsyncIteration:
+                self._exhausted = True
+                continue
+            self._pending[chunk.chunk_index] = chunk
+
+    async def _ensure_buffer(self, want: int) -> None:
+        while (self._buffer is None or self._buffer.num_rows < want) and not self._exhausted:
+            table = await self._pull_one_chunk_table()
+            if table is None:
+                break
+            if self.schema is None:
+                self.schema = table.schema
+            self._buffer = (
+                table
+                if self._buffer is None
+                else core.Table.from_batches(self._buffer.to_batches() + table.to_batches(), schema=self._buffer.schema)
+            )
+
+    async def fetchmany_arrow(self, size: int) -> core.Table:
+        await self._ensure_buffer(size)
+        if self._buffer is None or self._buffer.num_rows == 0:
+            return _empty_table(self.schema)
+        n = min(size, self._buffer.num_rows)
+        out = self._buffer.slice(0, n)
+        rest = self._buffer.slice(n)
+        self._buffer = rest if rest.num_rows else None
+        self.rownumber += n
+        return out
+
+    async def fetchall_arrow(self) -> core.Table:
+        while not self._exhausted:
+            await self._ensure_buffer((self._buffer.num_rows if self._buffer is not None else 0) + 1)
+        out = self._buffer if self._buffer is not None else _empty_table(self.schema)
+        self._buffer = None
+        self.rownumber += out.num_rows
+        return out
+
+    @staticmethod
+    def _table_to_rows(table: core.Table) -> list[Row]:
+        if table.num_rows == 0:
+            return []
+        columns = [table.column(i).combine_chunks().to_pylist() for i in range(table.num_columns)]
+        return list(zip(*columns, strict=True))
+
+    async def fetchone(self) -> Row | None:
+        rows = self._table_to_rows(await self.fetchmany_arrow(1))
+        return rows[0] if rows else None
+
+    async def fetchmany(self, size: int) -> list[Row]:
+        return self._table_to_rows(await self.fetchmany_arrow(size))
+
+    async def fetchall(self) -> list[Row]:
+        return self._table_to_rows(await self.fetchall_arrow())
+
+
+class Cursor:
+    """One statement's worth of state: `execute()` (or `execute_streamed()`)
+    submits and waits for it, then fetchone/fetchmany/fetchall/fetchall_arrow/
+    fetchmany_arrow pull its result. Reusable across statements -- each
+    `execute()` call replaces the previous result set."""
+
+    def __init__(self, client: DatabricksClient) -> None:
+        self._client = client
+        self._result: _ResultSet | None = None
+        self.description: list[Description] | None = None
+
+    def execute_streamed(
+        self,
+        sql: str,
+        parameters: list[dict[str, Any]] | None = None,
+        *,
+        row_limit: int | None = None,
+        offset: int | None = None,
+        catalog: str | None = None,
+        schema: str | None = None,
+        total_timeout_s: float | None = None,
+    ) -> AsyncIterator[Any]:
+        """Like execute(), but yields HEARTBEAT while waiting on Databricks
+        instead of blocking silently -- for bridging e.g. an SSE connection
+        during a possible multi-minute cold warehouse start. Yields HEARTBEAT
+        zero or more times, then this same Cursor once ready to fetch."""
+        sql = windowed_sql(sql, row_limit=row_limit, offset=offset)
+
+        async def _gen() -> AsyncIterator[Any]:
+            fetch = fetch_arrow_chunks_with_manifest(
+                self._client, sql, catalog=catalog, schema=schema, parameters=parameters
+            )
+            async for item in await_with_heartbeat(fetch, total_timeout_s=total_timeout_s):
+                if item is HEARTBEAT:
+                    yield HEARTBEAT
+                    continue
+                _statement_id, manifest, chunk_iter = item
+                self.description = _description_from_manifest(manifest)
+                self._result = _ResultSet(schema=None, chunk_aiter=chunk_iter.__aiter__())
+                yield self
+
+        return _gen()
+
+    async def execute(
+        self,
+        sql: str,
+        parameters: list[dict[str, Any]] | None = None,
+        *,
+        row_limit: int | None = None,
+        offset: int | None = None,
+        catalog: str | None = None,
+        schema: str | None = None,
+        total_timeout_s: float | None = None,
+    ) -> Cursor:
+        """Submits `sql` and waits for it to complete -- a plain blocking
+        await, like a real DB-API cursor's execute(). `parameters`, if given,
+        is Databricks' own named-parameter format --
+        [{"name": ..., "value": ..., "type": ...}] bound against `:name`
+        markers in `sql`, not DB-API's `?`/`%s` placeholder style."""
+        async for _ in self.execute_streamed(
+            sql,
+            parameters,
+            row_limit=row_limit,
+            offset=offset,
+            catalog=catalog,
+            schema=schema,
+            total_timeout_s=total_timeout_s,
+        ):
+            pass
+        return self
+
+    def _require_result(self) -> _ResultSet:
+        if self._result is None:
+            raise RuntimeError("no active result set -- call execute() first")
+        return self._result
+
+    async def fetchone(self) -> Row | None:
+        return await self._require_result().fetchone()
+
+    async def fetchmany(self, size: int) -> list[Row]:
+        return await self._require_result().fetchmany(size)
+
+    async def fetchall(self) -> list[Row]:
+        return await self._require_result().fetchall()
+
+    async def fetchmany_arrow(self, size: int) -> core.Table:
+        return await self._require_result().fetchmany_arrow(size)
+
+    async def fetchall_arrow(self) -> core.Table:
+        return await self._require_result().fetchall_arrow()
+
+    def __aiter__(self) -> Cursor:
+        return self
+
+    async def __anext__(self) -> Row:
+        row = await self.fetchone()
+        if row is None:
+            raise StopAsyncIteration
+        return row
+
+
+class Connection:
+    """One Databricks SQL warehouse endpoint -- see DatabricksClient for the
+    constructor args (auth, timeouts, chunk_fetch_concurrency). Statements are
+    independent REST calls with no server-side session state, so `cursor()`
+    can be called as many times as you like; `close()` is a no-op kept for
+    context-manager parity with a real DB-API connection."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._client = DatabricksClient(*args, **kwargs)
+
+    def cursor(self) -> Cursor:
+        return Cursor(self._client)
+
+    async def close(self) -> None:
+        pass
+
+    async def __aenter__(self) -> Connection:
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.close()
+
+
+def connect(*args: Any, **kwargs: Any) -> Connection:
+    """Same signature as DatabricksClient -- `connect(host, warehouse_id,
+    token=..., ...)`."""
+    return Connection(*args, **kwargs)

arrowbricks-0.1.0.dist-info/METADATA

@@ -0,0 +1,114 @@
+Metadata-Version: 2.4
+Name: arrowbricks
+Version: 0.1.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
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: arro3-core>=0.8
+Requires-Dist: arro3-io>=0.8
+Requires-Dist: httpx>=0.27
+Description-Content-Type: text/markdown
+
+# arrowbricks
+
+Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a `Cursor` shaped like [`databricks-sql-python`](https://github.com/databricks/databricks-sql-python)'s (`execute`, `fetchone`/`fetchmany`/`fetchall`, `fetchall_arrow`/`fetchmany_arrow`), or `stream_query_json` for streaming NDJSON. One Arrow engine ([arro3](https://github.com/kylebarron/arro3)), no DuckDB, no pandas/pyarrow.
+
+- Single responsibility: Databricks to Arrow via arro3. No embedded query engine -- that's [duckbricks](https://github.com/bmsuisse/duckbricks), built on top of this.
+- Bring-your-own-auth -- a static token or your own token-refresh callable. No cloud-SDK dependency baked in.
+- Result-order preserved even though chunks can complete out of order over the network.
+- Chunks are fetched lazily as `fetchone`/`fetchmany`/`fetchall` actually need them, not all upfront.
+- Heartbeats between slow chunks (`execute_streamed`/`stream_query_json`), so a caller streaming this over e.g. SSE never goes silent during a cold warehouse start.
+
+## Install
+
+```bash
+pip install arrowbricks
+```
+
+Dependencies: `httpx` + `arro3-core` + `arro3-io`. That's the whole tree.
+
+## Quickstart
+
+```python
+import asyncio
+from arrowbricks import connect
+
+
+async def main():
+    conn = connect(
+        host="adb-1234567890.1.azuredatabricks.net",
+        warehouse_id="abcd1234efgh5678",
+        token="dapi...",  # or token_provider=... -- see Auth below
+    )
+    cursor = conn.cursor()
+
+    await cursor.execute("SELECT * FROM my_catalog.my_schema.my_table LIMIT 100")
+    async for row in cursor:
+        print(row)
+
+    await cursor.execute("SELECT * FROM my_catalog.my_schema.my_table LIMIT 100")
+    table = await cursor.fetchall_arrow()  # an arro3 Table
+
+
+asyncio.run(main())
+```
+
+For streaming NDJSON (e.g. a FastAPI SSE endpoint, first row out as soon as its chunk arrives):
+
+```python
+from arrowbricks import HEARTBEAT, DatabricksClient, stream_query_json
+
+client = DatabricksClient(host=..., warehouse_id=..., token=...)
+
+async for item in stream_query_json(client, "SELECT * FROM my_catalog.my_schema.big_table"):
+    if item is HEARTBEAT:
+        continue  # forward as an SSE keep-alive comment, e.g.
+    print(item)  # one ready-to-send JSON string per row
+```
+
+## Why not `databricks-sql-connector`?
+
+The [official driver](https://github.com/databricks/databricks-sql-python) is the right choice if you need full DB-API 2.0 compatibility over Databricks' Thrift/ODBC-style protocol. If you just want a query result as Arrow/JSON in your own async app, it drags in a lot for that: `pandas`, `thrift`, `openpyxl`, `pybreaker`, `pyjwt`, `oauthlib`, `lz4`, `requests`, `urllib3` as hard dependencies. arrowbricks talks to the plain REST Statement Execution API instead, and its whole dependency tree is `httpx` + `arro3-core` + `arro3-io`. The `Cursor` API is deliberately shaped like the official driver's so switching between them is mostly a constructor change, but arrowbricks is async throughout (`execute`, `fetchone`, etc. are all coroutines) -- there's no sync escape hatch.
+
+## Why not `duckbricks`?
+
+[duckbricks](https://github.com/bmsuisse/duckbricks) does the same Databricks-to-Arrow work, then goes further: it uses a real embedded DuckDB engine to materialize results into your own DuckDB connection/table (`feed_select_to_duckdb_table`), or push a DuckDB query's result *up* to Databricks (`feed_duckdb_table_to_databricks`). If you need that -- a real local SQL engine sitting on top, not just "run this query, get Arrow/JSON back" -- use duckbricks; it depends on arrowbricks for the Databricks/Arrow half. If you don't need DuckDB at all, arrowbricks alone is the smaller, single-responsibility half.
+
+## Auth
+
+`connect`/`DatabricksClient` take either:
+
+- `token: str` -- a static personal access token or pre-issued OAuth token, or
+- `token_provider` -- a callable (sync or async) returning a token string, called on every request.
+
+arrowbricks has no opinion on *how* you get a token and no cloud-SDK dependency of its own. If your provider is expensive to call, cache/refresh inside it -- arrowbricks does no caching on your behalf.
+
+```python
+conn = connect(host=..., warehouse_id=..., token_provider=my_token_provider)
+```
+
+## API
+
+- `connect(host, warehouse_id, *, token=None, token_provider=None, ...) -> Connection`
+- `Connection.cursor() -> Cursor`
+- `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.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` 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).
+- `DatabricksClient(host, warehouse_id, *, token=None, token_provider=None, ...)` -- the lower-level client `Connection` wraps. `client.execute_json_statement(sql, ...)` for plain JSON rows with no Arrow parse at all; `client.upload_volume_file(volume_path, data)`/`client.delete_volume_file(volume_path)` for the Files API.
+- `write_ipc_stream(table_or_chunk, buf)` -- thin wrapper around `arro3.io.write_ipc_stream` that always writes uncompressed bodies (see below).
+
+`Cursor.execute`/`execute_streamed`/`stream_query_json` all accept `catalog`, `schema`, `row_limit`, `offset`, and `total_timeout_s`.
+
+## A note on Arrow IPC compression
+
+`write_ipc_stream` (and everything in this package that serializes Arrow-IPC bytes) always writes **uncompressed** bodies. arro3's own default (`compression="LZ4"`) is transparently decompressed by DuckDB's Arrow reader, but not necessarily by every other Arrow IPC reader -- notably, `duckdb-wasm`'s browser-side decoder silently fails to parse LZ4-compressed bodies. If you're producing bytes that might be consumed by something other than a Python DuckDB connection, this default matters.
+
+## License
+
+MIT

arrowbricks-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+arrowbricks/__init__.py,sha256=ztesGsBdZnd88gwezxhpay2ZpfiJIZt85BCenkySFEE,471
+arrowbricks/_streaming.py,sha256=FkOYLqLUANvYTMiNISUUkQ_qcnXlAGErd1J_ogsookM,11217
+arrowbricks/client.py,sha256=QZWEpZ_qGazFyq1nbCs5MDMnKmoCkudyQGv7nYNyqas,16950
+arrowbricks/cursor.py,sha256=ql5rpkufyHY81s0UpduLLf8S6bjogzVDxNRQ7kcrwpU,10306
+arrowbricks-0.1.0.dist-info/METADATA,sha256=O90txWepq1W__nkNN28qo4VkzrMOSQQ0F24a3573J2s,7340
+arrowbricks-0.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+arrowbricks-0.1.0.dist-info/licenses/LICENSE,sha256=OGvcpI5L_-dM9pJ30WhYStEa2cIhvRo5GZ2i5_YPl9k,1070
+arrowbricks-0.1.0.dist-info/RECORD,,

arrowbricks-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any

arrowbricks-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 BMS Suisse AG
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.