arrowbricks 3.0.0__tar.gz → 3.0.1__tar.gz

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 3.0.0
+Version: 3.0.1
 Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
 Provides-Extra: arro3
 License-File: LICENSE

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "3.0.0"
+version = "3.0.1"
 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. Rust/PyO3 core throughout -- zero required runtime dependencies."
 readme = "README.md"
 license = "MIT"

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/Cargo.lock

@@ -240,7 +240,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "3.0.0"
+version = "3.0.1"
 dependencies = [
  "arrow",
  "arrow-json",

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "3.0.0"
+version = "3.0.1"
 edition = "2024"
 readme = "README.md"
 

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/src/client.rs

@@ -980,7 +980,14 @@ impl DbClient {
             .map_err(join_error)?
     }
 
-    async fn ensure_warehouse_running(&self) -> Result<(), ApiError> {
+    /// Shared by both protocols -- plain REST against `/api/2.0/sql/warehouses/{id}`,
+    /// nothing SEA- or Thrift-specific about it. SEA's own `submit_and_poll`
+    /// has always called this; the Thrift path (`pipeline::execute_lazy_thrift`)
+    /// didn't, which meant a stopped warehouse got no proactive wake on
+    /// `protocol="thrift"` (now the default) -- statement submission would
+    /// eventually surface an error instead, with no `warehouse_start_timeout`
+    /// wait for it to come up first. Fixed by calling this from both.
+    pub(crate) async fn ensure_warehouse_running(&self) -> Result<(), ApiError> {
         {
             let confirmed = *self.warehouse_confirmed_running_at.lock().unwrap();
             if let Some(at) = confirmed {
@@ -1661,12 +1668,26 @@ impl DbClient {
                         };
                         match fetched {
                             Ok(blobs) => {
+                                // `meta.row_count` is the manifest's declared count for
+                                // the whole `chunk_index`, not per-blob -- safe to use as
+                                // `decode_chunk_item`'s truncation bound (same "server can
+                                // over-deliver past its declared count" protection Thrift's
+                                // resultLinks and arrowBatches paths both already have, see
+                                // AGENTS.md) only when there's exactly one blob, where
+                                // "the whole chunk's count" and "this blob's count" are the
+                                // same number. A chunk_index resolving to more than one
+                                // blob is a real but rare/defensive-coding case (see
+                                // `ChunkMeta::pre_resolved_links`'s own doc comment) whose
+                                // true per-blob row split isn't known here -- truncating
+                                // the first blob to the *whole* chunk's count would be
+                                // wrong, so those are left untruncated rather than guessed.
+                                let truncate_to = if blobs.len() == 1 { meta.row_count } else { None };
                                 for blob in blobs {
                                     let item = ChunkItem {
                                         blob,
                                         row_count: meta.row_count,
                                         chunk_index: meta.chunk_index,
-                                        truncate_to: None,
+                                        truncate_to,
                                     };
                                     if worker_tx.send(Ok(item)).await.is_err() {
                                         return Ok(());

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/src/pipeline.rs

@@ -377,32 +377,6 @@ fn hex_encode(bytes: &[u8]) -> String {
     s
 }
 
-/// Submit -> (maybe poll) -> start background chunk fetching for the
-/// `protocol="thrift"` backend -- see `client::Protocol::Thrift`'s own doc
-/// comment for the wire-format background and why this is faster than SEA
-/// for small queries. Produces the exact same `ResultStream` shape as
-/// `execute_lazy`/`execute_lazy_prefer_inline` (same `fetchmany_arrow`/
-/// `fetchall_arrow`/`schema` contract, same `ReorderBuffer`/`decode_chunk`
-/// underneath), so `PyResultSet` and everything above it needs zero
-/// Thrift-specific handling.
-///
-/// Session handling mirrors the SEA session pool (`client::SessionPool`)
-/// exactly, with one necessary difference: Thrift's `TExecuteStatementReq`
-/// *requires* a `sessionHandle` (unlike SEA's optional `session_id`), so
-/// pool exhaustion/creation failure can't fall back to a session-less
-/// submission the way SEA does -- instead, a fresh, unpooled session is
-/// opened just for this one call and closed again immediately afterward.
-/// This is a genuinely necessary fallback, not a shortcut: without it, a
-/// caller who exhausts `MAX_SESSIONS_PER_KEY` concurrent statements would
-/// get a hard failure instead of the same "a bit more session-creation
-/// overhead, but still works" degradation SEA gets.
-///
-/// A session is only needed for the initial `ExecuteStatement` call --
-/// every subsequent RPC (`GetOperationStatus`/`FetchResults`/
-/// `CloseOperation`) addresses the operation directly by its own handle, so
-/// the session is checked back in (or closed, if it was a throwaway) right
-/// after `ExecuteStatement` returns, not held for the statement's whole
-/// lifetime the way one might expect from a "session" name.
 /// What's known once a Thrift statement has reached a terminal state and is
 /// ready to hand off to `drive_thrift_fetch_loop` -- a small bundle so
 /// `execute_lazy_thrift`'s own error paths (see its doc comment) have one
@@ -492,8 +466,31 @@ async fn submit_and_await_thrift_statement(
 }
 
 /// Submit -> (maybe poll) -> start background chunk fetching for the
-/// `protocol="thrift"` backend. See `client::Protocol::Thrift`'s own doc
-/// comment for the wire-format background.
+/// `protocol="thrift"` backend -- see `client::Protocol::Thrift`'s own doc
+/// comment for the wire-format background and why this is faster than SEA
+/// for small queries. Produces the exact same `ResultStream` shape as
+/// `execute_lazy`/`execute_lazy_prefer_inline` (same `fetchmany_arrow`/
+/// `fetchall_arrow`/`schema` contract, same `ReorderBuffer`/`decode_chunk`
+/// underneath), so `PyResultSet` and everything above it needs zero
+/// Thrift-specific handling.
+///
+/// Session handling mirrors the SEA session pool (`client::SessionPool`)
+/// exactly, with one necessary difference: Thrift's `TExecuteStatementReq`
+/// *requires* a `sessionHandle` (unlike SEA's optional `session_id`), so
+/// pool exhaustion/creation failure can't fall back to a session-less
+/// submission the way SEA does -- instead, a fresh, unpooled session is
+/// opened just for this one call and closed again immediately afterward.
+/// This is a genuinely necessary fallback, not a shortcut: without it, a
+/// caller who exhausts `MAX_SESSIONS_PER_KEY` concurrent statements would
+/// get a hard failure instead of the same "a bit more session-creation
+/// overhead, but still works" degradation SEA gets.
+///
+/// A session is only needed for the initial `ExecuteStatement` call --
+/// every subsequent RPC (`GetOperationStatus`/`FetchResults`/
+/// `CloseOperation`) addresses the operation directly by its own handle, so
+/// the session is checked back in (or closed, if it was a throwaway) right
+/// after `ExecuteStatement` returns, not held for the statement's whole
+/// lifetime the way one might expect from a "session" name.
 ///
 /// **A session checked out for this call (whether pooled or a throwaway,
 /// see `client::ThriftSessionPool`'s doc comment) must stay open until the
@@ -526,6 +523,7 @@ pub async fn execute_lazy_thrift(
     schema: Option<&str>,
     parameters: Option<Value>,
 ) -> Result<ResultStream, ApiError> {
+    client.ensure_warehouse_running().await?;
     let pooled = client.thrift_checkout_session(catalog, schema).await;
     let (session, from_pool) = match pooled {
         Some(s) => (s, true),
@@ -1390,11 +1388,25 @@ pub struct JsonResult {
 /// Casting by the manifest's column type_name, if wanted, is left to the
 /// caller -- same pass-through the Python original leaves to its own
 /// caller.
-fn decode_json_chunk(blob: &Bytes) -> Result<Vec<Vec<Option<String>>>, ApiError> {
-    serde_json::from_slice(&blob[..]).map_err(|e| ApiError {
+///
+/// `truncate_to`, if `Some(n)` and the chunk decodes to more than `n` rows,
+/// slices to exactly the first `n` -- the JSON_ARRAY counterpart to
+/// `decode_chunk_item`'s Arrow-IPC truncation (see its own doc comment and
+/// `fetch_chunks_with_backpressure`'s `truncate_to` wiring): a JSON_ARRAY
+/// chunk is delivered through the exact same cloud-fetch blob-storage
+/// mechanism as an Arrow one, just decoded differently, so the same
+/// "server can over-deliver past its declared row count" risk applies here
+/// too. No batch-boundary complexity to worry about here (unlike the Arrow
+/// case) -- rows are already a flat `Vec`, so this is a plain `truncate`.
+fn decode_json_chunk(blob: &Bytes, truncate_to: Option<i64>) -> Result<Vec<Vec<Option<String>>>, ApiError> {
+    let mut rows: Vec<Vec<Option<String>>> = serde_json::from_slice(&blob[..]).map_err(|e| ApiError {
         message: format!("bad JSON_ARRAY chunk: {e}"),
         transient: false,
-    })
+    })?;
+    if let Some(n) = truncate_to.filter(|n| *n > 0) {
+        rows.truncate(n as usize);
+    }
+    Ok(rows)
 }
 
 /// JSON_ARRAY counterpart to `run_pipeline` -- same submit/poll/fetch/
@@ -1421,7 +1433,10 @@ pub async fn run_json_pipeline(
 
     let mut decode_handles = Vec::with_capacity(num_chunks);
     while let Some(item) = reorder.next().await? {
-        decode_handles.push(tokio::task::spawn_blocking(move || decode_json_chunk(&item.blob)));
+        let truncate_to = item.truncate_to;
+        decode_handles.push(tokio::task::spawn_blocking(move || {
+            decode_json_chunk(&item.blob, truncate_to)
+        }));
     }
 
     let mut rows = Vec::new();

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/tests/wiremock_pipeline.rs

@@ -336,6 +336,67 @@ async fn lazy_pipeline_skips_the_extra_resolution_get_for_a_pre_resolved_chunk()
     assert_ids_in_order(&batches, 10);
 }
 
+/// Regression test for the same "server can over-deliver past its declared
+/// row count" behavior Thrift's resultLinks/arrowBatches paths already
+/// guard against (see AGENTS.md) -- found missing on this SEA path during a
+/// later review pass, not yet observed to actually trigger against a real
+/// warehouse (unlike the Thrift case, which was), but structurally
+/// identical and cheap to close defensively: the manifest declares 3 rows
+/// for chunk_index 0, the actual file backing it encodes 5. A single link
+/// per chunk_index is the common case where `fetch_chunks_with_backpressure`
+/// now passes the declared count through as `ChunkItem::truncate_to`.
+#[tokio::test]
+async fn lazy_pipeline_truncates_a_sea_chunk_that_overshoots_its_declared_row_count() {
+    let server = MockServer::start().await;
+
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/warehouses/{WAREHOUSE_ID}")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"state": "RUNNING"})))
+        .mount(&server)
+        .await;
+
+    let uri = server.uri();
+    Mock::given(method("POST"))
+        .and(path("/api/2.0/sql/statements"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "statement_id": STATEMENT_ID,
+            "status": {"state": "SUCCEEDED"},
+            "manifest": {
+                "chunks": [{"chunk_index": 0, "row_count": 3}],
+                "schema": {"columns": [
+                    {"name": "id", "type_name": "LONG"},
+                    {"name": "label", "type_name": "STRING"},
+                ]},
+            },
+            "result": {
+                "external_links": [{"chunk_index": 0, "external_link": format!("{uri}/_data/chunk-0")}]
+            },
+        })))
+        .mount(&server)
+        .await;
+
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-0$"))
+        .respond_with(
+            // Declares 3 rows above, but the file itself has 5 -- must be
+            // sliced down to 3, not returned in full.
+            ResponseTemplate::new(200).set_body_raw(build_chunk_bytes(0, 5), "application/vnd.apache.arrow.stream"),
+        )
+        .mount(&server)
+        .await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token").with_protocol(Protocol::Sea));
+    let mut stream = execute_lazy(client, "SELECT * FROM t", None, None, None).await.unwrap();
+    let (batches, _schema) = stream.fetchall_arrow().await.unwrap();
+
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(
+        total_rows, 3,
+        "must keep exactly the manifest's declared row count, not however many rows the file actually encoded"
+    );
+    assert_ids_in_order(&batches, 3);
+}
+
 /// Regression test for a bug caught in code review before it shipped:
 /// `pre_resolved` was originally a plain `HashMap<i64, String>`, which keeps
 /// only the *last* entry when `result.external_links` has more than one for
@@ -739,6 +800,70 @@ async fn json_pipeline_happy_path() {
     assert_eq!(ids, (0..15).collect::<Vec<_>>());
 }
 
+/// Regression test for the same "server can over-deliver past its declared
+/// row count" behavior as the Arrow-format entries in this file and in
+/// AGENTS.md -- a JSON_ARRAY chunk is delivered through the exact same
+/// cloud-fetch blob-storage mechanism as an Arrow one, just decoded
+/// differently, so the same risk applies here. The manifest declares 3 rows
+/// for chunk_index 0, the actual file backing it has 5.
+#[tokio::test]
+async fn json_pipeline_truncates_a_chunk_that_overshoots_its_declared_row_count() {
+    let server = MockServer::start().await;
+
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/warehouses/{WAREHOUSE_ID}")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"state": "RUNNING"})))
+        .mount(&server)
+        .await;
+
+    let uri = server.uri();
+    Mock::given(method("POST"))
+        .and(path("/api/2.0/sql/statements"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "statement_id": STATEMENT_ID,
+            "status": {"state": "SUCCEEDED"},
+            "manifest": {
+                "chunks": [{"chunk_index": 0, "row_count": 3}],
+                "schema": {"columns": [
+                    {"name": "id", "type_name": "LONG"},
+                    {"name": "label", "type_name": "STRING"},
+                ]},
+            },
+            "result": {
+                "external_links": [{"chunk_index": 0, "external_link": format!("{uri}/_data/chunk-0")}]
+            },
+        })))
+        .mount(&server)
+        .await;
+
+    // Declares 3 rows above, but the file itself has 5.
+    let rows: Vec<_> = (0..5)
+        .map(|id: i64| json!([id.to_string(), format!("row_{id}")]))
+        .collect();
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-0$"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(rows))
+        .mount(&server)
+        .await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token").with_protocol(Protocol::Sea));
+    let result = run_json_pipeline(client, "SELECT * FROM t", None, None, None)
+        .await
+        .unwrap();
+
+    assert_eq!(
+        result.rows.len(),
+        3,
+        "must keep exactly the manifest's declared row count, not however many rows the file actually encoded"
+    );
+    let ids: Vec<i64> = result
+        .rows
+        .iter()
+        .map(|r| r[0].as_ref().unwrap().parse().unwrap())
+        .collect();
+    assert_eq!(ids, vec![0, 1, 2]);
+}
+
 #[tokio::test]
 async fn json_pipeline_survives_reverse_arrival() {
     let server = MockServer::start().await;

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/tests/wiremock_thrift.rs

@@ -361,7 +361,22 @@ fn thrift_client(server: &MockServer) -> Arc<DbClient> {
     Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token").with_protocol(Protocol::Thrift))
 }
 
+/// Every `execute_lazy_thrift` call opens with `ensure_warehouse_running`
+/// (a plain REST GET, shared with SEA -- see `DbClient::ensure_warehouse_running`'s
+/// own doc comment for why the Thrift path didn't have this until now)
+/// before it ever touches a session, so every test that submits a Thrift
+/// statement needs this mocked -- bundled into the near-universal
+/// "set up a session" helper below rather than repeated at every call site.
+async fn mount_warehouse_running(server: &MockServer) {
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/warehouses/{WAREHOUSE_ID}")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"state": "RUNNING"})))
+        .mount(server)
+        .await;
+}
+
 async fn mount_open_session_always(server: &MockServer, guid: &'static [u8]) -> Arc<AtomicUsize> {
+    mount_warehouse_running(server).await;
     let calls = Arc::new(AtomicUsize::new(0));
     let calls_for_mock = calls.clone();
     Mock::given(method("POST"))

{arrowbricks-3.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/tests_py/thrift_mock.py

@@ -35,6 +35,7 @@ from __future__ import annotations
 
 import http.server
 import io
+import json
 import re
 import threading
 from collections.abc import Callable
@@ -355,6 +356,19 @@ class ThriftMockServer:
                 self.wfile.write(resp_bytes)
 
             def do_GET(self) -> None:  # noqa: N802
+                # Every execute_lazy_thrift call opens with ensure_warehouse_running
+                # (a plain REST GET, shared with SEA -- see DbClient::ensure_warehouse_running's
+                # own doc comment) before it ever touches a session -- built in here,
+                # not test-registered, since every test that submits a Thrift
+                # statement needs it, same as tests/conftest.py's SEA mock_warehouse.
+                if self.path == f"/api/2.0/sql/warehouses/{outer.warehouse_id}":
+                    body = json.dumps({"state": "RUNNING"}).encode()
+                    self.send_response(200)
+                    self.send_header("Content-Type", "application/json")
+                    self.send_header("Content-Length", str(len(body)))
+                    self.end_headers()
+                    self.wfile.write(body)
+                    return
                 for pattern, responder in outer._get_routes:
                     m = pattern.match(self.path)
                     if m: