arrowbricks 0.2.0__tar.gz → 0.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 0.2.0
+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

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

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "0.2.0"
+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.2.0 → 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.2.0 → 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.2.0 → 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.2.0 → arrowbricks-0.3.0}/uv.lock

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