arrowbricks 3.0.1__tar.gz → 3.0.2__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "3.0.1"
+version = "3.0.2"
 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.1 → arrowbricks-3.0.2}/rust/arrowbricks_core/Cargo.lock

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

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

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "3.0.1"
+version = "3.0.2"
 edition = "2024"
 readme = "README.md"
 
@@ -23,8 +23,13 @@ chrono = { version = "0.4", default-features = false, features = ["std"] }
 # Decompresses cloud-fetch chunk bytes when the server honors our
 # `result_compression: LZ4_FRAME` request (see client.rs's execute_statement)
 # -- the same LZ4 Frame format databricks-sql-python's `lz4.frame` decodes.
-# Pure Rust (no C toolchain dependency), matching this crate's reasoning for
-# avoiding aws-lc-sys-style native-build pain elsewhere.
+# Pure Rust, no C toolchain dependency -- unlike reqwest's `rustls` feature
+# below, which *does* pull one in (`__rustls-aws-lc-rs` -> aws-lc-rs ->
+# aws-lc-sys, confirmed via `cargo tree -i aws-lc-sys`; this crate has no
+# `ring` feature available in reqwest 0.13 to avoid it with). A previous
+# version of this comment claimed the opposite -- that this crate avoids
+# aws-lc-sys-style native-build pain elsewhere too -- found wrong on review;
+# see AGENTS.md's build-prerequisites note.
 lz4_flex = { version = "0.11", default-features = false, features = ["frame", "std"] }
 pyo3 = { version = "0.29.1", features = ["abi3-py311"] }
 # pyo3-arrow's only optional feature is buffer_protocol (on by default),
@@ -53,6 +58,13 @@ extension-module = ["pyo3/extension-module"]
 tokio = { version = "1.53.1", features = ["rt-multi-thread", "macros", "test-util", "net", "io-util"] }
 wiremock = "0.6.5"
 
+# Currently clean on all three by luck, not enforcement -- pinning them here
+# so a future PR that introduces one fails CI instead of silently landing.
+[lints.clippy]
+redundant_clone = "warn"
+large_enum_variant = "warn"
+needless_collect = "warn"
+
 [profile.release]
 # opt-level stays at the implicit default (3, speed) rather than "s"/"z" --
 # perf was the original ask and the benchmarks were run against opt-level 3;

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

