From 7dfdfe64019485f7ee625394d007b160fcb9cb42 Mon Sep 17 00:00:00 2001 From: sanskar singh bhardwaj <95009647+sanskar-singh-2403@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:26:25 +0530 Subject: [PATCH] fix(remote): surface masked merge_insert stream errors under HTTP2 (#2339) (#3711) ## Summary Fixes #2339. `merge_insert()` on the remote client could mask the real cause of a mid-stream input error, reporting only: > stream error sent by user: unexpected internal error ## Root cause There were two divergent streaming-write code paths in the remote client: - `add()` uses `RemoteInsertExec`, which streams the request body through a `tokio::sync::oneshot` error side-channel and drains it before reporting the HTTP result. If the input stream errors mid-body, the original error is recovered. - `merge_insert()` used a legacy path (`send_streaming` -> `reader_as_body`) that piped arrow `Some(Err(e))` straight into the HTTP2 request body. Hyper swallows body-stream errors under HTTP2 (see hyperium/hyper#2547), so the original error was lost and only the generic transport error surfaced. ## Fix Consolidate both write paths onto the side-channel mechanism instead of patching the legacy path: - Generalize `RemoteInsertExec` into `RemoteWriteExec`, carrying a `WriteOp` enum (`Insert { overwrite }` | `MergeInsert { query, timeout }`) that selects the endpoint, query params, request-timeout header, and response parsing. The executor returns a `WriteResult` enum (`Add` | `Merge`) with typed accessors, and `with_new_children` still resets the result so the rescannable retry loop is unaffected. - Route `merge_insert()` through `RemoteWriteExec`. The public API only accepts a `RecordBatchReader` (not rescannable), so the reader is buffered into a `Vec` before the retry loop to preserve the previous retry-on-retryable-status behaviour. This mirrors what the old `send_streaming(with_retry=true)` path already did. - Remove the now-unused `send_streaming` / `reader_as_body` / `buffer_reader` / `make_reader` helpers. Multipart stays insert-only (the server has no multipart merge_insert endpoint), so that hot path is behaviorally unchanged. ## Testing - Added `test_merge_insert_input_error_surfaces_original`, which drives an erroring input through the single-request `merge_insert` path and asserts the original error (`boom`) is surfaced rather than the masked HTTP error. Confirmed it fails without the side-channel drain (it then reports a masked `500 ... request or response body error`). - Full suite green: `cargo test -p lancedb --lib --features remote` -> 694 passed, 0 failed. Includes the existing `test_merge_insert_retries_on_409`, confirming retry behaviour is preserved. --- rust/lancedb/src/remote/table.rs | 194 ++++++--------- rust/lancedb/src/remote/table/insert.rs | 307 ++++++++++++++++++------ 2 files changed, 310 insertions(+), 191 deletions(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 685e2b54e..b0d7dacc5 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3,10 +3,13 @@ pub mod insert; -use self::insert::RemoteInsertExec; +use self::insert::{RemoteWriteExec, WriteOp}; use crate::expr::expr_to_sql_string; use crate::table::write_progress::FinishOnDrop; - +// Used by the test module below (single-request write requests set these +// headers directly in test handlers); kept at module scope so both the +// library and its tests can name them. +#[cfg(test)] use super::ARROW_STREAM_CONTENT_TYPE; use super::client::RequestResultExt; use super::client::{HttpSend, RestfulLanceDbClient, Sender}; @@ -44,7 +47,7 @@ use crate::{ merge::MergeInsertBuilder, }, }; -use arrow_array::{RecordBatch, RecordBatchIterator, RecordBatchReader}; +use arrow_array::RecordBatchReader; use arrow_ipc::reader::FileReader; use arrow_schema::{DataType, SchemaRef}; use async_trait::async_trait; @@ -53,6 +56,7 @@ use datafusion_common::DataFusionError; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ExecutionPlan, RecordBatchStream, SendableRecordBatchStream}; use futures::{StreamExt, TryStreamExt}; +#[cfg(test)] use http::header::CONTENT_TYPE; use http::{HeaderName, StatusCode}; use lance::arrow::json::{JsonDataType, JsonSchema}; @@ -455,50 +459,6 @@ impl RemoteTable { } } - fn reader_as_body(data: Box) -> Result { - // TODO: Once Phalanx supports compression, we should use it here. - let mut writer = arrow_ipc::writer::StreamWriter::try_new( - Vec::new(), - &RecordBatchReader::schema(&*data), - )?; - - // Mutex is just here to make it sync. We shouldn't have any contention. - let mut data = Mutex::new(data); - let body_iter = std::iter::from_fn(move || match data.get_mut().unwrap().next() { - Some(Ok(batch)) => { - writer.write(&batch).ok()?; - let buffer = std::mem::take(writer.get_mut()); - Some(Ok(buffer)) - } - Some(Err(e)) => Some(Err(e)), - None => { - writer.finish().ok()?; - let buffer = std::mem::take(writer.get_mut()); - Some(Ok(buffer)) - } - }); - let body_stream = futures::stream::iter(body_iter); - Ok(reqwest::Body::wrap_stream(body_stream)) - } - - /// Buffer the reader into memory - async fn buffer_reader( - reader: &mut R, - ) -> Result<(SchemaRef, Vec)> { - let schema = reader.schema(); - let mut batches = Vec::new(); - for batch in reader { - batches.push(batch?); - } - Ok((schema, batches)) - } - - /// Create a new RecordBatchReader from buffered data - fn make_reader(schema: SchemaRef, batches: Vec) -> impl RecordBatchReader { - let iter = batches.into_iter().map(Ok); - RecordBatchIterator::new(iter, schema) - } - async fn send(&self, req: RequestBuilder, with_retry: bool) -> Result<(String, Response)> { let res = if with_retry { self.client.send_with_retry(req, None, true).await? @@ -508,34 +468,6 @@ impl RemoteTable { Ok(res) } - /// Send the request with streaming body. - /// This will use retries if with_retry is set and the number of configured retries is > 0. - /// If retries are enabled, the stream will be buffered into memory. - async fn send_streaming( - &self, - req: RequestBuilder, - mut data: Box, - with_retry: bool, - ) -> Result<(String, Response)> { - if !with_retry || self.client.retry_config.retries == 0 { - let body = Self::reader_as_body(data)?; - return self.client.send(req.body(body)).await; - } - - // to support retries, buffer into memory and clone the batches on each retry - let (schema, batches) = Self::buffer_reader(&mut *data).await?; - let make_body = Box::new(move || { - let reader = Self::make_reader(schema.clone(), batches.clone()); - Self::reader_as_body(Box::new(reader)) - }); - let res = self - .client - .send_with_retry(req, Some(make_body), false) - .await?; - - Ok(res) - } - pub(super) async fn handle_table_not_found( table_name: &str, response: reqwest::Response, @@ -1238,12 +1170,14 @@ impl RemoteTable { let _guard = output.tracker.as_ref().map(|t| t.track_task()); - let mut insert: Arc = Arc::new(RemoteInsertExec::new( + let mut insert: Arc = Arc::new(RemoteWriteExec::new( self.name.clone(), self.identifier.clone(), self.client.clone(), output.plan, - output.overwrite, + WriteOp::Insert { + overwrite: output.overwrite, + }, output.tracker.clone(), self.branch.clone(), )); @@ -1258,7 +1192,7 @@ impl RemoteTable { match result { Ok(_) => { let add_result = (insert.as_ref() as &dyn std::any::Any) - .downcast_ref::>() + .downcast_ref::>() .and_then(|i| i.add_result()) .unwrap_or(AddResult { version: 0 }); @@ -1359,7 +1293,7 @@ impl RemoteTable { )?, ) as Arc; - let insert = Arc::new(RemoteInsertExec::new_multipart( + let insert = Arc::new(RemoteWriteExec::new_multipart( self.name.clone(), self.identifier.clone(), self.client.clone(), @@ -2354,48 +2288,66 @@ impl BaseTable for RemoteTable { self.check_mutable().await?; let timeout = params.timeout; - let query = MergeInsertRequest::try_from(params)?; - let mut request = self - .client - .post(&format!("/v1/table/{}/merge_insert/", self.identifier)) - .query(&query) - .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE); - request = self.apply_branch_query(request); - if let Some(timeout) = timeout { - // (If it doesn't fit into u64, it's not worth sending anyways.) - if let Ok(timeout_ms) = u64::try_from(timeout.as_millis()) { - request = request.header(REQUEST_TIMEOUT_HEADER, timeout_ms); + // Drive merge_insert through the same RemoteWriteExec streaming path as + // add(). This routes the request body through the error side-channel so + // an input stream error (e.g. NaN rejection) surfaces with its original + // message instead of the masked HTTP error Hyper produces when a request + // body stream fails under HTTP2 (issue #2339). The branch, request + // timeout header, and merge query params are all applied inside the exec. + // + // The public merge_insert API only accepts a `RecordBatchReader`, which + // is not rescannable and so could not be retried directly. To preserve + // the previous retry-on-retryable-status behaviour, buffer the reader + // into memory first: a `Vec` is rescannable, so the outer + // loop can re-execute the plan (and re-stream the body) on each retry. + // This mirrors the old `send_streaming(with_retry=true)` path, which + // likewise buffered the reader to support retries. + let batches = new_data.collect::, _>>()?; + let source: Box = Box::new(batches); + let rescannable = source.rescannable(); + let input: Arc = + Arc::new(crate::table::datafusion::scannable_exec::ScannableExec::new(source, None)); + + let mut merge: Arc = Arc::new(RemoteWriteExec::new( + self.name.clone(), + self.identifier.clone(), + self.client.clone(), + input, + WriteOp::MergeInsert { query, timeout }, + None, + self.branch.clone(), + )); + + let mut retry_counter = crate::remote::retry::RetryCounter::new( + &self.client.retry_config, + uuid::Uuid::new_v4().to_string(), + ); + + loop { + let stream = execute_plan(merge.clone(), Default::default())?; + let result: Result> = stream.try_collect().await.map_err(Error::from); + + match result { + Ok(_) => { + let merge_result = (merge.as_ref() as &dyn std::any::Any) + .downcast_ref::>() + .and_then(|m| m.merge_result()) + .unwrap_or_default(); + + self.track_write_version(merge_result.version); + return Ok(merge_result); + } + Err(err) if rescannable && self.is_retryable_write_error(&err) => { + retry_counter.increment_from_error(err)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + merge = merge.reset_state()?; + continue; + } + Err(err) => return Err(err), } } - - let (request_id, response) = self.send_streaming(request, new_data, true).await?; - - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - if body.trim().is_empty() { - // Backward compatible with old servers - return Ok(MergeResult { - version: 0, - num_deleted_rows: 0, - num_inserted_rows: 0, - num_updated_rows: 0, - num_attempts: 0, - num_rows: 0, - }); - } - - let merge_insert_response: MergeResult = - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse merge_insert response: {}", e).into(), - request_id, - status_code: None, - })?; - - self.track_write_version(merge_insert_response.version); - Ok(merge_insert_response) } async fn set_unenforced_primary_key(&self, _columns: &[&str]) -> Result<()> { @@ -2862,20 +2814,20 @@ impl BaseTable for RemoteTable { write_params: lance::dataset::WriteParams, ) -> Result> { let overwrite = matches!(write_params.mode, lance::dataset::WriteMode::Overwrite); - Ok(Arc::new(insert::RemoteInsertExec::new( + Ok(Arc::new(insert::RemoteWriteExec::new( self.name.clone(), self.identifier.clone(), self.client.clone(), input, - overwrite, + WriteOp::Insert { overwrite }, None, self.branch.clone(), ))) } } -#[derive(Serialize)] -struct MergeInsertRequest { +#[derive(Serialize, Clone, Debug)] +pub(crate) struct MergeInsertRequest { on: String, when_matched_update_all: bool, when_matched_update_all_filt: Option, diff --git a/rust/lancedb/src/remote/table/insert.rs b/rust/lancedb/src/remote/table/insert.rs index 4b0aa5480..67ea7765d 100644 --- a/rust/lancedb/src/remote/table/insert.rs +++ b/rust/lancedb/src/remote/table/insert.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! DataFusion ExecutionPlan for inserting data into remote LanceDB tables. +//! DataFusion ExecutionPlan for streaming writes (add / merge_insert) to +//! remote LanceDB tables. use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -23,28 +24,57 @@ use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter; use crate::Error; use crate::remote::ARROW_STREAM_CONTENT_TYPE; use crate::remote::client::{HttpSend, RestfulLanceDbClient, Sender}; -use crate::remote::table::RemoteTable; -use crate::table::AddResult; +use crate::remote::table::{MergeInsertRequest, REQUEST_TIMEOUT_HEADER, RemoteTable}; use crate::table::datafusion::insert::COUNT_SCHEMA; use crate::table::write_progress::WriteProgressTracker; +use crate::table::{AddResult, MergeResult}; -/// ExecutionPlan for inserting data into a remote LanceDB table. +/// The write operation a [`RemoteWriteExec`] performs. Both variants share the +/// same Arrow-IPC streaming body and error side-channel; only the target +/// endpoint, query parameters, and parsed result type differ. +#[derive(Debug, Clone)] +pub(crate) enum WriteOp { + /// `add`: stream to `/v1/table/{id}/insert/`, optionally overwriting. + Insert { overwrite: bool }, + /// `merge_insert`: stream to `/v1/table/{id}/merge_insert/` with the merge + /// parameters carried as query params. Multipart is not supported for this + /// operation (the server has no multipart merge_insert endpoint), so an + /// `upload_id` combined with this op is a programming error. + MergeInsert { + query: MergeInsertRequest, + timeout: Option, + }, +} + +/// The parsed server response for a completed write, discriminated by the +/// operation that produced it. +#[derive(Debug, Clone)] +pub(crate) enum WriteResult { + Add(AddResult), + Merge(MergeResult), +} + +/// ExecutionPlan for streaming a write (add or merge_insert) to a remote +/// LanceDB table. /// -/// Streams data as Arrow IPC to `/v1/table/{id}/insert/` endpoint. +/// Streams data as Arrow IPC to the endpoint selected by [`WriteOp`]. Both +/// operations reuse the same error side-channel so an input stream error (e.g. +/// NaN rejection) surfaces with its original message rather than the masked +/// HTTP error Hyper produces when a request body stream fails under HTTP2. /// /// When `upload_id` is set, inserts are staged as part of a multipart write /// session and the plan supports multiple partitions for parallel uploads. /// Without `upload_id`, the plan requires a single partition and commits -/// immediately. +/// immediately. Multipart applies to `add` only. #[derive(Debug)] -pub struct RemoteInsertExec { +pub struct RemoteWriteExec { table_name: String, identifier: String, client: RestfulLanceDbClient, input: Arc, - overwrite: bool, + op: WriteOp, properties: Arc, - add_result: Arc>>, + result: Arc>>, metrics: ExecutionPlanMetricsSet, upload_id: Option, tracker: Option>, @@ -61,27 +91,28 @@ pub struct RemoteInsertExec { max_request_duration: Option, } -impl RemoteInsertExec { - /// Create a new single-partition RemoteInsertExec. +impl RemoteWriteExec { + /// Create a new single-partition RemoteWriteExec. pub fn new( table_name: String, identifier: String, client: RestfulLanceDbClient, input: Arc, - overwrite: bool, + op: WriteOp, tracker: Option>, branch: Option, ) -> Self { Self::new_inner( - table_name, identifier, client, input, overwrite, None, tracker, branch, None, None, + table_name, identifier, client, input, op, None, tracker, branch, None, None, ) } - /// Create a multi-partition RemoteInsertExec for use with multipart writes. + /// Create a multi-partition RemoteWriteExec for use with multipart writes. /// /// Each partition's insert is staged under the given `upload_id` without /// committing. The caller is responsible for calling the complete (or abort) - /// endpoint after all partitions finish. + /// endpoint after all partitions finish. Multipart is insert-only, so the + /// op is fixed to [`WriteOp::Insert`]. #[allow(clippy::too_many_arguments)] pub fn new_multipart( table_name: String, @@ -100,7 +131,7 @@ impl RemoteInsertExec { identifier, client, input, - overwrite, + WriteOp::Insert { overwrite }, Some(upload_id), tracker, branch, @@ -115,7 +146,7 @@ impl RemoteInsertExec { identifier: String, client: RestfulLanceDbClient, input: Arc, - overwrite: bool, + op: WriteOp, upload_id: Option, tracker: Option>, branch: Option, @@ -140,9 +171,9 @@ impl RemoteInsertExec { identifier, client, input, - overwrite, + op, properties: Arc::new(properties), - add_result: Arc::new(Mutex::new(None)), + result: Arc::new(Mutex::new(None)), metrics: ExecutionPlanMetricsSet::new(), upload_id, tracker, @@ -152,14 +183,30 @@ impl RemoteInsertExec { } } - /// Get the add result after execution. - // TODO: this will be used when we wire this up to Table::add(). - #[allow(dead_code)] + /// Get the add result after execution, if this exec ran an insert. pub fn add_result(&self) -> Option { - self.add_result + match self + .result .lock() .unwrap_or_else(|e| e.into_inner()) .clone() + { + Some(WriteResult::Add(r)) => Some(r), + _ => None, + } + } + + /// Get the merge result after execution, if this exec ran a merge_insert. + pub fn merge_result(&self) -> Option { + match self + .result + .lock() + .unwrap_or_else(|e| e.into_inner()) + .clone() + { + Some(WriteResult::Merge(r)) => Some(r), + _ => None, + } } /// Stream the input into an HTTP body as an Arrow IPC stream, capturing any @@ -464,24 +511,24 @@ impl PartRequestCtx<'_, S> { } } -impl DisplayAs for RemoteInsertExec { +impl DisplayAs for RemoteWriteExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!( - f, - "RemoteInsertExec: table={}, overwrite={}", - self.table_name, self.overwrite - ) + write!(f, "RemoteWriteExec: table={}, op=", self.table_name)?; + match &self.op { + WriteOp::Insert { overwrite } => write!(f, "insert, overwrite={}", overwrite), + WriteOp::MergeInsert { .. } => write!(f, "merge_insert"), + } } DisplayFormatType::TreeRender => { - write!(f, "RemoteInsertExec") + write!(f, "RemoteWriteExec") } } } } -impl ExecutionPlan for RemoteInsertExec { +impl ExecutionPlan for RemoteWriteExec { fn name(&self) -> &str { Self::static_name() } @@ -516,15 +563,18 @@ impl ExecutionPlan for RemoteInsertExec { ) -> DataFusionResult> { if children.len() != 1 { return Err(DataFusionError::Internal( - "RemoteInsertExec requires exactly one child".to_string(), + "RemoteWriteExec requires exactly one child".to_string(), )); } + // Building a fresh exec (with a new, empty `result`) is what makes the + // outer rescannable retry loop work: `reset_state()` clears the captured + // result so a re-execution starts clean. Ok(Arc::new(Self::new_inner( self.table_name.clone(), self.identifier.clone(), self.client.clone(), children[0].clone(), - self.overwrite, + self.op.clone(), self.upload_id.clone(), self.tracker.clone(), self.branch.clone(), @@ -540,11 +590,19 @@ impl ExecutionPlan for RemoteInsertExec { ) -> DataFusionResult { if self.upload_id.is_none() && partition != 0 { return Err(DataFusionError::Internal( - "RemoteInsertExec only supports single partition execution without upload_id" + "RemoteWriteExec only supports single partition execution without upload_id" .to_string(), )); } + // Multipart is insert-only: the server has no multipart merge_insert + // endpoint, so a merge_insert with an upload_id is a programming error. + if self.upload_id.is_some() && matches!(self.op, WriteOp::MergeInsert { .. }) { + return Err(DataFusionError::Internal( + "merge_insert does not support multipart (upload_id) writes".to_string(), + )); + } + let input_stream = self.input.execute(partition, context)?; let input_schema = input_stream.schema(); let input_stream: SendableRecordBatchStream = @@ -556,8 +614,8 @@ impl ExecutionPlan for RemoteInsertExec { )); let client = self.client.clone(); let identifier = self.identifier.clone(); - let overwrite = self.overwrite; - let add_result = self.add_result.clone(); + let op = self.op.clone(); + let result_slot = self.result.clone(); let table_name = self.table_name.clone(); let upload_id = self.upload_id.clone(); let tracker = self.tracker.clone(); @@ -568,10 +626,12 @@ impl ExecutionPlan for RemoteInsertExec { let stream = futures::stream::once(async move { // Multipart writes with a byte budget split the partition into // several bounded, still-streamed requests so no single request - // stays open long enough to hit the client read timeout. + // stays open long enough to hit the client read timeout. This path + // is insert-only (guarded above). if let (Some(upload_id), Some(max_bytes)) = (upload_id.as_deref(), max_bytes_per_request) { + let overwrite = matches!(op, WriteOp::Insert { overwrite: true }); let ctx = PartRequestCtx { client: &client, identifier: &identifier, @@ -592,16 +652,36 @@ impl ExecutionPlan for RemoteInsertExec { )?); } - let mut request = client - .post(&format!("/v1/table/{}/insert/", identifier)) - .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE); + // Build the request for the selected operation. Both endpoints take + // an Arrow-IPC streaming body and reuse the same error side-channel. + let mut request = match &op { + WriteOp::Insert { overwrite } => { + let mut request = client + .post(&format!("/v1/table/{}/insert/", identifier)) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE); + if *overwrite { + request = request.query(&[("mode", "overwrite")]); + } + if let Some(ref uid) = upload_id { + request = request.query(&[("upload_id", uid.as_str())]); + } + request + } + WriteOp::MergeInsert { query, timeout } => { + let mut request = client + .post(&format!("/v1/table/{}/merge_insert/", identifier)) + .query(query) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE); + if let Some(timeout) = timeout { + // (If it doesn't fit into u64, it's not worth sending anyways.) + if let Ok(timeout_ms) = u64::try_from(timeout.as_millis()) { + request = request.header(REQUEST_TIMEOUT_HEADER, timeout_ms); + } + } + request + } + }; - if overwrite { - request = request.query(&[("mode", "overwrite")]); - } - if let Some(ref uid) = upload_id { - request = request.query(&[("upload_id", uid.as_str())]); - } if let Some(ref b) = branch { request = request.query(&[("branch", b.as_str())]); } @@ -635,6 +715,8 @@ impl ExecutionPlan for RemoteInsertExec { // If the request failed due to an input stream error, surface the // original error (e.g. NaN rejection) instead of the HTTP error. + // This is the crux of the #2339 fix: Hyper silently swallows body + // stream errors under HTTP2, so we recover the original here. if let Ok(stream_err) = error_rx.try_recv() { return Err(stream_err); } @@ -642,7 +724,7 @@ impl ExecutionPlan for RemoteInsertExec { let (request_id, response) = result?; // For multipart writes, the staging response is not the final - // version. Only parse AddResult for non-multipart inserts. + // version. Only parse the result for non-multipart writes. if upload_id.is_none() { let body_text = response.text().await.map_err(|e| { DataFusionError::External(Box::new(Error::Http { @@ -652,21 +734,44 @@ impl ExecutionPlan for RemoteInsertExec { })) })?; - let parsed_result = if body_text.trim().is_empty() { - // Backward compatible with old servers - AddResult { version: 0 } - } else { - serde_json::from_str(&body_text).map_err(|e| { - DataFusionError::External(Box::new(Error::Http { - source: format!("Failed to parse add response: {}", e).into(), - request_id: request_id.clone(), - status_code: None, - })) - })? + let parsed_result = match &op { + WriteOp::Insert { .. } => { + let add = if body_text.trim().is_empty() { + // Backward compatible with old servers + AddResult { version: 0 } + } else { + serde_json::from_str(&body_text).map_err(|e| { + DataFusionError::External(Box::new(Error::Http { + source: format!("Failed to parse add response: {}", e).into(), + request_id: request_id.clone(), + status_code: None, + })) + })? + }; + WriteResult::Add(add) + } + WriteOp::MergeInsert { .. } => { + let merge = if body_text.trim().is_empty() { + // Backward compatible with old servers + MergeResult::default() + } else { + serde_json::from_str(&body_text).map_err(|e| { + DataFusionError::External(Box::new(Error::Http { + source: format!("Failed to parse merge_insert response: {}", e) + .into(), + request_id: request_id.clone(), + status_code: None, + })) + })? + }; + WriteResult::Merge(merge) + } }; - let mut res_lock = add_result.lock().map_err(|_| { - DataFusionError::Execution("Failed to acquire lock for add_result".to_string()) + let mut res_lock = result_slot.lock().map_err(|_| { + DataFusionError::Execution( + "Failed to acquire lock for write result".to_string(), + ) })?; *res_lock = Some(parsed_result); } else { @@ -680,7 +785,7 @@ impl ExecutionPlan for RemoteInsertExec { })?; } - // Return a single batch with count 0 (actual count is tracked in add_result) + // Return a single batch with count 0 (actual count is tracked in result) let count_array: ArrayRef = Arc::new(UInt64Array::from(vec![0u64])); let batch = RecordBatch::try_new(COUNT_SCHEMA.clone(), vec![count_array])?; Ok::<_, DataFusionError>(batch) @@ -711,9 +816,11 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; - use super::RemoteInsertExec; + use super::RemoteWriteExec; + use super::WriteOp; use crate::Table; use crate::remote::ARROW_STREAM_CONTENT_TYPE; + use crate::remote::table::MergeInsertRequest; use crate::table::datafusion::BaseTableAdapter; fn schema_json() -> &'static str { @@ -1028,7 +1135,7 @@ mod tests { let input = input_plan_from_batches(schema, batches).await; // A 1-byte budget forces every batch into its own part. - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1068,7 +1175,7 @@ mod tests { // A large byte budget and no time limit keep the whole partition in a // single part. - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1109,7 +1216,7 @@ mod tests { // A large byte budget but a tiny duration budget: writing and sending // one batch already takes longer than the limit, so each batch is cut // into its own part on the time check rather than the byte check. - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1144,7 +1251,7 @@ mod tests { // write relies on another partition having data to commit. let input = input_plan_from_batches(schema, vec![]).await; - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1184,7 +1291,7 @@ mod tests { let input = input_plan_from_batches(schema, batches).await; // A 1-byte budget forces every batch into its own part. - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1233,7 +1340,7 @@ mod tests { ]; let input = input_plan_from_partitions(schema, partitions).await; - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1267,7 +1374,7 @@ mod tests { // A large byte budget keeps the good batch and the following error in // the same part, exercising the mid-part abort path. let input: Arc = Arc::new(ErroringExec::new()); - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client, @@ -1297,6 +1404,66 @@ mod tests { ); } + #[tokio::test] + async fn test_merge_insert_input_error_surfaces_original() { + // Regression test for #2339 on the single-request merge_insert path. + // When the input stream errors mid-body, Hyper masks it under HTTP2 as a + // generic "stream error sent by user" message. The error side-channel + // must recover and surface the original DataFusion error instead. + use futures::StreamExt; + + let client = crate::remote::client::test_utils::client_with_handler(|request| { + assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); + http::Response::builder() + .status(200) + .body( + r#"{"version": 2, "num_updated_rows": 0, "num_inserted_rows": 0, "num_deleted_rows": 0}"# + .to_string(), + ) + .unwrap() + }); + + let query = MergeInsertRequest { + on: "id".to_string(), + when_matched_update_all: false, + when_matched_update_all_filt: None, + when_not_matched_insert_all: false, + when_not_matched_by_source_delete: false, + when_not_matched_by_source_delete_filt: None, + use_index: true, + use_lsm: None, + }; + + let input: Arc = Arc::new(ErroringExec::new()); + let exec = RemoteWriteExec::new( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + WriteOp::MergeInsert { + query, + timeout: None, + }, + None, + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + let mut err = None; + while let Some(item) = stream.next().await { + if let Err(e) = item { + err = Some(e); + break; + } + } + + let err = err.expect("expected the input stream error to surface"); + assert!( + err.to_string().contains("boom"), + "expected original input error, got: {err}" + ); + } + #[tokio::test] async fn test_multipart_records_progress_within_a_part() { use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker}; @@ -1327,7 +1494,7 @@ mod tests { // A large byte budget keeps all three batches in one part; smooth // progress therefore requires bytes to be reported per chunk rather than // once when the part completes. - let exec = RemoteInsertExec::new_multipart( + let exec = RemoteWriteExec::new_multipart( "my_table".to_string(), "my_table".to_string(), client,