From 7a3709b3768f3d1a496e1baf5832493de25e6c87 Mon Sep 17 00:00:00 2001
From: adenzhou1350 <aden1350@outlook.com>
Date: Fri, 23 Oct 2026 02:45:02 +0800
Subject: [PATCH] fix(remote): preserve request IDs in multipart response
 errors (#4456)

Keep the request ID when a multipart create/complete response has a
missing or incorrectly typed `upload_id`/`version`. These errors
currently use an empty ID, unlike JSON parse errors, so callers lose the
ID needed to correlate the failed request.

Add six cases through `Table.add(...).write_parallelism(2).execute()`,
capturing the actual outgoing `x-request-id`: four missing/wrong-type
cases fail on unchanged main, while two invalid-JSON controls already
pass. After the fix all six pass, and all 15 multipart tests pass.

Native Linux validation (locked dependencies, `ci` profile):
- `cargo test --offline --profile ci --locked -p lancedb
--no-default-features --features remote --lib
remote::table::tests::test_multipart_write`: 15 passed.
- `cargo fmt --all -- --check`: passed.
- `cargo clippy --offline --profile ci --locked -p lancedb
--no-default-features --features remote --lib --tests -- -D warnings`:
passed.

The committed source bytes match the tested source. HTTP endpoints are
mocked; this does not claim a live Cloud or full-repository test run.
---
 rust/lancedb/src/remote/table.rs | 69 ++++++++++++++++++++++++++++++--
 1 file changed, 65 insertions(+), 4 deletions(-)

diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs
index 9e2641c22..84c1b3fef 100644
--- a/rust/lancedb/src/remote/table.rs
+++ b/rust/lancedb/src/remote/table.rs
@@ -1277,7 +1277,7 @@ impl<S: HttpSend> RemoteTable<S> {
         let body = response.text().await.err_to_http(request_id.clone())?;
         let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http {
             source: format!("Failed to parse multipart create response: {}", e).into(),
-            request_id,
+            request_id: request_id.clone(),
             status_code: None,
         })?;
         parsed["upload_id"]
@@ -1285,7 +1285,7 @@ impl<S: HttpSend> RemoteTable<S> {
             .map(|s| s.to_string())
             .ok_or_else(|| Error::Http {
                 source: "Missing upload_id in multipart create response".into(),
-                request_id: String::new(),
+                request_id,
                 status_code: None,
             })
     }
@@ -1304,12 +1304,12 @@ impl<S: HttpSend> RemoteTable<S> {
         let body = response.text().await.err_to_http(request_id.clone())?;
         let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http {
             source: format!("Failed to parse multipart complete response: {}", e).into(),
-            request_id,
+            request_id: request_id.clone(),
             status_code: None,
         })?;
         let version = parsed["version"].as_u64().ok_or_else(|| Error::Http {
             source: "Missing version in multipart complete response".into(),
-            request_id: String::new(),
+            request_id,
             status_code: None,
         })?;
         Ok(AddResult { version })
@@ -12596,6 +12596,67 @@ mod tests {
         );
     }
 
+    #[rstest]
+    #[case("create", "{}")]
+    #[case("create", r#"{"upload_id": 1}"#)]
+    #[case("create", "not json")]
+    #[case("complete", "{}")]
+    #[case("complete", r#"{"version": "invalid"}"#)]
+    #[case("complete", "not json")]
+    #[tokio::test]
+    async fn test_multipart_write_error_preserves_request_id(
+        #[case] failed_stage: &str,
+        #[case] body: &str,
+    ) {
+        let failed_stage = failed_stage.to_string();
+        let body = body.to_string();
+        let request_id = Arc::new(std::sync::Mutex::new(String::new()));
+        let observed_request_id = request_id.clone();
+        let table = Table::new_with_handler_version(
+            "my_table",
+            semver::Version::new(0, 4, 0),
+            move |request| {
+                let path = request.url().path();
+                if path == "/v1/table/my_table/describe/" {
+                    return simple_describe_response();
+                }
+                if path == format!("/v1/table/my_table/multipart_write/{failed_stage}") {
+                    *observed_request_id.lock().unwrap() = request.headers()["x-request-id"]
+                        .to_str()
+                        .unwrap()
+                        .to_string();
+                    return http::Response::builder()
+                        .status(200)
+                        .body(body.clone())
+                        .unwrap();
+                }
+                let response = match path {
+                    "/v1/table/my_table/multipart_write/create" => r#"{"upload_id":"upload-1"}"#,
+                    "/v1/table/my_table/insert/" => r#"{"version":1}"#,
+                    "/v1/table/my_table/multipart_write/abort" => "",
+                    _ => panic!("Unexpected request path: {path}"),
+                };
+                http::Response::builder()
+                    .status(200)
+                    .body(response.to_string())
+                    .unwrap()
+            },
+        );
+        let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
+        let error = table
+            .add(vec![batch])
+            .write_parallelism(2)
+            .execute()
+            .await
+            .unwrap_err();
+        let expected = request_id.lock().unwrap();
+        assert!(!expected.is_empty());
+        match error {
+            Error::Http { request_id, .. } => assert_eq!(request_id, *expected),
+            other => panic!("Expected Http error, got {other:?}"),
+        }
+    }
+
     #[tokio::test]
     async fn test_multipart_write_progress() {
         let callback_count = Arc::new(AtomicUsize::new(0));