arrowbricks 3.1.4__tar.gz → 3.2.0__tar.gz

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 3.1.4
+Version: 3.2.0
 Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
 Provides-Extra: arro3
 License-File: LICENSE
@@ -160,9 +160,9 @@ conn = connect(host=..., warehouse_id=..., token=..., on_event=log_query)
 
 ## Cancellation
 
-When a chunk download times out (`total_timeout_s`) or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while it's in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
+When `total_timeout_s` elapses or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while a query is in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
 
-This covers `Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, `client.stream_query_json(...)`, and the lower-level `._core.Client`'s own streamed APIs -- everything that wraps the chunk-download phase in the Rust-level heartbeat. Neither `total_timeout_s` nor a bare `task.cancel()` on `Cursor.execute()`/`execute_streamed()` itself (the initial submit/poll wait) triggers a cancel -- by the time a query's chunks are being fetched at all, the statement has typically already finished running server-side, so there's usually nothing left to cancel there in practice.
+This covers both phases of a query: the submit/poll wait while the statement is still running (`Cursor.execute()`/`execute_streamed()`, and the start of `client.stream_query_json(...)`), and the chunk download afterwards (`Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, the rest of `stream_query_json`). `stream_query_json`'s `total_timeout_s` is one budget across both phases. On `protocol="sea"` there is one short blind spot: the submit request itself can wait server-side for up to `wait_timeout` (default 30s) before Databricks returns a statement id, and a query abandoned inside that window can't be cancelled.
 
 ## Errors
 

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/README.md

@@ -148,9 +148,9 @@ conn = connect(host=..., warehouse_id=..., token=..., on_event=log_query)
 
 ## Cancellation
 
-When a chunk download times out (`total_timeout_s`) or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while it's in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
+When `total_timeout_s` elapses or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while a query is in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
 
-This covers `Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, `client.stream_query_json(...)`, and the lower-level `._core.Client`'s own streamed APIs -- everything that wraps the chunk-download phase in the Rust-level heartbeat. Neither `total_timeout_s` nor a bare `task.cancel()` on `Cursor.execute()`/`execute_streamed()` itself (the initial submit/poll wait) triggers a cancel -- by the time a query's chunks are being fetched at all, the statement has typically already finished running server-side, so there's usually nothing left to cancel there in practice.
+This covers both phases of a query: the submit/poll wait while the statement is still running (`Cursor.execute()`/`execute_streamed()`, and the start of `client.stream_query_json(...)`), and the chunk download afterwards (`Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, the rest of `stream_query_json`). `stream_query_json`'s `total_timeout_s` is one budget across both phases. On `protocol="sea"` there is one short blind spot: the submit request itself can wait server-side for up to `wait_timeout` (default 30s) before Databricks returns a statement id, and a query abandoned inside that window can't be cancelled.
 
 ## Errors
 

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "3.1.4"
+version = "3.2.0"
 description = "Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a DB-API-ish Cursor (fetchone/fetchmany/fetchall/fetchall_arrow) or NDJSON streaming. Rust/PyO3 core throughout -- zero required runtime dependencies."
 readme = "README.md"
 license = "MIT"

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/Cargo.lock

@@ -178,7 +178,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "3.1.4"
+version = "3.2.0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/Cargo.toml

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

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/client/model.rs

@@ -6,6 +6,7 @@
 //! statement gets you before any chunk is fetched (`ChunkMeta`,
 //! `StatementSubmitResult`, `InlineOrExternal`, `ChunkItem`).
 
+use std::sync::Mutex;
 use std::sync::atomic::{AtomicU8, AtomicU32, AtomicU64, AtomicUsize, Ordering};
 
 use bytes::Bytes;
