diff --git a/Cargo.lock b/Cargo.lock index 61fe8e3..7ff5bb1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -123,50 +123,34 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5ec52ba94edeed950e4a41f75d35376df196e8cb04437f7280a5aa49f20f796" +checksum = "f3f15b4c6b148206ff3a2b35002e08929c2462467b62b9c02036d9c34f9ef994" dependencies = [ "arrow-arith", - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", + "arrow-array", + "arrow-buffer", "arrow-cast", - "arrow-data 54.3.1", - "arrow-ipc 54.3.1", + "arrow-data", + "arrow-ipc", "arrow-ord", "arrow-row", - "arrow-schema 54.3.1", + "arrow-schema", "arrow-select", "arrow-string", ] [[package]] name = "arrow-arith" -version = "54.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc766fdacaf804cb10c7c70580254fcdb5d55cdfda2bc57b02baf5223a3af9e" -dependencies = [ - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", - "chrono", - "num", -] - -[[package]] -name = "arrow-array" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a12fcdb3f1d03f69d3ec26ac67645a8fe3f878d77b5ebb0b15d64a116c212985" +checksum = "30feb679425110209ae35c3fbf82404a39a4c0436bb3ec36164d8bffed2a4ce4" dependencies = [ - "ahash", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", "chrono", - "half", - "hashbrown 0.15.5", "num", ] @@ -177,26 +161,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "70732f04d285d49054a48b72c54f791bb3424abae92d27aafdf776c98af161c8" dependencies = [ "ahash", - "arrow-buffer 55.2.0", - "arrow-data 55.2.0", - "arrow-schema 55.2.0", + "arrow-buffer", + "arrow-data", + "arrow-schema", "chrono", "half", "hashbrown 0.15.5", "num", ] -[[package]] -name = "arrow-buffer" -version = "54.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "263f4801ff1839ef53ebd06f99a56cecd1dbaf314ec893d93168e2e860e0291c" -dependencies = [ - "bytes", - "half", - "num", -] - [[package]] name = "arrow-buffer" version = "55.2.0" @@ -210,14 +183,14 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ede6175fbc039dfc946a61c1b6d42fd682fcecf5ab5d148fbe7667705798cac9" +checksum = "e4f12eccc3e1c05a766cafb31f6a60a46c2f8efec9b74c6e0648766d30686af8" dependencies = [ - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", "arrow-select", "atoi", "base64", @@ -228,88 +201,57 @@ dependencies = [ "ryu", ] -[[package]] -name = "arrow-data" -version = "54.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61cfdd7d99b4ff618f167e548b2411e5dd2c98c0ddebedd7df433d34c20a4429" -dependencies = [ - "arrow-buffer 54.3.1", - "arrow-schema 54.3.1", - "half", - "num", -] - [[package]] name = "arrow-data" version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8de1ce212d803199684b658fc4ba55fb2d7e87b213de5af415308d2fee3619c2" dependencies = [ - "arrow-buffer 55.2.0", - "arrow-schema 55.2.0", + "arrow-buffer", + "arrow-schema", "half", "num", ] -[[package]] -name = "arrow-ipc" -version = "54.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62ff528658b521e33905334723b795ee56b393dbe9cf76c8b1f64b648c65a60c" -dependencies = [ - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", - "flatbuffers 24.12.23", -] - [[package]] name = "arrow-ipc" version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9ea5967e8b2af39aff5d9de2197df16e305f47f404781d3230b2dc672da5d92" dependencies = [ - "arrow-array 55.2.0", - "arrow-buffer 55.2.0", - "arrow-data 55.2.0", - "arrow-schema 55.2.0", - "flatbuffers 25.12.19", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "flatbuffers", ] [[package]] name = "arrow-ord" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0a3334a743bd2a1479dbc635540617a3923b4b2f6870f37357339e6b5363c21" +checksum = "6506e3a059e3be23023f587f79c82ef0bcf6d293587e3272d20f2d30b969b5a7" dependencies = [ - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", "arrow-select", ] [[package]] name = "arrow-row" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d1d7a7291d2c5107e92140f75257a99343956871f3d3ab33a7b41532f79cb68" +checksum = "52bf7393166beaf79b4bed9bfdf19e97472af32ce5b6b48169d321518a08cae2" dependencies = [ - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", "half", ] -[[package]] -name = "arrow-schema" -version = "54.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39cfaf5e440be44db5413b75b72c2a87c1f8f0627117d110264048f2969b99e9" - [[package]] name = "arrow-schema" version = "55.2.0" @@ -318,28 +260,28 @@ checksum = "af7686986a3bf2254c9fb130c623cdcb2f8e1f15763e7c71c310f0834da3d292" [[package]] name = "arrow-select" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69efcd706420e52cd44f5c4358d279801993846d1c2a8e52111853d61d55a619" +checksum = "dd2b45757d6a2373faa3352d02ff5b54b098f5e21dccebc45a21806bc34501e5" dependencies = [ "ahash", - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", "num", ] [[package]] name = "arrow-string" -version = "54.3.1" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a21546b337ab304a32cfc0770f671db7411787586b45b78b4593ae78e64e2b03" +checksum = "0377d532850babb4d927a06294314b316e23311503ed580ec6ce6a0158f49d40" dependencies = [ - "arrow-array 54.3.1", - "arrow-buffer 54.3.1", - "arrow-data 54.3.1", - "arrow-schema 54.3.1", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", "arrow-select", "memchr", "num", @@ -423,12 +365,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "bitflags" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" - [[package]] name = "bitflags" version = "2.12.1" @@ -734,7 +670,7 @@ version = "0.28.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" dependencies = [ - "bitflags 2.12.1", + "bitflags", "crossterm_winapi", "mio", "parking_lot", @@ -750,7 +686,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" dependencies = [ - "bitflags 2.12.1", + "bitflags", "crossterm_winapi", "derive_more", "document-features", @@ -981,23 +917,13 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" -[[package]] -name = "flatbuffers" -version = "24.12.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" -dependencies = [ - "bitflags 1.3.2", - "rustc_version", -] - [[package]] name = "flatbuffers" version = "25.12.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" dependencies = [ - "bitflags 2.12.1", + "bitflags", "rustc_version", ] @@ -1242,9 +1168,9 @@ name = "hotdata" version = "0.1.0" source = "git+https://github.com/hotdata-dev/sdk-rust?rev=8d4018fb899ba52228db44eaffa6caa0eb5b603f#8d4018fb899ba52228db44eaffa6caa0eb5b603f" dependencies = [ - "arrow-array 55.2.0", - "arrow-ipc 55.2.0", - "arrow-schema 55.2.0", + "arrow-array", + "arrow-ipc", + "arrow-schema", "async-trait", "bytes", "futures-core", @@ -1605,7 +1531,7 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6654738b8024300cf062d04a1c13c10c8e2cea598ec1c47dc9b6641159429756" dependencies = [ - "bitflags 2.12.1", + "bitflags", "crossterm 0.29.0", "dyn-clone", "fuzzy-matcher", @@ -1958,7 +1884,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" dependencies = [ - "bitflags 2.12.1", + "bitflags", "cfg-if", "cfg_aliases", "libc", @@ -2075,7 +2001,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags 2.12.1", + "bitflags", ] [[package]] @@ -2117,7 +2043,7 @@ version = "0.10.80" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a45fa2aa886c42762255da344f0a0d313e254066c46aad76f300c3d3da62d967" dependencies = [ - "bitflags 2.12.1", + "bitflags", "cfg-if", "foreign-types", "libc", @@ -2462,7 +2388,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags 2.12.1", + "bitflags", ] [[package]] @@ -2650,7 +2576,7 @@ version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags 2.12.1", + "bitflags", "errno", "libc", "linux-raw-sys 0.4.15", @@ -2663,7 +2589,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags 2.12.1", + "bitflags", "errno", "libc", "linux-raw-sys 0.12.1", @@ -2725,7 +2651,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2812,7 +2738,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags 2.12.1", + "bitflags", "core-foundation 0.10.1", "core-foundation-sys", "libc", @@ -3162,7 +3088,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" dependencies = [ - "bitflags 2.12.1", + "bitflags", "core-foundation 0.9.4", "system-configuration-sys", ] @@ -3465,7 +3391,7 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ - "bitflags 2.12.1", + "bitflags", "bytes", "futures-util", "http", @@ -3759,7 +3685,7 @@ version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ - "bitflags 2.12.1", + "bitflags", "hashbrown 0.15.5", "indexmap 2.14.0", "semver", @@ -3825,7 +3751,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4193,7 +4119,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", - "bitflags 2.12.1", + "bitflags", "indexmap 2.14.0", "log", "serde", diff --git a/Cargo.toml b/Cargo.toml index e1b1a68..7a5c1d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,7 @@ reqwest = { version = "0.13", features = ["blocking", "json"] } rayon = "1.10" serde = { version = "1", features = ["derive"] } serde_json = "1" -arrow = { version = "54", default-features = false, features = ["ipc"] } +arrow = { version = "55", default-features = false, features = ["ipc"] } serde_yaml = "0.9" base64 = "0.22" crossterm = "0.28" diff --git a/src/query.rs b/src/query.rs index 8e2a737..f9dd783 100644 --- a/src/query.rs +++ b/src/query.rs @@ -2,8 +2,6 @@ use crate::sdk::Api; use serde::Deserialize; use serde_json::Value; -const ACCEPT_ARROW: &str = "application/vnd.apache.arrow.stream"; - #[derive(Deserialize)] pub struct QueryResponse { pub result_id: Option, @@ -162,35 +160,21 @@ fn arrow_cell(col: &dyn arrow::array::Array, row: usize) -> Value { } } -/// Decode an Arrow IPC stream into a `QueryResponse` suitable for display. -fn arrow_ipc_to_query_response(bytes: Vec, result_id: String) -> QueryResponse { - use arrow::ipc::reader::StreamReader; - use std::io::Cursor; - - let reader = match StreamReader::try_new(Cursor::new(&bytes), None) { - Ok(r) => r, - Err(e) => { - eprintln!("error reading Arrow IPC stream: {e}"); - std::process::exit(1); - } - }; - - let columns: Vec = reader - .schema() +/// Convert an SDK-decoded [`hotdata::ArrowResult`] into a `QueryResponse` +/// suitable for display. +fn arrow_result_to_query_response( + result: hotdata::ArrowResult, + result_id: String, +) -> QueryResponse { + let columns: Vec = result + .schema .fields() .iter() .map(|f| f.name().clone()) .collect(); let mut rows: Vec> = Vec::new(); - for batch_result in reader { - let batch = match batch_result { - Ok(b) => b, - Err(e) => { - eprintln!("error reading Arrow batch: {e}"); - std::process::exit(1); - } - }; + for batch in &result.batches { for row in 0..batch.num_rows() { rows.push( (0..batch.num_columns()) @@ -211,23 +195,14 @@ fn arrow_ipc_to_query_response(bytes: Vec, result_id: String) -> QueryRespon } } -/// Fetch `/results/{result_id}` as Arrow IPC and return a `QueryResponse`. +/// Fetch `/results/{result_id}` as Arrow and return a `QueryResponse`. /// -/// The Arrow stream is fetched through the SDK seam ([`Api::get_bytes`]) — same -/// auth/transport as every other call — but decoded here with the CLI's own -/// pinned `arrow` crate rather than the SDK's `get_result_arrow` (whose -/// `RecordBatch` comes from a different `arrow` major version). +/// Both transport and decode are owned by the SDK's `get_result_arrow` (via the +/// [`Api::get_result_arrow`] seam), so the CLI shares one `arrow` major version +/// with the SDK. pub(crate) fn fetch_arrow_result(api: &Api, result_id: &str) -> QueryResponse { - let (status, bytes) = api - .get_bytes(&format!("/results/{result_id}"), ACCEPT_ARROW) - .unwrap_or_else(|e| e.exit()); - if !status.is_success() { - use crossterm::style::Stylize; - let msg = String::from_utf8_lossy(&bytes); - eprintln!("{}", format!("error fetching result: {status} {msg}").red()); - std::process::exit(1); - } - arrow_ipc_to_query_response(bytes, result_id.to_owned()) + let result = api.get_result_arrow(result_id).unwrap_or_else(|e| e.exit()); + arrow_result_to_query_response(result, result_id.to_owned()) } pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &str) { diff --git a/src/sdk.rs b/src/sdk.rs index bf695a8..fba56a9 100644 --- a/src/sdk.rs +++ b/src/sdk.rs @@ -129,6 +129,33 @@ impl ApiError { } } + /// Map the SDK's [`hotdata::ArrowError`] (the Arrow result-fetch error type, + /// which is *not* an `Error`) into an [`ApiError`]. + /// + /// Status-bearing variants are preserved as [`ApiError::Status`] so the 4xx + /// re-auth hint in [`exit`](Self::exit) still fires; the statusless variants + /// (decode/transport/not-ready/failed) collapse to [`ApiError::Transport`] + /// carrying the SDK's own descriptive message. `ArrowError` is + /// `#[non_exhaustive]`, hence the wildcard arm. + pub fn from_arrow(err: hotdata::ArrowError) -> Self { + use hotdata::ArrowError; + match err { + ArrowError::NotFound => ApiError::Status { + status: reqwest::StatusCode::NOT_FOUND, + body: "result not found".to_string(), + }, + ArrowError::InvalidParams { ref message } => ApiError::Status { + status: reqwest::StatusCode::BAD_REQUEST, + body: message.clone(), + }, + ArrowError::Http { status, ref body } => ApiError::Status { + status, + body: body.clone(), + }, + other => ApiError::Transport(other.to_string()), + } + } + /// Print the standard error and exit, reproducing `ApiClient::fail_response`. /// /// On a 4xx, re-probe the auth status so a masked 404/403 is upgraded into @@ -599,44 +626,17 @@ impl Api { }) } - /// Issue an authenticated `GET {base}/v1{path}` with a custom `Accept` - /// header through the SDK `Configuration`, returning the raw status + body - /// bytes. + /// Fetch `/v1/results/{id}` as Arrow IPC and decode it through the SDK's + /// `get_result_arrow`, returning the fully-buffered [`hotdata::ArrowResult`]. /// - /// The seam's binary-body counterpart to [`get_json`](Self::get_json): used - /// for the Arrow IPC result fetch (`/results/{id}`), where the CLI decodes - /// the stream itself with its own pinned `arrow` crate version rather than - /// the SDK's `get_result_arrow` (which returns a `RecordBatch` from a - /// different `arrow` major). The seam still owns auth/transport (same - /// reqwest client, bearer via the `token_provider`, `X-Workspace-Id`). - pub fn get_bytes( - &self, - path: &str, - accept: &str, - ) -> Result<(reqwest::StatusCode, Vec), ApiError> { - let cfg = self.client.configuration(); - let url = format!("{}/v1{path}", cfg.base_path); - let database_id = self.database_id.clone(); - let session_id = self.session_id.clone(); - let accept = accept.to_string(); - rt().block_on(async move { - let mut req = cfg - .client - .request(reqwest::Method::GET, &url) - .header(reqwest::header::ACCEPT, accept); - req = apply_seam_headers(req, cfg, session_id.as_deref(), database_id.as_deref()).await; - - let resp = req - .send() - .await - .map_err(|e| ApiError::Transport(format!("error connecting to API: {e}")))?; - let status = resp.status(); - let bytes = resp - .bytes() - .await - .map_err(|e| ApiError::Transport(format!("error connecting to API: {e}")))?; - Ok((status, bytes.to_vec())) - }) + /// The SDK owns transport (same reqwest client, bearer via the + /// `token_provider`, `X-Workspace-Id`/`X-Session-Id`) and decode. Its + /// `ArrowError` (the Arrow-path error type, which is not an `Error`) is + /// mapped to [`ApiError`] via [`from_arrow`](ApiError::from_arrow) so callers + /// keep the same `.exit()` handling. + pub fn get_result_arrow(&self, id: &str) -> Result { + rt().block_on(self.client.get_result_arrow(id, None, None)) + .map_err(ApiError::from_arrow) } // --- Sample migrated call (workspace.rs uses this) ----------------------- @@ -1013,49 +1013,102 @@ mod tests { } #[test] - fn get_bytes_sends_bearer_workspace_accept_then_returns_body() { - // The seam's binary-body escape hatch (used by query.rs for the Arrow - // result fetch): carry the bearer + X-Workspace-Id + the custom Accept - // like every SDK call, and return the raw status + bytes for CLI-side - // Arrow decoding. + fn get_result_arrow_fetches_decodes_and_forwards_headers() { + use arrow::array::{Int64Array, RecordBatch, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::ipc::writer::StreamWriter; + + // Build a small Arrow IPC stream the mock server can hand back. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + let mut ipc: Vec = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + // The SDK's get_result_arrow carries the bearer + X-Workspace-Id like + // every SDK call and negotiates Arrow via ?format=arrow + Accept. let mut server = mockito::Server::new(); let m = server .mock("GET", "/v1/results/res_1") + .match_query(mockito::Matcher::UrlEncoded( + "format".into(), + "arrow".into(), + )) .match_header("Authorization", "Bearer test-jwt") .match_header("X-Workspace-Id", "ws-1") .match_header("Accept", "application/vnd.apache.arrow.stream") .with_status(200) .with_header("content-type", "application/vnd.apache.arrow.stream") - .with_body(&[0u8, 1, 2, 3][..]) + .with_body(ipc) .create(); let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); - let (status, bytes) = api - .get_bytes("/results/res_1", "application/vnd.apache.arrow.stream") - .expect("get_bytes should succeed"); - assert_eq!(status, reqwest::StatusCode::OK); - assert_eq!(bytes, vec![0u8, 1, 2, 3]); + let result = api + .get_result_arrow("res_1") + .expect("get_result_arrow should succeed"); + assert_eq!(result.num_rows(), 3); + assert_eq!(result.schema.fields().len(), 2); + assert_eq!(result.schema.field(0).name(), "id"); + assert_eq!(result.schema.field(1).name(), "name"); m.assert(); } #[test] - fn get_bytes_returns_non_success_status_with_body() { - // A failed Arrow fetch surfaces the status + body so the caller can - // print the server error (reproducing the old get_bytes control flow, - // which returned (status, bytes) rather than erroring on non-2xx). + fn get_result_arrow_maps_not_found_to_status() { + // A 404 surfaces as ApiError::Status so the CLI's 4xx re-auth hint path + // still fires. let mut server = mockito::Server::new(); let _m = server .mock("GET", "/v1/results/missing") + .match_query(mockito::Matcher::Any) .with_status(404) .with_body("not found") .create(); let api = Api::test_new(&server.url(), "test-jwt", None); - let (status, bytes) = api - .get_bytes("/results/missing", "application/vnd.apache.arrow.stream") - .expect("get_bytes returns Ok even on non-2xx"); - assert_eq!(status, reqwest::StatusCode::NOT_FOUND); - assert_eq!(String::from_utf8_lossy(&bytes), "not found"); + match api.get_result_arrow("missing").unwrap_err() { + ApiError::Status { status, .. } => { + assert_eq!(status, reqwest::StatusCode::NOT_FOUND); + } + other => panic!("expected ApiError::Status, got {other:?}"), + } + } + + #[test] + fn get_result_arrow_preserves_generic_http_status() { + // Statuses outside the SDK's explicitly-mapped set (404/400/409/202) + // come back as ArrowError::Http; from_arrow must preserve them so a + // 401/403 still reaches the CLI's 4xx re-auth hint. 403 stands in for + // that family. + let mut server = mockito::Server::new(); + let _m = server + .mock("GET", "/v1/results/forbidden") + .match_query(mockito::Matcher::Any) + .with_status(403) + .with_body("forbidden") + .create(); + + let api = Api::test_new(&server.url(), "test-jwt", None); + match api.get_result_arrow("forbidden").unwrap_err() { + ApiError::Status { status, body } => { + assert_eq!(status, reqwest::StatusCode::FORBIDDEN); + assert_eq!(body, "forbidden"); + } + other => panic!("expected ApiError::Status, got {other:?}"), + } } #[test]