From b2eb57198f1d9ebc1ea9f993cbfe46219b1284ff Mon Sep 17 00:00:00 2001 From: miracle Date: Thu, 18 Jun 2026 18:31:41 +0800 Subject: [PATCH] fix: stream XLSX table export to avoid OOM on large tables --- .../assets/database-drivers.manifest.json | 4 +- crates/dbx-core/src/database_capabilities.rs | 10 +- crates/dbx-core/src/table_export.rs | 28 +++- crates/dbx-core/src/xlsx_export.rs | 158 +++++++++++++++++- .../dbx-core/tests/database_capabilities.rs | 1 + 5 files changed, 181 insertions(+), 20 deletions(-) diff --git a/crates/dbx-core/assets/database-drivers.manifest.json b/crates/dbx-core/assets/database-drivers.manifest.json index 2fa17f1bb..3fd753d82 100644 --- a/crates/dbx-core/assets/database-drivers.manifest.json +++ b/crates/dbx-core/assets/database-drivers.manifest.json @@ -1474,8 +1474,8 @@ "supportLevel": "connect", "capabilities": { "queryExecution": false, - "metadataBrowse": false, - "objectBrowser": false, + "metadataBrowse": true, + "objectBrowser": true, "objectSource": false, "schemaSearch": false, "diagram": false, diff --git a/crates/dbx-core/src/database_capabilities.rs b/crates/dbx-core/src/database_capabilities.rs index 0db895f3b..2456ba19b 100644 --- a/crates/dbx-core/src/database_capabilities.rs +++ b/crates/dbx-core/src/database_capabilities.rs @@ -37,8 +37,14 @@ pub fn is_metadata_connection_scoped(db_type: &DatabaseType) -> bool { } pub fn skips_tcp_probe(db_type: &DatabaseType) -> bool { - matches!(db_type, DatabaseType::Sqlite | DatabaseType::DuckDb | DatabaseType::Turso | DatabaseType::Jdbc) - || is_agent_type(db_type) + matches!( + db_type, + DatabaseType::Sqlite + | DatabaseType::DuckDb + | DatabaseType::Turso + | DatabaseType::Jdbc + | DatabaseType::MessageQueue + ) || is_agent_type(db_type) } /// Database types whose connection backs onto a single local file (or may, in the diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index 38c548fcc..8002425c3 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -10,7 +10,7 @@ use crate::transfer::{ count_sql_with_where, execute_on_pool, execute_on_pool_with_max_rows, keyset_pagination_sql, pagination_sql_with_filter_order, }; -use crate::xlsx_export::{build_xlsx_workbook, XlsxWorksheetData}; +use crate::xlsx_export::{finish_streaming_xlsx_workbook, start_streaming_xlsx_workbook}; const DEFAULT_BATCH_SIZE: usize = 10_000; @@ -325,7 +325,13 @@ pub async fn export_table_data_core( } } "xlsx" => { - let mut all_rows: Vec> = Vec::new(); + // Create a dedicated file handle for the streaming XLSX writer + // instead of cloning the outer BufWriter's handle. This avoids + // sharing a file descriptor between two independent buffers. + let xlsx_file = + std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; + let mut writer = + start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some(&request.table_name), &col_names)?; loop { // Check cancellation between batches @@ -358,12 +364,14 @@ pub async fn export_table_data_core( break; } - all_rows.extend(result.rows); + for row in &result.rows { + writer.write_row(row).map_err(|e| format!("Failed to write XLSX row: {e}"))?; + } rows_exported += row_count as u64; if use_keyset { // Keyset pagination: track last PK values for next batch - if let Some(last_row) = all_rows.last() { + if let Some(last_row) = result.rows.last() { last_pk_values = pk_indices.iter().map(|&i| last_row[i].clone()).collect(); } } else { @@ -394,11 +402,12 @@ pub async fn export_table_data_core( error_message: None, }); - // Build XLSX workbook from accumulated rows - let workbook_data = - XlsxWorksheetData { sheet_name: Some(request.table_name.clone()), columns: col_names, rows: all_rows }; - let xlsx_bytes = build_xlsx_workbook(&workbook_data)?; - file.write_all(&xlsx_bytes).map_err(|e| format!("Failed to write XLSX file: {e}"))?; + // Explicitly flush the XLSX writer's BufWriter so IO errors + // (e.g. disk-full) are surfaced rather than silently swallowed + // by Drop. + let mut xlsx_buf = + finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; + xlsx_buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; } "json" => { file.write_all(b"[\n").map_err(|e| format!("Failed to write JSON: {e}"))?; @@ -636,6 +645,7 @@ pub async fn export_table_data_core( mod tests { use super::*; use crate::database_export::{clear_export_cancelled, set_export_cancelled}; + use crate::xlsx_export::{build_xlsx_workbook, XlsxWorksheetData}; use serde_json::json; // ----------------------------------------------------------------------- diff --git a/crates/dbx-core/src/xlsx_export.rs b/crates/dbx-core/src/xlsx_export.rs index 2620e84f7..702612950 100644 --- a/crates/dbx-core/src/xlsx_export.rs +++ b/crates/dbx-core/src/xlsx_export.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::io::{Cursor, Write}; +use std::io::{Cursor, Seek, Write}; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -10,6 +10,131 @@ pub struct XlsxWorksheetData { pub rows: Vec>, } +/// Streaming XLSX writer that incrementally writes rows to a ZIP-backed +/// workbook. This avoids accumulating all rows in memory before building the +/// final file, drastically reducing peak memory for large exports. +pub struct StreamingXlsxWriter { + zip: zip::ZipWriter, + columns: Vec, + next_row_number: usize, +} + +/// Estimate column widths from header names only (used by the streaming path +/// where full row data is not available up-front). Each width is clamped to +/// [10, 60] to stay within reasonable bounds. +fn estimate_header_widths(columns: &[String]) -> Vec { + columns.iter().map(|col| (col.chars().count() + 2).clamp(10, 60)).collect() +} + +/// Build the `` XML fragment from a width slice. +fn cols_xml(widths: &[usize]) -> String { + widths + .iter() + .enumerate() + .map(|(index, width)| { + format!("", index + 1, index + 1, width) + }) + .collect() +} + +/// Build a single `` XML fragment for the header row (row 1). +pub(crate) fn header_row_xml(columns: &[String]) -> String { + format!( + "{}", + columns + .iter() + .enumerate() + .map(|(index, col)| cell_xml(Some(&Value::String(col.clone())), 0, index, Some(1))) + .collect::() + ) +} + +/// Build a single `` XML fragment for a data row. +pub(crate) fn data_row_xml(row_number: usize, columns: &[String], row: &[Value]) -> String { + let cells = columns + .iter() + .enumerate() + .map(|(col_index, _)| cell_xml(row.get(col_index), row_number - 1, col_index, None)) + .collect::(); + format!("{cells}") +} + +fn write_zip_entry(zip: &mut zip::ZipWriter, path: &str, content: &str) -> Result<(), String> { + let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); + zip.start_file(path, options).map_err(|err| err.to_string())?; + zip.write_all(content.as_bytes()).map_err(|err| err.to_string()) +} + +/// Start a new streaming XLSX workbook. The ZIP skeleton, worksheet header, +/// column widths (estimated from header names) and the header row are written +/// immediately. Callers then feed data rows via [`StreamingXlsxWriter::write_row`] +/// and finalize with [`StreamingXlsxWriter::finish`]. +pub(crate) fn start_streaming_xlsx_workbook( + writer: W, + sheet_name: Option<&str>, + columns: &[String], +) -> Result, String> { + let sheet_name = normalize_sheet_name(sheet_name); + let widths = estimate_header_widths(columns); + + let mut zip = zip::ZipWriter::new(writer); + write_zip_entry(&mut zip, "[Content_Types].xml", content_types_xml())?; + write_zip_entry(&mut zip, "_rels/.rels", root_rels_xml())?; + write_zip_entry(&mut zip, "xl/workbook.xml", &workbook_xml(&sheet_name))?; + write_zip_entry(&mut zip, "xl/_rels/workbook.xml.rels", workbook_rels_xml())?; + write_zip_entry(&mut zip, "xl/styles.xml", styles_xml())?; + + // Begin the sheet1.xml entry with header, frozen pane, column widths and + // the header row. + let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); + zip.start_file("xl/worksheets/sheet1.xml", options).map_err(|err| err.to_string())?; + + let sheet_header = format!( + concat!( + "", + "", + "", + "", + "", + "", + "{cols}", + "" + ), + cols = cols_xml(&widths), + ); + zip.write_all(sheet_header.as_bytes()).map_err(|err| err.to_string())?; + zip.write_all(header_row_xml(columns).as_bytes()).map_err(|err| err.to_string())?; + + Ok(StreamingXlsxWriter { zip, columns: columns.to_vec(), next_row_number: 2 }) +} + +impl StreamingXlsxWriter { + /// Append a single data row to the worksheet. + pub fn write_row(&mut self, row: &[Value]) -> Result<(), String> { + self.zip + .write_all(data_row_xml(self.next_row_number, &self.columns, row).as_bytes()) + .map_err(|err| err.to_string())?; + self.next_row_number += 1; + Ok(()) + } + + /// Finalize the worksheet and close the ZIP archive. Returns the + /// underlying writer so callers can flush / close it as needed. + pub fn finish(mut self) -> Result { + let row_count = self.next_row_number.saturating_sub(1); + let range = sheet_range(self.columns.len(), row_count); + self.zip + .write_all(format!("").as_bytes()) + .map_err(|err| err.to_string())?; + self.zip.finish().map_err(|err| err.to_string()) + } +} + +/// Convenience wrapper that finalizes a streaming workbook. +pub(crate) fn finish_streaming_xlsx_workbook(writer: StreamingXlsxWriter) -> Result { + writer.finish() +} + fn escape_xml(value: &str) -> String { let mut result = String::with_capacity(value.len()); for ch in value.chars() { @@ -36,11 +161,7 @@ fn column_name(index: usize) -> String { out.push((b'A' + rem as u8) as char); n = (n - 1) / 26; } - // Safety: all pushed chars are ASCII, so reversing by bytes is correct - unsafe { - out.as_mut_vec().reverse(); - } - out + out.chars().rev().collect() } fn cell_ref(row_index: usize, col_index: usize) -> String { @@ -262,8 +383,10 @@ pub fn build_xlsx_workbook(data: &XlsxWorksheetData) -> Result, String> #[cfg(test)] mod tests { - use super::{build_xlsx_workbook, XlsxWorksheetData}; + use super::{build_xlsx_workbook, start_streaming_xlsx_workbook, XlsxWorksheetData}; + use calamine::{open_workbook_auto, Reader}; use serde_json::json; + use std::fs; #[test] fn builds_xlsx_zip_with_sheet_data() { @@ -296,4 +419,25 @@ mod tests { let text = String::from_utf8_lossy(&workbook); assert!(text.contains("name=\"bad name with chars and-a-very-\"")); } + + #[test] + fn streams_xlsx_rows_to_a_readable_workbook() { + let path = std::env::temp_dir().join(format!("dbx-stream-test-{}.xlsx", uuid::Uuid::new_v4())); + { + let file = fs::File::create(&path).expect("create temp xlsx"); + let mut writer = + start_streaming_xlsx_workbook(file, Some("Streamed"), &["id".to_string(), "name".to_string()]) + .expect("start workbook"); + writer.write_row(&[json!(1), json!("Ada")]).expect("write row"); + writer.write_row(&[json!(2), json!("Bob")]).expect("write row"); + drop(writer.finish().expect("finish workbook")); + } + + let mut workbook = open_workbook_auto(&path).expect("open workbook"); + let range = workbook.worksheet_range("Streamed").expect("read worksheet"); + assert_eq!(range.get_value((0, 0)).expect("header"), &calamine::Data::String("id".to_string())); + assert_eq!(range.get_value((1, 0)).expect("row1"), &calamine::Data::Float(1.0)); + assert_eq!(range.get_value((2, 1)).expect("row2"), &calamine::Data::String("Bob".to_string())); + let _ = fs::remove_file(&path); + } } diff --git a/crates/dbx-core/tests/database_capabilities.rs b/crates/dbx-core/tests/database_capabilities.rs index 43f7ea517..5d73564bf 100644 --- a/crates/dbx-core/tests/database_capabilities.rs +++ b/crates/dbx-core/tests/database_capabilities.rs @@ -188,6 +188,7 @@ fn skips_tcp_probe_for_local_file_plugin_and_agent_types() { assert!(skips_tcp_probe(&DatabaseType::Gbase)); assert!(skips_tcp_probe(&DatabaseType::Databend)); assert!(skips_tcp_probe(&DatabaseType::InfluxDb)); + assert!(skips_tcp_probe(&DatabaseType::MessageQueue)); assert!(!skips_tcp_probe(&DatabaseType::Postgres)); assert!(!skips_tcp_probe(&DatabaseType::Mysql)); assert!(!skips_tcp_probe(&DatabaseType::Gaussdb));