@@ -55,7 +56,7 @@ pub struct ColumnDescription {
 /// `Drop for HeartbeatWait`/`Drop for HeartbeatStream`) can fire
 /// `DbClient::cancel_statement` without either of those generic structs
 /// needing to know anything protocol-specific themselves.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum CancelHandle {
     Sea { statement_id: String },
     Thrift { operation: thrift::OperationHandle },
@@ -113,6 +114,17 @@ pub struct QueryStatsAccumulator {
     /// anything itself.
     warehouse_wait_bits: AtomicU64,
     outcome: AtomicU8,
+    /// The statement/operation handle while it's submitted but not yet
+    /// terminal -- set by the submit/poll loops (`client/sea.rs`'s
+    /// `submit_and_poll_inner`, `pipeline/thrift_exec.rs`'s
+    /// `submit_and_await_thrift_statement`) as soon as Databricks hands one
+    /// back, cleared once the statement reaches a terminal state. Whatever
+    /// is still here when `pipeline/stats.rs`'s `CancelInFlightOnDrop`
+    /// drops (the submit future was abandoned -- a `total_timeout_s`, a
+    /// Python-side `task.cancel()`/`asyncio.wait_for` -- or a poll failed
+    /// mid-wait) gets a best-effort server-side cancel, so a query nobody
+    /// is waiting for anymore stops running on the warehouse.
+    in_flight: Mutex<Option<CancelHandle>>,
 }
 
 const OUTCOME_UNSET: u8 = 0;
@@ -170,6 +182,18 @@ impl QueryStatsAccumulator {
     pub fn warehouse_wait_s(&self) -> f64 {
         f64::from_bits(self.warehouse_wait_bits.load(Ordering::Relaxed))
     }
+
+    pub fn set_in_flight(&self, handle: CancelHandle) {
+        *self.in_flight.lock().unwrap() = Some(handle);
+    }
+
+    pub fn clear_in_flight(&self) {
+        *self.in_flight.lock().unwrap() = None;
+    }
+
+    pub fn take_in_flight(&self) -> Option<CancelHandle> {
+        self.in_flight.lock().unwrap().take()
+    }
 }
 
 /// One query's worth of timing/counters, handed to `EventSink::on_event`

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/client/sea.rs

@@ -21,7 +21,8 @@ use super::DbClient;
 use super::POLL_INTERVAL;
 use super::error::{ApiError, join_error};
 use super::model::{
-    ChunkItem, ChunkMeta, ColumnDescription, InlineOrExternal, QueryStatsAccumulator, StatementSubmitResult,
+    CancelHandle, ChunkItem, ChunkMeta, ColumnDescription, InlineOrExternal, QueryStatsAccumulator,
+    StatementSubmitResult,
 };
 
 /// Typed response shapes -- replaces navigating a dynamic `serde_json::Value`
@@ -373,10 +374,14 @@ impl DbClient {
             }
         }
 
+        let mut checkin = SessionCheckin {
+            client: self,
+            catalog,
+            schema,
+            session_id,
+        };
         let result = self.submit_and_poll_inner(body, stats).await;
-        if let Some(id) = session_id {
-            self.checkin_session(catalog, schema, id, result.is_ok());
-        }
+        checkin.finish(result.is_ok());
         result
     }
 
