diff --git a/CHANGELOG.md b/CHANGELOG.md index e58a913..c3abb04 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `open_result_arrow` and `ArrowResultStream`: fetch a result as an Arrow IPC + stream decoded directly off the socket, so peak memory is one record batch + rather than the whole result. `get_result_arrow` and `stream_result_arrow` + both collect the entire body first and are unchanged. Also available as + `Client::open_result_arrow`. ## [0.17.0] - 2026-09-10 diff --git a/Cargo.toml b/Cargo.toml index c854007..00793b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,10 @@ bytes = "^1" futures-core = "^0.3" log = "^0.4" arrow-ipc = { version = "59", optional = true } +# Only for `Buffer`, which `arrow-ipc`'s push decoder names in its public +# signature and does not re-export. Backs the zero-copy handoff from a reqwest +# body chunk (`bytes::Bytes`) to `StreamDecoder::decode`. +arrow-buffer = { version = "59", optional = true } arrow-array = { version = "59", optional = true } arrow-schema = { version = "59", optional = true } @@ -40,7 +44,7 @@ arrow-schema = { version = "59", optional = true } default = ["native-tls"] native-tls = ["reqwest/native-tls"] rustls = ["reqwest/rustls"] -arrow = ["dep:arrow-ipc", "dep:arrow-array", "dep:arrow-schema"] +arrow = ["dep:arrow-ipc", "dep:arrow-array", "dep:arrow-schema", "dep:arrow-buffer"] [dev-dependencies] tokio = { version = "^1.46.0", features = ["rt-multi-thread", "macros", "time"] } diff --git a/README.md b/README.md index ceea558..959a1fe 100644 --- a/README.md +++ b/README.md @@ -246,7 +246,19 @@ async fn main() -> Result<(), Box> { } ``` -Both methods accept `offset` and `limit` for pagination, and both carry the same authentication and workspace headers as the generated operations. They return `ArrowError::NotReady` if the result is still pending or processing — poll `client.get_result(result_id, database_id)` until its status is `ready` first. `ArrowResult` also surfaces the `X-Total-Row-Count` header (`total_row_count`) and the `rel="next"` pagination `Link` (`next_link`). +A third, `open_result_arrow`, decodes straight off the socket: it pulls body chunks as batches are asked for, so peak memory is one record batch rather than the whole result. Use it for a result larger than memory — the other two collect the entire body before returning. Its schema is available before the first batch, and the pooled connection stays checked out until the stream is drained or dropped. + +```rust +let mut stream = client + .open_result_arrow(&result_id, &database_id, None, None) + .await?; +println!("columns: {:?}", stream.schema().fields()); +while let Some(batch) = stream.next_batch().await? { + // ... one batch at a time; the rest is still on the wire +} +``` + +All three accept `offset` and `limit` for pagination, and all carry the same authentication and workspace headers as the generated operations. They return `ArrowError::NotReady` if the result is still pending or processing — poll `client.get_result(result_id, database_id)` until its status is `ready` first. `ArrowResult` also surfaces the `X-Total-Row-Count` header (`total_row_count`) and the `rel="next"` pagination `Link` (`next_link`). To run a query and get its result as Arrow in a single call — submit, await `ready`, and decode — use `query_to_arrow`: diff --git a/src/arrow.rs b/src/arrow.rs index 0cccc92..3a1a56b 100644 --- a/src/arrow.rs +++ b/src/arrow.rs @@ -13,7 +13,7 @@ //! adds `Accept: application/vnd.apache.arrow.stream` plus `?format=arrow`, and //! decodes the resulting IPC stream with `arrow-ipc`. //! -//! Two entry points are provided: +//! Three entry points are provided: //! //! * [`get_result_arrow`] — buffers the full IPC stream and returns all //! [`RecordBatch`]es (the Rust analog of pyarrow `Table`; Rust has no @@ -22,6 +22,11 @@ //! [`RecordBatch`] at a time, mirroring pyarrow's //! `RecordBatchStreamReader`. The body is still collected once (reqwest's //! async body is not a blocking `Read`); decoding is then lazy per batch. +//! * [`open_result_arrow`] — returns an [`ArrowResultStream`] that decodes +//! straight off the socket, pulling body chunks as batches are asked for. +//! Peak memory is one record batch rather than the whole result, which makes +//! it the entry point for a result larger than memory. The trade is that the +//! pooled connection stays checked out until the stream is drained. //! //! Enable with the `arrow` cargo feature (mirrors Python's `[arrow]` extra): //! @@ -33,7 +38,8 @@ use std::fmt; use std::io::Cursor; use arrow_array::RecordBatch; -use arrow_ipc::reader::StreamReader; +use arrow_buffer::Buffer; +use arrow_ipc::reader::{StreamDecoder, StreamReader}; use arrow_schema::{ArrowError as IpcArrowError, SchemaRef}; use bytes::Bytes; @@ -297,6 +303,228 @@ pub async fn stream_result_arrow( }) } +/// A [`RecordBatch`] stream decoded straight off the HTTP response body. +/// +/// The difference from [`stream_result_arrow`] is where the bytes live. That one +/// collects the whole body first and then decodes lazily, so peak memory scales +/// with the size of the result. This one pulls body chunks from the socket as +/// batches are asked for, so peak memory is bounded by a single record batch — +/// which is what makes a result larger than memory readable at all. +/// +/// The cost is that the pooled connection stays checked out until the stream is +/// drained or dropped. Dropping early abandons the download; the endpoint is +/// streamed end-to-end server-side, so the server stops producing when the +/// client goes away. +/// +/// The schema is resolved during [`open_result_arrow`], so it is available +/// before the first batch — a caller writing a header row does not have to read +/// data to learn the column names. +#[derive(Debug)] +pub struct ArrowResultStream { + source: ChunkSource, + decoder: StreamDecoder, + /// Bytes pulled from the socket that the decoder has not consumed yet. + /// + /// Empty whenever [`StreamDecoder::decode`] last returned `None`: it + /// consumes the buffer it is given before asking for more, and only returns + /// early — with bytes left over — when it has a batch to hand back. + pending: Buffer, + /// A batch the schema read in [`open_result_arrow`] decoded as a side + /// effect, handed out by the first [`next_batch`](Self::next_batch) call so + /// it is not dropped. + buffered: Option, + schema: SchemaRef, + total_row_count: Option, + next_link: Option, + done: bool, +} + +impl ArrowResultStream { + /// The schema of the result, known before any batch is read. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + /// `X-Total-Row-Count`: rows in the full result, ignoring offset/limit. + pub fn total_row_count(&self) -> Option { + self.total_row_count + } + + /// The `rel="next"` `Link` URL, when a finite `limit` did not reach the end. + pub fn next_link(&self) -> Option<&str> { + self.next_link.as_deref() + } + + /// Decode the next [`RecordBatch`], pulling more body chunks as needed. + /// + /// Returns `Ok(None)` once the stream is exhausted. A body that ends in the + /// middle of an IPC message is an error ([`StreamDecoder::finish`]), not a + /// silent short read — a truncated download must not look like a complete + /// result. + pub async fn next_batch(&mut self) -> Result, ArrowError> { + if let Some(batch) = self.buffered.take() { + return Ok(Some(batch)); + } + if self.done { + return Ok(None); + } + loop { + if !self.pending.is_empty() { + if let Some(batch) = self.decoder.decode(&mut self.pending)? { + return Ok(Some(batch)); + } + } + match self.source.next_chunk().await? { + // `pending` is empty here (see the field comment), so this + // replaces rather than discards. `Buffer::from(Bytes)` is a + // refcount bump, not a copy. + Some(bytes) => self.pending = Buffer::from(bytes), + None => { + // `finish` first: a body that ended mid-message must keep + // erroring. Setting `done` before it would make a second + // call return `Ok(None)`, so a caller that logged the error + // and read on would see a truncated download as a clean end + // of stream. An exhausted body keeps yielding `None`, so the + // loop reaches `finish` again on every later call. + self.decoder.finish()?; + self.done = true; + return Ok(None); + } + } + } + } + + /// Drain the remaining batches into a single [`ArrowResult`]. + /// + /// Materializes the whole result, so it defeats the point of this reader for + /// a large one; it exists for callers that streamed only to avoid the double + /// buffering of [`get_result_arrow`] (body bytes plus decoded batches). + pub async fn read_all(mut self) -> Result { + let schema = self.schema.clone(); + let total_row_count = self.total_row_count; + let next_link = self.next_link.clone(); + let mut batches = Vec::new(); + while let Some(batch) = self.next_batch().await? { + batches.push(batch); + } + Ok(ArrowResult { + batches, + schema, + total_row_count, + next_link, + }) + } +} + +/// Open a ready result as an [`ArrowResultStream`] that decodes off the socket. +/// +/// Use this for a result too large to hold in memory: peak memory is one record +/// batch, where [`get_result_arrow`] and [`stream_result_arrow`] both scale with +/// the full result size. +/// +/// The request is built by the same code as [`get_result_arrow`], so the two are +/// identical on the wire. The schema is read before returning, which costs one +/// body chunk. +/// +/// # Errors +/// +/// Same status mapping as [`get_result_arrow`], plus [`ArrowError::Ipc`] if the +/// body ends before a schema message arrives. +pub async fn open_result_arrow( + configuration: &Configuration, + id: &str, + x_database_id: &str, + offset: Option, + limit: Option, +) -> Result { + let req = build_result_request(configuration, id, x_database_id, offset, limit).await?; + crate::http_log::log_request(&req); + let resp = configuration.client.execute(req).await?; + let status = resp.status(); + crate::http_log::log_response_status(status); + + if status != reqwest::StatusCode::OK { + return Err(map_error_response(id, resp).await); + } + + let total_row_count = parse_total_row_count(&resp); + let next_link = parse_next_link(&resp); + + open_from_source(ChunkSource::Body(resp), total_row_count, next_link).await +} + +/// Where an [`ArrowResultStream`] pulls its bytes from. +/// +/// The body variant is the only one in production. The test variant exists so +/// the pump below — the part that has to survive a chunk boundary landing inside +/// an IPC message — can be driven with chosen splits instead of whatever +/// chunking a socket happens to produce. +#[derive(Debug)] +enum ChunkSource { + Body(reqwest::Response), + #[cfg(test)] + Chunks(std::vec::IntoIter), +} + +impl ChunkSource { + async fn next_chunk(&mut self) -> Result, ArrowError> { + match self { + ChunkSource::Body(resp) => Ok(resp.chunk().await?), + #[cfg(test)] + ChunkSource::Chunks(chunks) => Ok(chunks.next()), + } + } +} + +/// Read the schema off `source`, then hand back a stream positioned to decode +/// batches. +/// +/// The schema is the first thing an Arrow IPC stream carries, so this normally +/// consumes one chunk — but a chunk boundary can fall inside the schema message, +/// and a single chunk can also carry the first batch along with it, which is why +/// a decoded batch is stashed rather than dropped. +async fn open_from_source( + mut source: ChunkSource, + total_row_count: Option, + next_link: Option, +) -> Result { + let mut decoder = StreamDecoder::new(); + let mut pending = Buffer::from(Bytes::new()); + let mut buffered = None; + let schema = loop { + if !pending.is_empty() { + if let Some(batch) = decoder.decode(&mut pending)? { + buffered = Some(batch); + } + } + if let Some(schema) = decoder.schema() { + break schema; + } + match source.next_chunk().await? { + Some(bytes) => pending = Buffer::from(bytes), + None => { + // An empty 200 body, or one cut short before the schema. Report + // it rather than presenting a zero-column result as complete. + decoder.finish()?; + return Err(ArrowError::Ipc(IpcArrowError::IpcError( + "result body ended before an Arrow schema message".to_string(), + ))); + } + } + }; + + Ok(ArrowResultStream { + source, + decoder, + pending, + buffered, + schema, + total_row_count, + next_link, + done: false, + }) +} + /// Apply the `X-Workspace-Id` API-key header, mirroring the generated /// `get_result` `isKeyInHeader` block so the Arrow path is scoped identically. fn apply_apikey_headers( @@ -327,6 +555,34 @@ async fn fetch_arrow_bytes( offset: Option, limit: Option, ) -> Result<(Bytes, Option, Option), ArrowError> { + let req = build_result_request(configuration, id, x_database_id, offset, limit).await?; + crate::http_log::log_request(&req); + let resp = configuration.client.execute(req).await?; + let status = resp.status(); + crate::http_log::log_response_status(status); + + if status == reqwest::StatusCode::OK { + let total_row_count = parse_total_row_count(&resp); + let next_link = parse_next_link(&resp); + let bytes = resp.bytes().await?; + return Ok((bytes, total_row_count, next_link)); + } + + Err(map_error_response(id, resp).await) +} + +/// Build the `GET /v1/results/{id}?format=arrow` request. +/// +/// Shared by every entry point in this module so the buffered and streaming +/// paths are byte-identical on the wire: same URL, query parameters, API-key +/// header, user-agent, `Accept`, and the bearer credential resolved per request. +async fn build_result_request( + configuration: &Configuration, + id: &str, + x_database_id: &str, + offset: Option, + limit: Option, +) -> Result { let uri_str = format!( "{}/v1/results/{id}", configuration.base_path, @@ -362,55 +618,60 @@ async fn fetch_arrow_bytes( req_builder = req_builder.header(reqwest::header::ACCEPT, ARROW_STREAM_MEDIA_TYPE); - let req = req_builder.build()?; - crate::http_log::log_request(&req); - let resp = configuration.client.execute(req).await?; - let status = resp.status(); - crate::http_log::log_response_status(status); - - if status == reqwest::StatusCode::OK { - let total_row_count = parse_total_row_count(&resp); - let next_link = parse_next_link(&resp); - let bytes = resp.bytes().await?; - return Ok((bytes, total_row_count, next_link)); - } + Ok(req_builder.build()?) +} +/// Map a non-200 `GET /v1/results/{id}` response to an [`ArrowError`]. +/// +/// Every path drains the body, both to produce the message and so the pooled +/// connection is returned rather than left for the pool to reclaim. +async fn map_error_response(id: &str, resp: reqwest::Response) -> ArrowError { + let status = resp.status(); match status { reqwest::StatusCode::ACCEPTED => { let retry_after = parse_retry_after(&resp); - let body = resp.text().await?; + let body = match resp.text().await { + Ok(body) => body, + Err(e) => return ArrowError::Reqwest(e), + }; crate::http_log::log_response_body(&body); let (result_status, result_id) = parse_status_and_id(&body, id); - Err(ArrowError::NotReady { + ArrowError::NotReady { status: result_status, result_id, retry_after, - }) + } } reqwest::StatusCode::CONFLICT => { - let body = resp.text().await?; + let body = match resp.text().await { + Ok(body) => body, + Err(e) => return ArrowError::Reqwest(e), + }; crate::http_log::log_response_body(&body); let error_message = parse_error_message(&body); - Err(ArrowError::Failed { error_message }) + ArrowError::Failed { error_message } } reqwest::StatusCode::NOT_FOUND => { // Drain the body so the connection is returned to the pool. let body = resp.text().await.unwrap_or_default(); crate::http_log::log_response_body(&body); - Err(ArrowError::NotFound) + ArrowError::NotFound } reqwest::StatusCode::BAD_REQUEST => { - let message = resp.text().await?; + let message = match resp.text().await { + Ok(message) => message, + Err(e) => return ArrowError::Reqwest(e), + }; crate::http_log::log_response_body(&message); - Err(ArrowError::InvalidParams { message }) + ArrowError::InvalidParams { message } } other => { let body = resp.text().await.unwrap_or_default(); crate::http_log::log_response_body(&body); - Err(ArrowError::Http { + ArrowError::Http { status: other, body, - }) + } } } } @@ -688,6 +949,292 @@ mod tests { assert!(matches!(err, Some(ArrowError::Ipc(_)))); } + // --- streaming reader (open_result_arrow / ArrowResultStream) ----------- + + /// Open a stream over `chunks` using the production pump. + async fn open_chunks(chunks: Vec) -> Result { + open_from_source(ChunkSource::Chunks(chunks.into_iter()), None, None).await + } + + /// Drain a stream into (batches, total rows). + async fn drain(mut stream: ArrowResultStream) -> Vec { + let mut batches = Vec::new(); + while let Some(batch) = stream.next_batch().await.expect("decode should succeed") { + batches.push(batch); + } + batches + } + + /// The whole point: the streaming reader must produce exactly what the + /// buffered reader produces. Anything else is a silent difference between + /// two ways of reading the same result. + #[tokio::test] + async fn streaming_reader_matches_the_buffered_reader() { + let (ipc, schema) = make_ipc_stream(); + + let buffered = StreamReader::try_new(Cursor::new(Bytes::from(ipc.clone())), None) + .unwrap() + .collect::, _>>() + .unwrap(); + + let streamed = drain(open_chunks(vec![Bytes::from(ipc)]).await.unwrap()).await; + + assert_eq!(streamed, buffered); + assert_eq!(streamed.iter().map(|b| b.num_rows()).sum::(), 5); + assert_eq!(streamed[0].schema(), schema); + } + + /// A chunk boundary can land anywhere — inside the schema message, inside a + /// batch header, inside array data. One-byte chunks put a boundary at every + /// possible position at once, which is the adversarial case a socket only + /// produces occasionally. + #[tokio::test] + async fn streaming_reader_survives_one_byte_chunks() { + let (ipc, _schema) = make_ipc_stream(); + let expected = StreamReader::try_new(Cursor::new(Bytes::from(ipc.clone())), None) + .unwrap() + .collect::, _>>() + .unwrap(); + + let chunks: Vec = ipc.iter().map(|b| Bytes::copy_from_slice(&[*b])).collect(); + let streamed = drain(open_chunks(chunks).await.unwrap()).await; + + assert_eq!(streamed, expected); + } + + /// Splitting at every offset in turn, so no single boundary position is + /// left untested by luck of where the one-byte case happens to succeed. + #[tokio::test] + async fn streaming_reader_survives_a_split_at_every_offset() { + let (ipc, _schema) = make_ipc_stream(); + let expected = StreamReader::try_new(Cursor::new(Bytes::from(ipc.clone())), None) + .unwrap() + .collect::, _>>() + .unwrap(); + + for split in 1..ipc.len() { + let chunks = vec![ + Bytes::copy_from_slice(&ipc[..split]), + Bytes::copy_from_slice(&ipc[split..]), + ]; + let streamed = drain(open_chunks(chunks).await.unwrap()).await; + assert_eq!(streamed, expected, "mismatch when split at byte {split}"); + } + } + + /// A caller writing a CSV header needs the column names before it has read + /// any data. `open` resolves the schema, so this holds even for a result + /// whose rows never arrive. + #[tokio::test] + async fn streaming_reader_knows_the_schema_before_any_batch() { + let (ipc, schema) = make_ipc_stream(); + let stream = open_chunks(vec![Bytes::from(ipc)]).await.unwrap(); + assert_eq!(stream.schema(), schema); + } + + #[tokio::test] + async fn streaming_reader_yields_zero_batches_for_an_empty_result() { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let mut buf: Vec = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, &schema).unwrap(); + writer.finish().unwrap(); + } + + let stream = open_chunks(vec![Bytes::from(buf)]).await.unwrap(); + assert_eq!( + stream.schema(), + schema, + "header info survives an empty result" + ); + assert!(drain(stream).await.is_empty()); + } + + /// A download cut short must not read as a complete result. This is the + /// failure the buffered path got for free (a short body fails to parse) and + /// that a streaming reader has to assert deliberately. + #[tokio::test] + async fn a_body_cut_short_is_an_error_not_a_short_read() { + let (ipc, _schema) = make_ipc_stream(); + // Drop the trailing end-of-stream marker and part of the last batch. + let truncated = &ipc[..ipc.len() - 16]; + + let mut stream = open_chunks(vec![Bytes::copy_from_slice(truncated)]) + .await + .expect("the schema is intact, so opening succeeds"); + + let mut err = None; + loop { + match stream.next_batch().await { + Ok(Some(_)) => continue, + Ok(None) => break, + Err(e) => { + err = Some(e); + break; + } + } + } + assert!( + matches!(err, Some(ArrowError::Ipc(_))), + "a truncated body must surface as an IPC error, got {err:?}" + ); + } + + /// A truncated body must keep failing. If the stream marked itself done + /// before confirming a clean end, a caller that logged the first error and + /// read on would be told the stream ended normally — a short download + /// silently becoming a complete result. + /// `read_all` must carry every batch *and* the metadata headers into the + /// returned `ArrowResult`. A dropped field here would hand back + /// metadata-free results on every call. + #[tokio::test] + async fn streaming_reader_read_all_carries_batches_and_metadata() { + let (ipc, schema) = make_ipc_stream(); + let expected = StreamReader::try_new(Cursor::new(Bytes::from(ipc.clone())), None) + .unwrap() + .collect::, _>>() + .unwrap(); + + let stream = open_from_source( + ChunkSource::Chunks(vec![Bytes::from(ipc)].into_iter()), + Some(5), + Some("https://api.hotdata.dev/v1/results/abc?offset=5".to_string()), + ) + .await + .unwrap(); + + let result = stream.read_all().await.expect("draining should succeed"); + assert_eq!( + result.batches, expected, + "every batch must reach the result" + ); + assert_eq!(result.num_rows(), 5); + assert_eq!(result.schema, schema); + assert_eq!(result.total_row_count, Some(5), "X-Total-Row-Count dropped"); + assert_eq!( + result.next_link.as_deref(), + Some("https://api.hotdata.dev/v1/results/abc?offset=5"), + "the rel=next Link dropped" + ); + } + + #[tokio::test] + async fn a_cut_short_body_keeps_erroring_on_every_later_call() { + let (ipc, _schema) = make_ipc_stream(); + let truncated = &ipc[..ipc.len() - 16]; + + let mut stream = open_chunks(vec![Bytes::copy_from_slice(truncated)]) + .await + .expect("the schema is intact, so opening succeeds"); + + // Drain to the failure. + let first_err = loop { + match stream.next_batch().await { + Ok(Some(_)) => continue, + Ok(None) => panic!("a truncated body must not report a clean end of stream"), + Err(e) => break e, + } + }; + assert!(matches!(first_err, ArrowError::Ipc(_))); + + // Reading on must report the same failure, not `Ok(None)`. + for attempt in 0..3 { + match stream.next_batch().await { + Err(ArrowError::Ipc(_)) => {} + other => panic!( + "call {attempt} after a truncated body must repeat the error, got {other:?}" + ), + } + } + } + + #[tokio::test] + async fn a_body_with_no_schema_message_is_an_error() { + let err = open_chunks(vec![]).await.err(); + assert!( + matches!(err, Some(ArrowError::Ipc(_))), + "an empty body must not open as a zero-column result, got {err:?}" + ); + } + + /// End-to-end over HTTP: the streaming opener forwards the same headers and + /// query parameters as the buffered fetch, and decodes the same rows. + #[tokio::test] + async fn open_result_arrow_forwards_headers_and_decodes() { + use wiremock::matchers::{header, method, path, query_param}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let (ipc, _schema) = make_ipc_stream(); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/results/res_1")) + .and(query_param("format", "arrow")) + .and(header("X-Database-Id", "db_x")) + .and(header("accept", ARROW_STREAM_MEDIA_TYPE)) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", ARROW_STREAM_MEDIA_TYPE) + .insert_header("X-Total-Row-Count", "5") + .set_body_bytes(ipc), + ) + .mount(&server) + .await; + + let mut configuration = Configuration::new(); + configuration.base_path = server.uri(); + + let stream = open_result_arrow(&configuration, "res_1", "db_x", None, None) + .await + .expect("streaming open should succeed"); + assert_eq!(stream.total_row_count(), Some(5)); + assert_eq!( + drain(stream) + .await + .iter() + .map(|b| b.num_rows()) + .sum::(), + 5 + ); + } + + /// The non-200 mapping moved into a shared helper during this change; both + /// entry points must still report a not-ready result as such. + #[tokio::test] + async fn open_result_arrow_maps_a_not_ready_result() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/results/res_pending")) + .respond_with( + ResponseTemplate::new(202) + .insert_header("Retry-After", "2") + .set_body_string(r#"{"status":"processing","result_id":"res_pending"}"#), + ) + .mount(&server) + .await; + + let mut configuration = Configuration::new(); + configuration.base_path = server.uri(); + + let err = open_result_arrow(&configuration, "res_pending", "db_x", None, None) + .await + .err(); + match err { + Some(ArrowError::NotReady { + status, + result_id, + retry_after, + }) => { + assert_eq!(status, "processing"); + assert_eq!(result_id, "res_pending"); + assert_eq!(retry_after, Some(2)); + } + other => panic!("expected NotReady, got {other:?}"), + } + } + #[test] fn link_header_parses_rel_next() { let h = "; rel=\"next\""; diff --git a/src/client.rs b/src/client.rs index 8a49110..6081d7f 100644 --- a/src/client.rs +++ b/src/client.rs @@ -561,6 +561,24 @@ impl Client { crate::arrow::stream_result_arrow(&self.configuration, id, database_id, offset, limit).await } + /// Open a result as an Arrow IPC stream decoded off the socket. + /// + /// Peak memory is one record batch rather than the whole result, so this is + /// the path for a result larger than memory. See + /// [`crate::arrow::open_result_arrow`] for the connection trade-off. + /// + /// Requires the `arrow` cargo feature. + #[cfg(feature = "arrow")] + pub async fn open_result_arrow( + &self, + id: &str, + database_id: &str, + offset: Option, + limit: Option, + ) -> Result { + crate::arrow::open_result_arrow(&self.configuration, id, database_id, offset, limit).await + } + // --- Resource handles ----------------------------------------------------- // // Grouped, ergonomic accessors over the generated `apis::*_api` free diff --git a/src/lib.rs b/src/lib.rs index 34b1870..73553ae 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -36,8 +36,8 @@ pub use apis::configuration::{ApiKey, BasicAuth, Configuration}; pub use apis::Error; #[cfg(feature = "arrow")] pub use arrow::{ - get_result_arrow, stream_result_arrow, ArrowBatchStream, ArrowError, ArrowResult, - ARROW_STREAM_MEDIA_TYPE, + get_result_arrow, open_result_arrow, stream_result_arrow, ArrowBatchStream, ArrowError, + ArrowResult, ArrowResultStream, ARROW_STREAM_MEDIA_TYPE, }; pub use auth::{BearerTokenError, BearerTokenProvider}; #[cfg(feature = "arrow")] diff --git a/src/query.rs b/src/query.rs index e979b01..aaeb9bf 100644 --- a/src/query.rs +++ b/src/query.rs @@ -274,9 +274,13 @@ pub enum ResultError { deadline: Duration, }, /// Auto-follow would materialize more than the guard allows, on either axis. - /// Stream the result instead via - /// [`Client::stream_result_arrow`](crate::Client::stream_result_arrow), or - /// raise (or set to `None`) the relevant guard. + /// Read the result instead via + /// [`Client::open_result_arrow`](crate::Client::open_result_arrow), which + /// decodes off the socket and holds one record batch at a time, or raise + /// (or set to `None`) the relevant guard. + /// + /// Not `stream_result_arrow`: that one collects the whole body before + /// decoding, so it would materialize exactly what this guard refused. TooLarge { /// The result id that exceeded the guard. result_id: String, @@ -338,7 +342,7 @@ impl std::fmt::Display for ResultError { write!( f, "result {result_id} exceeds the auto-materialize limit: {desc}. \ - Stream it with Client::stream_result_arrow, or raise (or set to \ + Read it with Client::open_result_arrow, or raise (or set to \ None) {}.", kind.knob() )