arrowbricks 5.0.1__tar.gz → 5.0.2__tar.gz

{arrowbricks-5.0.1 → arrowbricks-5.0.2}/PKG-INFO

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

{arrowbricks-5.0.1 → arrowbricks-5.0.2}/pyproject.toml

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

@@ -177,7 +177,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "5.0.1"
+version = "5.0.2"
 dependencies = [
  "arrow-array",
  "arrow-buffer",

{arrowbricks-5.0.1 → arrowbricks-5.0.2}/rust/arrowbricks_core/Cargo.toml

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

{arrowbricks-5.0.1 → arrowbricks-5.0.2}/rust/arrowbricks_core/src/client/download.rs

@@ -67,6 +67,21 @@ pub(crate) fn decompress_lz4_frame(compressed: &Bytes) -> Result<Bytes, ApiError
     Ok(Bytes::from(out))
 }
 
+/// Reads the body into one buffer sized from `Content-Length`. Frames are
+/// freed as they are consumed, so a chunk peaks at about its own size instead
+/// of the ~2x `Response::bytes()` reaches while joining collected frames.
+async fn read_body(mut resp: reqwest::Response) -> Result<Bytes, reqwest::Error> {
+    // Cap the up-front reservation so a bogus header can't demand gigabytes.
+    let Some(len) = resp.content_length().map(|n| n.min(1 << 30) as usize) else {
+        return resp.bytes().await;
+    };
+    let mut buf = Vec::with_capacity(len);
+    while let Some(frame) = resp.chunk().await? {
+        buf.extend_from_slice(&frame);
+    }
+    Ok(Bytes::from(buf))
+}
+
 impl DbClient {
     /// Unauthenticated -- external links are presigned blob-storage URLs,
     /// same as `_fetch_link_bytes` in the Python client. `compressed` decodes
@@ -99,7 +114,7 @@ impl DbClient {
                 let text = resp.text().await.unwrap_or_default();
                 return Err(ApiError::from_status(status, &text, true));
             }
-            let bytes = resp.bytes().await.map_err(|e| ApiError::from_reqwest(e, true))?;
+            let bytes = read_body(resp).await.map_err(|e| ApiError::from_reqwest(e, true))?;
             // Counted here, before decompression -- "downloaded" means bytes
             // actually received off the wire, which is exactly the smaller,
             // (usually) LZ4-compressed size `compress_results` exists to