@@ -390,6 +395,9 @@ impl DbClient {
             .authed_json(reqwest::Method::POST, &url, Some(&body), Some(stats))
             .await?;
 
+        stats.set_in_flight(CancelHandle::Sea {
+            statement_id: data.statement_id.clone(),
+        });
         while !matches!(
             data.status.state.as_str(),
             "SUCCEEDED" | "FAILED" | "CANCELED" | "CLOSED"
@@ -400,6 +408,7 @@ impl DbClient {
                 .authed_json(reqwest::Method::GET, &poll_url, None, Some(stats))
                 .await?;
         }
+        stats.clear_in_flight();
 
         match data.status.state.as_str() {
             "FAILED" => {
@@ -599,6 +608,31 @@ impl DbClient {
 /// panic case would let that worker's unfetched work vanish with no error at
 /// all: the channel closing normally looks to the consumer exactly like a
 /// complete, successful result instead of a truncated one.
+/// Returns `submit_and_poll`'s session to the pool even when its future is
+/// dropped mid-poll (a timeout or cancellation) -- without it the pool's
+/// reservation for that key leaks, and after `MAX_SESSIONS_PER_KEY` such
+/// drops every later query for the key runs session-less.
+struct SessionCheckin<'a> {
+    client: &'a DbClient,
+    catalog: Option<&'a str>,
+    schema: Option<&'a str>,
+    session_id: Option<String>,
+}
+
+impl SessionCheckin<'_> {
+    fn finish(&mut self, keep: bool) {
+        if let Some(id) = self.session_id.take() {
+            self.client.checkin_session(self.catalog, self.schema, id, keep);
+        }
+    }
+}
+
+impl Drop for SessionCheckin<'_> {
+    fn drop(&mut self) {
+        self.finish(false);
+    }
+}
+
 async fn join_first_error(handles: Vec<tokio::task::JoinHandle<Result<(), ApiError>>>) -> Option<ApiError> {
     let mut first_err = None;
     for h in handles {

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/heartbeat.rs

@@ -175,6 +175,10 @@ impl<T: Send + 'static> HeartbeatWait<T> {
         self
     }
 
+    pub fn deadline(&self) -> Option<Instant> {
+        self.deadline
+    }
+
     /// One step: `Ok(Some(Tick::Heartbeat))` if still waiting,
     /// `Ok(Some(Tick::Ready(value)))` exactly once when the wrapped future
     /// completes, `Ok(None)` if already exhausted (caller should raise
@@ -254,6 +258,14 @@ impl<T: Send + 'static> HeartbeatStream<T> {
         }
     }
 
+    /// Shares an existing deadline (e.g. `HeartbeatWait::deadline` of the
+    /// submit/poll wait that preceded this stream) instead of starting a
+    /// fresh `total_timeout_s` budget at construction.
+    pub fn with_deadline(mut self, deadline: Option<Instant>) -> Self {
+        self.deadline = deadline;
+        self
+    }
+
     /// See `HeartbeatWait::with_cancel`'s own doc comment -- identical
     /// contract.
     pub fn with_cancel(mut self, on_cancel: impl FnOnce(bool) + Send + 'static) -> Self {

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/lib.rs

@@ -911,13 +911,10 @@ impl PyFetchallArrowStreamedIter {
     }
 }
 
-/// Not-yet-started vs. running state for `PyNdjsonStreamIter`. The
-/// submit/poll/spawn-workers step (`Pending` -> `Running`) happens on the
-/// iterator's first `__anext__` call, un-heartbeated -- matching
-/// `stream_query_json`'s pre-cutover behavior (its submit/poll wait was never
-/// heartbeat-wrapped, only its chunk loop was). The `total_timeout_s` budget
-/// starts counting from `Running`, not from construction, for the same
-/// reason.
+/// State for `PyNdjsonStreamIter`. The iterator's first `__anext__` starts
+/// the submit/poll/spawn-workers step (`Pending` -> `Submitting`), which is
+/// heartbeated and timed like the chunk pulls in `Running`; one
+/// `total_timeout_s` budget, counted from that first call, covers both.
 enum PyNdjsonStreamState {
     Pending {
         client: Arc<DbClient>,
@@ -928,6 +925,15 @@ enum PyNdjsonStreamState {
         total_timeout_s: Option<f64>,
         non_finite_as_string: bool,
     },
+    /// Submit/poll in flight under the same `total_timeout_s` budget (and
+    /// heartbeats) as the chunk pulls that follow -- a timeout here drops
+    /// the submit future, which cancels the statement server-side (see
+    /// `pipeline/stats.rs`'s `CancelInFlightOnDrop`).
+    Submitting {
+        client: Arc<DbClient>,
+        wait: HeartbeatWait<NdjsonStream>,
+        total_timeout_s: Option<f64>,
+    },
     Running {
         stream: Arc<AsyncMutex<NdjsonStream>>,
         heartbeat: HeartbeatStream<Vec<String>>,
@@ -970,22 +976,54 @@ impl PyNdjsonStreamIter {
                         else {
                             unreachable!()
                         };
-                        let client_for_cancel = client.clone();
-                        let stream = pipeline::execute_ndjson_stream(
+                        let submit_client = client.clone();
+                        let wait = HeartbeatWait::new(
+                            async move {
+                                pipeline::execute_ndjson_stream(
+                                    submit_client,
+                                    &statement,
+                                    catalog.as_deref(),
+                                    schema.as_deref(),
+                                    parameters,
+                                    non_finite_as_string,
+                                )
+                                .await
+                            },
+                            total_timeout_s,
+                        );
+                        *guard = PyNdjsonStreamState::Submitting {
                             client,
-                            &statement,
-                            catalog.as_deref(),
-                            schema.as_deref(),
-                            parameters,
-                            non_finite_as_string,
-                        )
-                        .await
-                        .map_err(api_error_to_pyerr)?;
-                        let heartbeat = HeartbeatStream::new(total_timeout_s).with_cancel(pipeline::cancel_hook(
-                            client_for_cancel,
-                            stream.cancel_handle.clone(),
-                            stream.stats.clone(),
-                        ));
+                            wait,
+                            total_timeout_s,
+                        };
+                    }
+                    PyNdjsonStreamState::Submitting { wait, .. } => {
+                        let stream = match wait.tick().await {
+                            Ok(Some(Tick::Heartbeat)) => {
+                                return Python::attach(|py| heartbeat_singleton(py).map(|h| h.into_any()));
+                            }
+                            Ok(Some(Tick::Ready(stream))) => stream,
+                            Ok(None) => unreachable!("Submitting is replaced as soon as its wait yields Ready"),
+                            Err(e) => {
+                                *guard = PyNdjsonStreamState::Done;
+                                return Err(api_error_to_pyerr(e));
+                            }
+                        };
+                        let PyNdjsonStreamState::Submitting {
+                            client,
+                            wait,
+                            total_timeout_s,
+                        } = std::mem::replace(&mut *guard, PyNdjsonStreamState::Done)
+                        else {
+                            unreachable!()
+                        };
+                        let heartbeat = HeartbeatStream::new(total_timeout_s)
+                            .with_deadline(wait.deadline())
+                            .with_cancel(pipeline::cancel_hook(
+                                client,
+                                stream.cancel_handle.clone(),
+                                stream.stats.clone(),
+                            ));
                         *guard = PyNdjsonStreamState::Running {
                             stream: Arc::new(AsyncMutex::new(stream)),
                             heartbeat,

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/pipeline/ndjson.rs

@@ -300,11 +300,10 @@ impl NdjsonStream {
 }
 
 /// Submit -> poll -> start background chunk fetching for the chunk-at-a-time
-/// stream above. Matches Python's `fetch_arrow_chunks_with_manifest`: this
-/// await itself is never heartbeat-wrapped (only the per-chunk pulls that
-/// follow are) -- `stream_query_json` only wraps its chunk iterator, not
-/// this initial submit/poll wait, so this crate preserves that same gap
-/// rather than "fixing" it during the port.
+/// stream above. `lib.rs`'s `PyNdjsonStreamIter` runs this under a
+/// `HeartbeatWait` sharing the stream's `total_timeout_s` budget, so the
+/// submit/poll wait gets heartbeats and a timeout too; abandoning it
+/// cancels the statement server-side (`CancelInFlightOnDrop`).
 ///
 /// **Branches on `client.protocol`, same as `PyDbClient::execute`'s own
 /// dispatch (`lib.rs`) -- found missing entirely during a later review

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/pipeline/sea.rs

@@ -22,7 +22,7 @@ use crate::client::{
 };
 
 use super::reorder::{ReorderBuffer, decode_chunk_item};
-use super::stats::{PoisonOnDrop, StatsReporter, report_submit_error};
+use super::stats::{CancelInFlightOnDrop, PoisonOnDrop, StatsReporter, report_submit_error};
 
 pub struct ExecuteResult {
     pub statement_id: String,
@@ -217,6 +217,7 @@ pub(crate) async fn submit_sea_and_report(
     stats: &QueryStatsAccumulator,
 ) -> Result<(StatementSubmitResult, f64), ApiError> {
     let submit_t0 = Instant::now();
+    let _cancel_guard = CancelInFlightOnDrop { client, stats };
     match client
         .execute_arrow_statement(statement, catalog, schema, parameters, stats)
         .await
@@ -326,6 +327,10 @@ pub async fn execute_lazy_prefer_inline(
     // submission rather than starting fresh -- see that fallback's own doc
     // comment.
     let submit_t0 = Instant::now();
+    let cancel_guard = CancelInFlightOnDrop {
+        client: &client,
+        stats: &stats,
+    };
     let outcome = match client
         .execute_arrow_statement_prefer_inline(statement, catalog, schema, parameters.clone(), &stats)
         .await
@@ -336,6 +341,8 @@ pub async fn execute_lazy_prefer_inline(
             return Err(e);
         }
     };
+    // The statement is terminal by now; released so `client` can move below.
+    drop(cancel_guard);
     let submit_to_ready_s = submit_t0.elapsed().as_secs_f64();
 
     let (statement_id, rows, columns) = match outcome {

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/pipeline/stats.rs

@@ -45,6 +45,30 @@ pub fn cancel_hook(
     }
 }
 
+/// Held across a statement's submit/poll wait: if it drops while the
+/// statement is still in flight (see `QueryStatsAccumulator::in_flight`),
+/// fires the same fire-and-forget server-side cancel `cancel_hook` does.
+/// `heartbeat.rs`'s hooks only exist once a `ResultStream`/`NdjsonStream`
+/// has been built, i.e. after the statement is already terminal -- without
+/// this, a timeout or cancellation during the wait itself (the phase a
+/// long-running query actually spends its time in) released the caller but
+/// left the statement running on the warehouse for nobody.
+pub(crate) struct CancelInFlightOnDrop<'a> {
+    pub(crate) client: &'a Arc<DbClient>,
+    pub(crate) stats: &'a QueryStatsAccumulator,
+}
+
+impl Drop for CancelInFlightOnDrop<'_> {
+    fn drop(&mut self) {
+        if let Some(handle) = self.stats.take_in_flight() {
+            let client = self.client.clone();
+            pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+                client.cancel_statement(&handle).await;
+            });
+        }
+    }
+}
+
 /// Per-`ResultStream`/`NdjsonStream` bookkeeping needed to build and
 /// dispatch exactly one `QueryStatsData` per query, at completion -- see
 /// `client::QueryStatsAccumulator`/`EventSink` for the counters/dispatch

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/pipeline/thrift_exec.rs

@@ -26,7 +26,7 @@ use crate::thrift;
 
 use super::reorder::ReorderBuffer;
 use super::sea::ResultStream;
-use super::stats::{StatsReporter, report_submit_error};
+use super::stats::{CancelInFlightOnDrop, StatsReporter, report_submit_error};
 
 /// Encodes 16 raw bytes (a THandleIdentifier's guid) as lowercase hex --
 /// used only to give a Thrift result stream a human-readable `statement_id`
@@ -55,6 +55,34 @@ struct ThriftStatementReady {
     already_closed: bool,
 }
 
+/// Releases and closes `submit_thrift_and_start_fetch`'s session if its
+/// future is dropped during submit/poll (a timeout or cancellation), which
+/// otherwise skipped both: the pool reservation leaked and the session
+/// stayed open server-side until its TTL.
+struct ThriftSessionOnDrop<'a> {
+    client: &'a Arc<DbClient>,
+    catalog: Option<&'a str>,
+    schema: Option<&'a str>,
+    session: Option<thrift::SessionHandle>,
+    from_pool: bool,
+}
+
+impl Drop for ThriftSessionOnDrop<'_> {
+    fn drop(&mut self) {
+        let Some(session) = self.session.take() else {
+            return;
+        };
+        if self.from_pool {
+            self.client
+                .thrift_checkin_session(self.catalog, self.schema, session.clone(), false);
+        }
+        let client = self.client.clone();
+        pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+            client.thrift_close_session_raw(&session).await;
+        });
+    }
+}
+
 async fn submit_and_await_thrift_statement(
     client: &Arc<DbClient>,
     session: &thrift::SessionHandle,
@@ -108,12 +136,18 @@ async fn submit_and_await_thrift_statement(
     }
 
     if !already_finished {
+        let _cancel_guard = CancelInFlightOnDrop { client, stats };
+        stats.set_in_flight(CancelHandle::Thrift {
+            operation: operation.clone(),
+        });
         loop {
             let status = client.thrift_get_operation_status_raw(&operation, stats).await?;
             if let Some(e) = status.terminal_error() {
+                stats.clear_in_flight();
                 return Err(ApiError::statement_failed(format!("Thrift statement failed: {e}")));
             }
             if status.is_finished() {
+                stats.clear_in_flight();
                 break;
             }
             tokio::time::sleep(crate::client::THRIFT_POLL_INTERVAL).await;
@@ -306,7 +340,16 @@ pub(crate) async fn submit_thrift_and_start_fetch(
         },
     };
 
+    let mut session_guard = ThriftSessionOnDrop {
+        client: &client,
+        catalog,
+        schema,
+        session: Some(session.clone()),
+        from_pool,
+    };
     let ready = submit_and_await_thrift_statement(&client, &session, statement, parameters.as_ref(), &stats).await;
+    session_guard.session = None;
+    drop(session_guard);
     let submit_to_ready_s = submit_t0.elapsed().as_secs_f64();
 
     // Exactly one of these two arms ever touches `session` -- a pooled

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/tests/wiremock_pipeline.rs

@@ -1615,3 +1615,97 @@ async fn sea_dropping_the_heartbeat_wait_mid_fetch_fires_cancel_statement() {
         "a bare drop (not through tick()'s own timeout branch) must record outcome=cancelled, not timeout"
     );
 }
+
+async fn mount_running_warehouse_and_statement(server: &MockServer, state: &str) {
+    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;
+    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": state, "error": {"error_code": "BAD_SQL", "message": "nope"}},
+        })))
+        .mount(server)
+        .await;
+}
+
+/// Abandoning the submit/poll wait itself (a `total_timeout_s` or a Python
+/// `task.cancel()` drops this future mid-poll) must cancel the statement
+/// server-side -- `heartbeat.rs`'s hooks don't exist yet at this point.
+#[tokio::test]
+async fn sea_abandoning_the_submit_poll_wait_fires_cancel_statement() {
+    let server = MockServer::start().await;
+    mount_running_warehouse_and_statement(&server, "PENDING").await;
+    let cancel_calls = mount_cancel_statement_ok(&server).await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token").with_protocol(Protocol::Sea));
+    let result = tokio::time::timeout(
+        std::time::Duration::from_millis(300),
+        execute_lazy(client, "SELECT * FROM t", None, None, None),
+    )
+    .await;
+    assert!(
+        result.is_err(),
+        "the statement never leaves PENDING, so the wait must time out"
+    );
+
+    wait_for_calls(&cancel_calls, 1).await;
+    assert_eq!(cancel_calls.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn sea_terminal_statement_does_not_fire_cancel_statement() {
+    let server = MockServer::start().await;
+    mount_running_warehouse_and_statement(&server, "FAILED").await;
+    let cancel_calls = mount_cancel_statement_ok(&server).await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token").with_protocol(Protocol::Sea));
+    let result = execute_lazy(client, "SELECT * FROM t", None, None, None).await;
+    assert!(result.is_err());
+
+    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+    assert_eq!(
+        cancel_calls.load(Ordering::SeqCst),
+        0,
+        "an already-terminal statement has nothing left to cancel"
+    );
+}
+
+/// A submit future dropped mid-poll must still release its pooled session's
+/// reservation -- leaking one per abandonment would, after
+/// `MAX_SESSIONS_PER_KEY` of them, leave every later query session-less.
+#[tokio::test]
+async fn abandoning_the_submit_poll_wait_releases_the_session_reservation() {
+    let server = MockServer::start().await;
+    mount_running_warehouse_and_statement(&server, "PENDING").await;
+    mount_cancel_statement_ok(&server).await;
+    let session_calls = Arc::new(AtomicUsize::new(0));
+    let session_calls_for_mock = session_calls.clone();
+    Mock::given(method("POST"))
+        .and(path("/api/2.0/sql/sessions"))
+        .respond_with(move |_req: &wiremock::Request| {
+            let n = session_calls_for_mock.fetch_add(1, Ordering::SeqCst);
+            ResponseTemplate::new(200).set_body_json(json!({"session_id": format!("sess-{n}")}))
+        })
+        .mount(&server)
+        .await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token").with_protocol(Protocol::Sea));
+    for _ in 0..=MAX_SESSIONS_PER_KEY {
+        let result = tokio::time::timeout(
+            std::time::Duration::from_millis(100),
+            execute_lazy(client.clone(), "SELECT * FROM t", Some("cat1"), None, None),
+        )
+        .await;
+        assert!(result.is_err());
+    }
+
+    assert_eq!(
+        session_calls.load(Ordering::SeqCst),
+        MAX_SESSIONS_PER_KEY + 1,
+        "every abandoned attempt must release its reservation, so each one can create a session"
+    );
+}

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/rust/arrowbricks_core/tests/wiremock_thrift.rs

@@ -1826,3 +1826,125 @@ async fn thrift_dropping_the_heartbeat_wait_mid_fetch_fires_cancel_operation() {
         "a bare drop (not through tick()'s own timeout branch) must record outcome=cancelled, not timeout"
     );
 }
+
+/// Thrift counterpart of `wiremock_pipeline.rs`'s
+/// `sea_abandoning_the_submit_poll_wait_fires_cancel_statement`.
+#[tokio::test]
+async fn thrift_abandoning_the_submit_poll_wait_fires_cancel_operation() {
+    let server = MockServer::start().await;
+    mount_open_session_always(&server, b"sess").await;
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("ExecuteStatement"))
+        .respond_with(ResponseTemplate::new(200).set_body_raw(
+            build_execute_statement_resp(b"op-poll", b"opsecret-poll", None),
+            "application/x-thrift",
+        ))
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("GetOperationStatus"))
+        .respond_with(ResponseTemplate::new(200).set_body_raw(
+            build_get_operation_status_resp(operation_state::RUNNING, None),
+            "application/x-thrift",
+        ))
+        .mount(&server)
+        .await;
+    let cancel_calls = mount_cancel_operation_ok(&server).await;
+
+    let result = tokio::time::timeout(
+        std::time::Duration::from_millis(500),
+        execute_lazy_thrift(thrift_client(&server), "SELECT * FROM t", None, None, None),
+    )
+    .await;
+    assert!(
+        result.is_err(),
+        "the operation never finishes, so the wait must time out"
+    );
+
+    wait_for_calls(&cancel_calls, 1).await;
+    assert_eq!(cancel_calls.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn thrift_polled_terminal_error_does_not_fire_cancel_operation() {
+    let server = MockServer::start().await;
+    mount_open_session_always(&server, b"sess").await;
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("ExecuteStatement"))
+        .respond_with(ResponseTemplate::new(200).set_body_raw(
+            build_execute_statement_resp(b"op-err", b"opsecret-err", None),
+            "application/x-thrift",
+        ))
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("GetOperationStatus"))
+        .respond_with(ResponseTemplate::new(200).set_body_raw(
+            build_get_operation_status_resp(operation_state::ERROR, Some("boom")),
+            "application/x-thrift",
+        ))
+        .mount(&server)
+        .await;
+    mount_close_operation_ok(&server).await;
+    let cancel_calls = mount_cancel_operation_ok(&server).await;
+
+    let result = execute_lazy_thrift(thrift_client(&server), "SELECT * FROM t", None, None, None).await;
+    assert!(result.is_err());
+
+    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+    assert_eq!(
+        cancel_calls.load(Ordering::SeqCst),
+        0,
+        "an already-terminal operation has nothing left to cancel"
+    );
+}
+
+#[tokio::test]
+async fn thrift_abandoning_the_submit_poll_wait_closes_its_session() {
+    let server = MockServer::start().await;
+    mount_open_session_always(&server, b"sess").await;
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("ExecuteStatement"))
+        .respond_with(ResponseTemplate::new(200).set_body_raw(
+            build_execute_statement_resp(b"op-sess", b"opsecret-sess", None),
+            "application/x-thrift",
+        ))
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("GetOperationStatus"))
+        .respond_with(ResponseTemplate::new(200).set_body_raw(
+            build_get_operation_status_resp(operation_state::RUNNING, None),
+            "application/x-thrift",
+        ))
+        .mount(&server)
+        .await;
+    mount_cancel_operation_ok(&server).await;
+    let close_calls = Arc::new(AtomicUsize::new(0));
+    let close_calls_for_mock = close_calls.clone();
+    Mock::given(method("POST"))
+        .and(path(thrift_path()))
+        .and(IsThriftRpc("CloseSession"))
+        .respond_with(move |_req: &Request| {
+            close_calls_for_mock.fetch_add(1, Ordering::SeqCst);
+            ResponseTemplate::new(200).set_body_raw(build_close_session_resp(), "application/x-thrift")
+        })
+        .mount(&server)
+        .await;
+
+    let result = tokio::time::timeout(
+        std::time::Duration::from_millis(500),
+        execute_lazy_thrift(thrift_client(&server), "SELECT * FROM t", None, None, None),
+    )
+    .await;
+    assert!(result.is_err());
+
+    wait_for_calls(&close_calls, 1).await;
+    assert_eq!(close_calls.load(Ordering::SeqCst), 1);
+}

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/src/arrowbricks/_streaming.py