@@ -148,8 +148,7 @@ struct ResultBody {
     external_links: Vec<ResultLinkBody>,
     /// Only present for `disposition: INLINE` + `format: JSON_ARRAY` -- each
     /// row a `Vec<Option<String>>` (every non-null value a string,
-    /// Databricks' own JSON_ARRAY contract, same as `execute_json_statement`'s
-    /// normal `EXTERNAL_LINKS`+`JSON_ARRAY` chunks). Consumed by
+    /// Databricks' own JSON_ARRAY contract). Consumed by
     /// `json_convert::json_array_to_record_batch` on the `prefer_inline`
     /// fast path.
     #[serde(default)]
@@ -329,7 +328,7 @@ pub struct StatementSubmitResult {
 
 /// What `execute_arrow_statement_prefer_inline` gets you -- either the whole
 /// result inline as raw JSON_ARRAY rows (every non-null value a string,
-/// same contract as `execute_json_statement`; converting these into a real
+/// Databricks' own JSON_ARRAY contract; converting these into a real
 /// `RecordBatch` is `pipeline.rs`'s job via `json_convert`, not this
 /// Arrow-agnostic module's -- see this file's own module doc comment), or a
 /// normal `StatementSubmitResult` to fetch chunks for exactly as if
@@ -353,15 +352,18 @@ pub struct ChunkItem {
     pub blob: Bytes,
     pub row_count: Option<i64>,
     pub chunk_index: i64,
-    /// Set only by the Thrift cloud-fetch path (`pipeline.rs`'s
-    /// `fetch_thrift_link`) -- a declared row-count bound this chunk's own
-    /// decoded batches must be sliced down to if they exceed it, since a
-    /// Thrift `resultLinks` file (like SEA's own cloud-fetch files) can
+    /// A declared row-count bound this chunk's own decoded batches must be
+    /// sliced down to if they exceed it, since a cloud-fetch file can
     /// legitimately contain more rows than its own declared count for a
     /// `LIMIT`-bounded query (see `pipeline.rs`'s `decode_chunk_item`).
-    /// `None` for every other producer (SEA's own chunk fetch, Thrift's
-    /// inline `arrowBatches`) -- their row counts are never overshot the
-    /// same way, so there's nothing to slice.
+    /// Every real producer sets this now -- Thrift `resultLinks`
+    /// (`pipeline.rs`'s `fetch_thrift_link`), Thrift's inline `arrowBatches`
+    /// (`pipeline.rs`'s `run_thrift_fetch_loop`), and SEA's own chunk fetch
+    /// (`fetch_chunks_with_backpressure`, below, though only when a
+    /// `chunk_index` resolves to exactly one blob -- see that function's own
+    /// comment for why more than one blob is left untruncated). `None`
+    /// means no authoritative bound is known, not "this chunk was never
+    /// overshot" -- `decode_chunk_item` treats it as a no-op either way.
     pub truncate_to: Option<i64>,
 }
 
@@ -507,9 +509,21 @@ pub struct DbClient {
     /// rebuild the extension to do it. Also doubles, on the Thrift path, as
     /// `TExecuteStatementReq.canDecompressLZ4Result`.
     compress_results: bool,
-    session_pool: SessionPool,
+    session_pool: Pool<String>,
     pub protocol: Protocol,
-    thrift_session_pool: ThriftSessionPool,
+    /// Same checkout/checkin shape as `session_pool` above, plus: Thrift's
+    /// session model is *mandatory* per statement, unlike SEA's optional
+    /// `session_id`, and this crate has not independently confirmed whether
+    /// concurrent statements on one Thrift session are safe against a real
+    /// workspace the way the SEA crash was -- so this pool exists
+    /// defensively, following the same proven-safe pattern regardless of
+    /// whether the analogous crash reproduces here. Unlike SEA, there is no
+    /// "sessionless" fallback available at all -- `TExecuteStatementReq.
+    /// sessionHandle` is required by the protocol -- so a pool-exhaustion/
+    /// creation-failure `None` from `thrift_checkout_session` means the
+    /// caller must open one throwaway session for that single call and
+    /// close it again immediately after (see `execute_lazy_thrift`).
+    thrift_session_pool: Pool<thrift::SessionHandle>,
     /// Global budget for concurrent cloud-fetch HTTP requests, sized to
     /// `chunk_fetch_concurrency`. Every link download takes one permit; a
     /// download that finds spare permits (few links in flight -- exactly
@@ -529,24 +543,12 @@ pub struct DbClient {
 /// ever split into -- see `DbClient::download_slots`.
 pub const MAX_SPLIT_PARTS: usize = 8;
 
-/// Session pool for the Thrift backend -- same checkout/checkin shape and
-/// the exact same two hard constraints as `SessionPool` above (a session is
-/// created *for* one (catalog, schema) pair and can't be redirected; Thrift's
-/// session model is *mandatory* per statement, unlike SEA's optional
-/// `session_id`, and this crate has not independently confirmed whether
-/// concurrent statements on one Thrift session are safe against a real
-/// workspace the way the SEA crash was -- so this pool exists defensively,
-/// following the same proven-safe pattern regardless of whether the
-/// analogous crash reproduces here. Unlike SEA, there is no "sessionless"
-/// fallback available at all -- `TExecuteStatementReq.sessionHandle` is
-/// required by the protocol -- so a pool-exhaustion/creation-failure `None`
-/// here means the caller must open one throwaway session for that single
-/// call and close it again immediately after (see `execute_lazy_thrift`).
-#[derive(Default)]
-struct ThriftSessionPool {
-    idle: Mutex<HashMap<(Option<String>, Option<String>), Vec<thrift::SessionHandle>>>,
-    total: Mutex<HashMap<(Option<String>, Option<String>), usize>>,
-}
+/// `chunk_fetch_concurrency`'s own default -- see `DbClient::new`'s own
+/// comment on it for the measured numbers behind picking 64. Named so
+/// `download_slots` (sized to this same number, since it's a budget over
+/// the same worker count) can't drift from it the way a second bare `64`
+/// silently could.
+const DEFAULT_CHUNK_FETCH_CONCURRENCY: usize = 64;
 
 /// How long the Thrift path polls `GetOperationStatus` when a statement
 /// doesn't finish within its `getDirectResults` budget (see
@@ -621,13 +623,77 @@ const THRIFT_DIRECT_RESULTS_MAX_BYTES: i64 = 1024 * 1024 * 1024;
 /// still throws away a perfectly good session), but guarantees a session
 /// that might be in the same bad state behind the SparkSession-null crash
 /// above is never handed to a second caller.
-#[derive(Default)]
-struct SessionPool {
+type SessionKey = (Option<String>, Option<String>);
+
+/// Generic checkout/checkin pool keyed by (catalog, schema) -- shared by
+/// SEA's session ids (`Pool<String>`) and Thrift's session handles
+/// (`Pool<thrift::SessionHandle>`), which had this exact logic duplicated
+/// twice over before this. See `checkout_session`/`thrift_checkout_session`
+/// for the constraints that shape it (session pinned to one (catalog,
+/// schema) pair, safe for sequential reuse only, checkout never blocks on
+/// exhaustion).
+struct Pool<T> {
     // ponytail: fixed cap, not a constructor kwarg -- nothing's asked to
     // tune this yet; raise (or expose one) if a workload needs more
     // concurrent sessions per (catalog, schema) pair than this.
-    idle: Mutex<HashMap<(Option<String>, Option<String>), Vec<String>>>,
-    total: Mutex<HashMap<(Option<String>, Option<String>), usize>>,
+    idle: Mutex<HashMap<SessionKey, Vec<T>>>,
+    total: Mutex<HashMap<SessionKey, usize>>,
+}
+
+// Manual impl instead of `#[derive(Default)]`: the derive would wrongly
+// require `T: Default` even though neither field actually needs it.
+impl<T> Default for Pool<T> {
+    fn default() -> Self {
+        Self {
+            idle: Mutex::new(HashMap::new()),
+            total: Mutex::new(HashMap::new()),
+        }
+    }
+}
+
+impl<T> Pool<T> {
+    fn take(&self, key: &SessionKey) -> Option<T> {
+        self.idle.lock().unwrap().get_mut(key).and_then(Vec::pop)
+    }
+
+    /// Reserves a slot for `key` if under `MAX_SESSIONS_PER_KEY`. The
+    /// increment happens before the caller's own creation call so two
+    /// concurrent callers can't both squeeze past the cap; call `release` if
+    /// creation then fails.
+    fn reserve(&self, key: &SessionKey) -> bool {
+        let mut total = self.total.lock().unwrap();
+        let count = total.entry(key.clone()).or_insert(0);
+        if *count >= MAX_SESSIONS_PER_KEY {
+            return false;
+        }
+        *count += 1;
+        true
+    }
+
+    fn release(&self, key: &SessionKey) {
+        let mut total = self.total.lock().unwrap();
+        if let Some(count) = total.get_mut(key) {
+            *count = count.saturating_sub(1);
+        }
+    }
+
+    fn put(&self, key: SessionKey, item: T) {
+        self.idle.lock().unwrap().entry(key).or_default().push(item);
+    }
+
+    /// Returns `item` to the pool for reuse (`keep = true`) or discards it,
+    /// releasing its reservation instead (`keep = false`).
+    fn checkin(&self, key: SessionKey, item: T, keep: bool) {
+        if keep {
+            self.put(key, item);
+        } else {
+            self.release(&key);
+        }
+    }
+
+    fn drain_idle(&self) -> Vec<T> {
+        self.idle.lock().unwrap().drain().flat_map(|(_, v)| v).collect()
+    }
 }
 
 pub const MAX_SESSIONS_PER_KEY: usize = 8;
@@ -672,15 +738,15 @@ impl DbClient {
             // a claim that 64 is the true optimum. A caller with a very
             // large chunk count or a fast/low-latency link to the warehouse
             // may still want to raise it further.
-            chunk_fetch_concurrency: 64,
+            chunk_fetch_concurrency: DEFAULT_CHUNK_FETCH_CONCURRENCY,
             warehouse_start_timeout: Duration::from_secs(300),
             warehouse_confirmed_running_ttl: Duration::from_secs(30),
             warehouse_confirmed_running_at: Mutex::new(None),
             compress_results: true,
-            session_pool: SessionPool::default(),
+            session_pool: Pool::default(),
             protocol: Protocol::Thrift,
-            thrift_session_pool: ThriftSessionPool::default(),
-            download_slots: tokio::sync::Semaphore::new(64),
+            thrift_session_pool: Pool::default(),
+            download_slots: tokio::sync::Semaphore::new(DEFAULT_CHUNK_FETCH_CONCURRENCY),
         }
     }
 
@@ -690,11 +756,7 @@ impl DbClient {
         self
     }
 
-    /// Selects the wire protocol/backend -- `Protocol::Sea` (this bare
-    /// `DbClient` constructor's own internal starting value, see
-    /// `Protocol`'s own doc comment for why that's not the same thing as
-    /// "the user-facing default") or `Protocol::Thrift` (the actual
-    /// user-facing default as of this session's benchmarking work).
+    /// Selects the wire protocol/backend -- see `Protocol`'s own doc comment.
     pub fn with_protocol(mut self, protocol: Protocol) -> Self {
         self.protocol = protocol;
         self
@@ -830,11 +892,9 @@ impl DbClient {
     }
 
     /// Downloads one cloud-fetch link, splitting it across parallel HTTP
-    /// Range requests when (and only when) the shared `download_slots`
-    /// budget has room to spare -- i.e. when few links are in flight, which
-    /// is exactly the case a single-chunk result hits and where a lone TCP
-    /// stream leaves most of the link idle. See `download_slots`' own doc
-    /// comment for the measured win and why the large-result case is safe.
+    /// Range requests when the shared `download_slots` budget has room to
+    /// spare -- see that field's own doc comment for the measured win and
+    /// why the large-result case is safe.
     pub(crate) async fn fetch_link_bytes_budgeted(
         self: &Arc<Self>,
         url: &str,
@@ -843,8 +903,16 @@ impl DbClient {
         // One permit per download is mandatory; the worker pool already
         // bounds concurrent links to `chunk_fetch_concurrency`, so this
         // never blocks in practice -- it just makes the budget accounting
-        // exact.
-        let _base = self.download_slots.acquire().await;
+        // exact. `acquire()`'s `Err` case (the semaphore closed) can't
+        // happen -- nothing in this crate ever calls `.close()` on
+        // `download_slots` -- but propagating it as a real error instead of
+        // an `.expect()` costs nothing and avoids a panic if that ever
+        // changes.
+        let _base = self
+            .download_slots
+            .acquire()
+            .await
+            .map_err(|e| ApiError::permanent(format!("download_slots semaphore closed unexpectedly: {e}")))?;
         let want = (MAX_SPLIT_PARTS - 1).min(self.download_slots.available_permits());
         let extra = if want > 0 {
             self.download_slots.try_acquire_many(want as u32).ok()
@@ -960,8 +1028,33 @@ impl DbClient {
 
         let mut out = bytes::BytesMut::with_capacity(total.unwrap_or(head.len() as u64) as usize);
         out.extend_from_slice(&head);
-        for h in handles {
-            out.extend_from_slice(&h.await.map_err(join_error)??);
+        // Parts must concatenate in order, so this can't use `JoinSet`
+        // (which yields in completion order) without tracking indices --
+        // simpler to keep the `Vec` and explicitly `.abort()` every
+        // not-yet-awaited sibling the moment one part fails, rather than
+        // silently leaving them running (a bare `?` here would return
+        // early and just drop the rest, which does NOT cancel them --
+        // `JoinHandle::drop` detaches, it doesn't abort -- leaving up to
+        // `MAX_SPLIT_PARTS - 1` sibling Range downloads, each with their
+        // own `retry_call` backoff, still in flight for a link the caller
+        // has already given up on).
+        let mut iter = handles.into_iter();
+        while let Some(h) = iter.next() {
+            match h.await {
+                Ok(Ok(part)) => out.extend_from_slice(&part),
+                Ok(Err(e)) => {
+                    for remaining in iter {
+                        remaining.abort();
+                    }
+                    return Err(e);
+                }
+                Err(join_err) => {
+                    for remaining in iter {
+                        remaining.abort();
+                    }
+                    return Err(join_error(join_err));
+                }
+            }
         }
         let bytes = out.freeze();
         if let Some(t) = total
@@ -981,19 +1074,15 @@ impl DbClient {
     }
 
     /// 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.
+    /// nothing SEA- or Thrift-specific about it (see AGENTS.md for why the
+    /// Thrift path didn't always call this).
     pub(crate) async fn ensure_warehouse_running(&self) -> Result<(), ApiError> {
         {
             let confirmed = *self.warehouse_confirmed_running_at.lock().unwrap();
-            if let Some(at) = confirmed {
-                if at.elapsed() < self.warehouse_confirmed_running_ttl {
-                    return Ok(());
-                }
+            if let Some(at) = confirmed
+                && at.elapsed() < self.warehouse_confirmed_running_ttl
+            {
+                return Ok(());
             }
         }
 
@@ -1048,36 +1137,20 @@ impl DbClient {
     /// Hands back a pooled session for (`catalog`, `schema`) if one's idle,
     /// creates one if the pool for that key isn't at `MAX_SESSIONS_PER_KEY`
     /// yet, or `None` if neither -- the caller falls back to a plain
-    /// session-less submission in that case, see `SessionPool`'s own doc
-    /// comment for why this never blocks instead. The slot-reservation
-    /// increment happens *before* the `create_session` await so two
-    /// concurrent callers can't both squeeze past the cap; a failed creation
-    /// releases the reservation again.
+    /// session-less submission in that case, see `session_pool`'s own doc
+    /// comment for why this never blocks instead.
     async fn checkout_session(&self, catalog: Option<&str>, schema: Option<&str>) -> Option<String> {
         let key = (catalog.map(str::to_string), schema.map(str::to_string));
-        {
-            let mut idle = self.session_pool.idle.lock().unwrap();
-            if let Some(ids) = idle.get_mut(&key)
-                && let Some(id) = ids.pop()
-            {
-                return Some(id);
-            }
+        if let Some(id) = self.session_pool.take(&key) {
+            return Some(id);
         }
-        {
-            let mut total = self.session_pool.total.lock().unwrap();
-            let count = total.entry(key.clone()).or_insert(0);
-            if *count >= MAX_SESSIONS_PER_KEY {
-                return None;
-            }
-            *count += 1;
+        if !self.session_pool.reserve(&key) {
+            return None;
         }
         match self.create_session(catalog, schema).await {
             Ok(id) => Some(id),
             Err(_) => {
-                let mut total = self.session_pool.total.lock().unwrap();
-                if let Some(count) = total.get_mut(&key) {
-                    *count = count.saturating_sub(1);
-                }
+                self.session_pool.release(&key);
                 None
             }
         }
@@ -1085,24 +1158,11 @@ impl DbClient {
 
     /// Returns a session to the pool for reuse (`keep = true`, the statement
     /// it backed reached SUCCEEDED/FAILED/CANCELED cleanly) or discards it
-    /// (`keep = false`) -- see `SessionPool`'s doc comment for why any error
+    /// (`keep = false`) -- see `session_pool`'s doc comment for why any error
     /// discards rather than reuses.
     fn checkin_session(&self, catalog: Option<&str>, schema: Option<&str>, session_id: String, keep: bool) {
         let key = (catalog.map(str::to_string), schema.map(str::to_string));
-        if keep {
-            self.session_pool
-                .idle
-                .lock()
-                .unwrap()
-                .entry(key)
-                .or_default()
-                .push(session_id);
-        } else {
-            let mut total = self.session_pool.total.lock().unwrap();
-            if let Some(count) = total.get_mut(&key) {
-                *count = count.saturating_sub(1);
-            }
-        }
+        self.session_pool.checkin(key, session_id, keep);
     }
 
     /// Best-effort cleanup of every currently-idle pooled session -- meant
@@ -1113,11 +1173,7 @@ impl DbClient {
     /// this before every pending statement has finished is a caller
     /// ordering issue, not something this method can fix from inside.
     pub async fn close_all_sessions(&self) {
-        let ids: Vec<String> = {
-            let mut idle = self.session_pool.idle.lock().unwrap();
-            idle.drain().flat_map(|(_, v)| v).collect()
-        };
-        for id in ids {
+        for id in self.session_pool.drain_idle() {
             self.delete_session(&id).await;
         }
     }
@@ -1208,8 +1264,8 @@ impl DbClient {
     }
 
     /// Same checkout contract as SEA's `checkout_session` (see
-    /// `ThriftSessionPool`'s doc comment): `None` means the caller must open
-    /// its own throwaway session for this one call (Thrift has no
+    /// `thrift_session_pool`'s doc comment): `None` means the caller must
+    /// open its own throwaway session for this one call (Thrift has no
     /// session-less submission mode to fall back to).
     pub(crate) async fn thrift_checkout_session(
         &self,
@@ -1217,29 +1273,16 @@ impl DbClient {
         schema: Option<&str>,
     ) -> Option<thrift::SessionHandle> {
         let key = (catalog.map(str::to_string), schema.map(str::to_string));
-        {
-            let mut idle = self.thrift_session_pool.idle.lock().unwrap();
-            if let Some(ids) = idle.get_mut(&key)
-                && let Some(id) = ids.pop()
-            {
-                return Some(id);
-            }
+        if let Some(handle) = self.thrift_session_pool.take(&key) {
+            return Some(handle);
         }
-        {
-            let mut total = self.thrift_session_pool.total.lock().unwrap();
-            let count = total.entry(key.clone()).or_insert(0);
-            if *count >= MAX_SESSIONS_PER_KEY {
-                return None;
-            }
-            *count += 1;
+        if !self.thrift_session_pool.reserve(&key) {
+            return None;
         }
         match self.thrift_open_session_raw(catalog, schema).await {
             Ok(handle) => Some(handle),
             Err(_) => {
-                let mut total = self.thrift_session_pool.total.lock().unwrap();
-                if let Some(count) = total.get_mut(&key) {
-                    *count = count.saturating_sub(1);
-                }
+                self.thrift_session_pool.release(&key);
                 None
             }
         }
@@ -1253,30 +1296,13 @@ impl DbClient {
         keep: bool,
     ) {
         let key = (catalog.map(str::to_string), schema.map(str::to_string));
-        if keep {
-            self.thrift_session_pool
-                .idle
-                .lock()
-                .unwrap()
-                .entry(key)
-                .or_default()
-                .push(session);
-        } else {
-            let mut total = self.thrift_session_pool.total.lock().unwrap();
-            if let Some(count) = total.get_mut(&key) {
-                *count = count.saturating_sub(1);
-            }
-        }
+        self.thrift_session_pool.checkin(key, session, keep);
     }
 
     /// Best-effort close of every currently-idle pooled Thrift session --
     /// same contract as `close_all_sessions` (SEA).
     pub async fn close_all_thrift_sessions(&self) {
-        let sessions: Vec<thrift::SessionHandle> = {
-            let mut idle = self.thrift_session_pool.idle.lock().unwrap();
-            idle.drain().flat_map(|(_, v)| v).collect()
-        };
-        for s in sessions {
+        for s in self.thrift_session_pool.drain_idle() {
             self.thrift_close_session_raw(&s).await;
         }
     }
@@ -1312,7 +1338,6 @@ impl DbClient {
         op: &thrift::OperationHandle,
     ) -> Result<thrift::OperationStatusResp, ApiError> {
         let body = Bytes::from(thrift::build_get_operation_status(op));
-        // Idempotent -- read-only.
         let resp_bytes = self.thrift_call(body, true).await?;
         thrift::parse_get_operation_status(&resp_bytes).map_err(Self::thrift_parse_error)
     }
@@ -1445,24 +1470,6 @@ impl DbClient {
         }
     }
 
-    /// Like `execute_arrow_statement`, fixed to JSON_ARRAY -- each fetched
-    /// chunk's bytes are then a JSON array of rows, each row itself an array
-    /// of values where every non-null value is a *string* regardless of its
-    /// real column type (Databricks' own JSON_ARRAY contract; null stays
-    /// JSON null) -- casting by the manifest's column type_name, if wanted,
-    /// is left to the caller, same as the Python original does nothing extra
-    /// here either.
-    pub async fn execute_json_statement(
-        &self,
-        statement: &str,
-        catalog: Option<&str>,
-        schema: Option<&str>,
-        parameters: Option<Value>,
-    ) -> Result<StatementSubmitResult, ApiError> {
-        self.execute_statement(statement, "JSON_ARRAY", catalog, schema, parameters)
-            .await
-    }
-
     /// Shared by `execute_statement` and `execute_arrow_statement_prefer_inline`:
     /// POST the statement, poll until a terminal state, and turn FAILED/
     /// CANCELED into an `Err` -- everything both callers need before they
@@ -1470,7 +1477,7 @@ impl DbClient {
     ///
     /// `body` must *not* already carry `catalog`/`schema`/`session_id` --
     /// this method owns that decision: a pooled session for (`catalog`,
-    /// `schema`) if one's available (see `checkout_session`/`SessionPool`),
+    /// `schema`) if one's available (see `checkout_session`),
     /// falling back to setting `catalog`/`schema` directly on the body
     /// otherwise (Databricks rejects `session_id` combined with either
     /// field). The session, if any, is returned to the pool on a clean
@@ -1608,12 +1615,8 @@ impl DbClient {
             self.host, statement_id, chunk_index
         );
         let data: ChunkLinksBody = self.authed_json(reqwest::Method::GET, &url, None).await?;
-
-        let mut blobs = Vec::with_capacity(data.external_links.len());
-        for link in data.external_links {
-            blobs.push(self.fetch_link_bytes(&link.external_link, compressed).await?);
-        }
-        Ok(blobs)
+        let links: Vec<String> = data.external_links.into_iter().map(|l| l.external_link).collect();
+        self.fetch_pre_resolved_links(&links, compressed).await
     }
 
     /// Same shape as `fetch_chunk_index`, but for links the statement submit/
@@ -1668,19 +1671,9 @@ 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.
+                                // One blob ⇒ `meta.row_count` (the whole chunk_index's
+                                // declared count) and "this blob's count" are the same
+                                // number -- see `ChunkItem::truncate_to`'s own doc comment.
                                 let truncate_to = if blobs.len() == 1 { meta.row_count } else { None };
                                 for blob in blobs {
                                     let item = ChunkItem {
@@ -1758,8 +1751,6 @@ impl DbClient {
                 .timeout(self.http_timeout)
                 .send()
                 .await
-                // DELETE is naturally idempotent here too (404 is already
-                // treated as success below).
                 .map_err(|e| ApiError::from_reqwest(e, true))?;
             let status = resp.status();
             if status == StatusCode::NOT_FOUND || status.is_success() {