xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: stream XLSX table export to avoid OOM on large tables

This commit is contained in:
miracle 2026-06-18 18:31:41 +08:00 committed by GitHub
parent 320336e84e
commit b2eb57198f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 181 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
crates/dbx-core/assets/database-drivers.manifest.json
View File

@ -1474,8 +1474,8 @@
"supportLevel": "connect",
"capabilities": {
"queryExecution": false,
"metadataBrowse": false,
"objectBrowser": false,
"metadataBrowse": true,
"objectBrowser": true,
"objectSource": false,
"schemaSearch": false,
"diagram": false,

10
crates/dbx-core/src/database_capabilities.rs
View File

@ -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

28
crates/dbx-core/src/table_export.rs
View File

@ -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<Value>> = 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;
// -----------------------------------------------------------------------

158
crates/dbx-core/src/xlsx_export.rs
View File

@ -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<Vec<Value>>,
}
/// 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<W: Write + Seek> {
zip: zip::ZipWriter<W>,
columns: Vec<String>,
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<usize> {
columns.iter().map(|col| (col.chars().count() + 2).clamp(10, 60)).collect()
}
/// Build the `<cols>` XML fragment from a width slice.
fn cols_xml(widths: &[usize]) -> String {
widths
.iter()
.enumerate()
.map(|(index, width)| {
format!("<col min=\"{}\" max=\"{}\" width=\"{}\" customWidth=\"1\"/>", index + 1, index + 1, width)
})
.collect()
}
/// Build a single `<row>` XML fragment for the header row (row 1).
pub(crate) fn header_row_xml(columns: &[String]) -> String {
format!(
"<row r=\"1\">{}</row>",
columns
.iter()
.enumerate()
.map(|(index, col)| cell_xml(Some(&Value::String(col.clone())), 0, index, Some(1)))
.collect::<String>()
)
}
/// Build a single `<row>` 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::<String>();
format!("<row r=\"{row_number}\">{cells}</row>")
}
fn write_zip_entry<W: Write + Seek>(zip: &mut zip::ZipWriter<W>, 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<W: Write + Seek>(
writer: W,
sheet_name: Option<&str>,
columns: &[String],
) -> Result<StreamingXlsxWriter<W>, 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!(
"<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>",
"<worksheet xmlns=\"http://schemas.openxmlformats.org/spreadsheetml/2006/main\">",
"<sheetViews><sheetView workbookViewId=\"0\">",
"<pane ySplit=\"1\" topLeftCell=\"A2\" activePane=\"bottomLeft\" state=\"frozen\"/>",
"</sheetView></sheetViews>",
"<sheetFormatPr defaultRowHeight=\"15\"/>",
"<cols>{cols}</cols>",
"<sheetData>"
),
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<W: Write + Seek> StreamingXlsxWriter<W> {
/// 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<W, String> {
let row_count = self.next_row_number.saturating_sub(1);
let range = sheet_range(self.columns.len(), row_count);
self.zip
.write_all(format!("</sheetData><autoFilter ref=\"{range}\"/></worksheet>").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<W: Write + Seek>(writer: StreamingXlsxWriter<W>) -> Result<W, String> {
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<Vec<u8>, 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);
}
}

1
crates/dbx-core/tests/database_capabilities.rs
View File

@ -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));