@@ -13,7 +13,7 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
-from collections.abc import AsyncIterator, Awaitable
+from collections.abc import AsyncIterator, Awaitable, Iterator
 from typing import Any, BinaryIO, TypeVar, cast
 
 from . import _core
@@ -54,6 +54,22 @@ class QueryTimeout(_core.ArrowbricksError):
     this change keeps working unchanged either way."""
 
 
+@contextlib.contextmanager
+def rust_timeout_as_query_timeout() -> Iterator[None]:
+    """Translates the Rust-level heartbeat's `total_timeout_s` error -- a
+    plain `ArrowbricksError` with heartbeat.rs's literal `"Query exceeded
+    {secs}s timeout"` message -- into `QueryTimeout`, so callers see one
+    exception type whichever heartbeat implementation ran. Matched by a
+    prefix this crate controls end to end, so an unrelated error is never
+    misclassified."""
+    try:
+        yield
+    except _core.ArrowbricksError as exc:
+        if str(exc).startswith("Query exceeded"):
+            raise QueryTimeout(str(exc)) from exc
+        raise
+
+
 class _Heartbeat:
     __slots__ = ()
 
@@ -185,6 +201,11 @@ async def stream_query_json(
     top-level float columns are covered; one nested inside a STRUCT/ARRAY/MAP
     still comes back as `null` either way.
 
+    `total_timeout_s` bounds the whole statement -- the submit/poll wait
+    (heartbeats included) and every chunk download after it -- and raises
+    `QueryTimeout` once it elapses, firing a best-effort server-side cancel
+    so the statement doesn't keep running on the warehouse.
+
     Note this yields a whole chunk's rows at once -- Databricks' own chunk
     sizing already bounds how much that is."""
     if non_finite_floats not in ("null", "string"):
@@ -192,16 +213,17 @@ async def stream_query_json(
     sql = windowed_sql(sql, row_limit=row_limit, offset=offset)
     core_client = client._core_client  # noqa: SLF001 -- same package, see client.py
 
-    async for item in core_client.stream_ndjson_lines(
-        sql,
-        catalog=catalog,
-        schema=schema,
-        parameters=params,
-        total_timeout_s=total_timeout_s,
-        non_finite_as_string=(non_finite_floats == "string"),
-    ):
-        if item is _core.HEARTBEAT:
-            yield HEARTBEAT
-            continue
-        for line in cast("list[str]", item):
-            yield line
+    with rust_timeout_as_query_timeout():
+        async for item in core_client.stream_ndjson_lines(
+            sql,
+            catalog=catalog,
+            schema=schema,
+            parameters=params,
+            total_timeout_s=total_timeout_s,
+            non_finite_as_string=(non_finite_floats == "string"),
+        ):
+            if item is _core.HEARTBEAT:
+                yield HEARTBEAT
+                continue
+            for line in cast("list[str]", item):
+                yield line

{arrowbricks-3.1.4 → arrowbricks-3.2.0}/src/arrowbricks/cursor.py

@@ -20,7 +20,7 @@ from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
 
 from . import _core
-from ._streaming import HEARTBEAT, QueryTimeout, await_with_heartbeat, windowed_sql
+from ._streaming import HEARTBEAT, await_with_heartbeat, rust_timeout_as_query_timeout, windowed_sql
 from .client import DatabricksClient
 
 if TYPE_CHECKING:
@@ -309,7 +309,7 @@ class Cursor:
             # `execute_streamed`'s own `_gen()` pattern.
             self._require_empty_row_buffer("fetchall_arrow_streamed")
             result = self._require_result()
-            try:
+            with rust_timeout_as_query_timeout():
                 async for item in result.fetchall_arrow_streamed(total_timeout_s=total_timeout_s):
                     if item is _core.HEARTBEAT:
                         yield HEARTBEAT
@@ -317,21 +317,6 @@ class Cursor:
                         if self._schema is None:
                             self._schema = await result.schema()
                         yield item
-            except RuntimeError as exc:
-                # `heartbeat::HeartbeatWait`'s own `total_timeout_s` error
-                # (`format!("Query exceeded {secs}s timeout")`, heartbeat.rs)
-                # surfaces here as a plain `RuntimeError`, not this package's
-                # own `QueryTimeout` -- translated so callers relying on
-                # catching `QueryTimeout` (same contract `execute_streamed`/
-                # `fetchall_streamed` already promise) see the same exception
-                # type regardless of which heartbeat implementation is
-                # actually running underneath. Matched by a stable literal
-                # prefix this crate controls end to end (not string-matching
-                # someone else's error), so a real, unrelated `RuntimeError`
-                # from a genuine chunk-fetch failure is never misclassified.
-                if str(exc).startswith("Query exceeded"):
-                    raise QueryTimeout(str(exc)) from exc
-                raise
 
         return _gen()