github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(export): split XLSX sheets and stream worksheet XML

Browse Source
This commit is contained in:
Abeautifulsnow authored and GitHub committed 2026-07-31 19:45:37 +08:00
1 parent ecac4800c4
commit c43e8455c4
16 files changed
+1006 -250
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/ai_codex_cli.rs
+2 -1
View File
@@ -984,6 +984,7 @@ pub async fn run_codex_agent(
#[cfg(test)]
mod tests {
use super::codex_process_env_with_system_proxy;
#[cfg(not(windows))]
use super::shell_quote;
use super::{
@@ -996,7 +997,7 @@ mod tests {
DEFAULT_CODEX_MODELS,
};
#[cfg(not(windows))]
use super::{codex_process_env, codex_process_env_with_system_proxy, common_executable_dirs, merged_path_with_dir};
use super::{codex_process_env, common_executable_dirs, merged_path_with_dir};
#[cfg(windows)]
use super::{
direct_program_path, first_windows_program_path, program_path_candidates, resolve_codex_command,
crates/dbx-core/src/query_result_export.rs
+28 -92
View File
@@ -35,8 +35,6 @@ use sqlparser::parser::Parser;
use tokio_util::sync::CancellationToken;
const AGENT_UNBOUNDED_ROW_LIMIT: usize = i32::MAX as usize;
pub const XLSX_MAX_DATA_ROWS: usize = 1_048_575;
const XLSX_ROW_LIMIT_ERROR: &str = "XLSX supports at most 1,048,575 data rows. Use CSV export for the full result.";
const STREAMING_PAGINATION_UNSUPPORTED_ERROR: &str =
"Streaming export is unsupported for this query. Simplify it or use a supported driver.";
const AGENT_SESSION_MISSING_ERROR: &str =
@@ -241,11 +239,7 @@ fn progress(
status: ExportStatus,
error_message: Option<String>,
) -> TableExportProgress {
let total_rows = request.total_rows.map(|total| {
let format = request.format.to_lowercase();
let limit = effective_row_limit(&format, request);
limit.map_or(total, |limit| total.min(limit as u64))
});
let total_rows = request.total_rows.map(|total| request.row_limit.map_or(total, |limit| total.min(limit as u64)));
TableExportProgress {
export_id: request.export_id.clone(),
table_name: String::new(),
@@ -387,16 +381,8 @@ impl SqlInsertWriter {
}
}
fn effective_row_limit(format: &str, request: &QueryResultExportRequest) -> Option<usize> {
if format == "xlsx" {
Some(request.row_limit.map_or(XLSX_MAX_DATA_ROWS, |limit| limit.min(XLSX_MAX_DATA_ROWS)))
} else {
request.row_limit
}
}
fn xlsx_hard_limit_active(format: &str, request: &QueryResultExportRequest) -> bool {
format == "xlsx" && request.row_limit.is_none_or(|limit| limit > XLSX_MAX_DATA_ROWS)
fn effective_row_limit(request: &QueryResultExportRequest) -> Option<usize> {
request.row_limit
}
fn format_text_export_header(format: &str, columns: &[String]) -> String {
@@ -619,18 +605,8 @@ async fn export_query_result_core_inner(
}
let page_size = request.page_size.max(1);
let effective_row_limit = effective_row_limit(&format, request);
let xlsx_hard_limit_active = xlsx_hard_limit_active(&format, request);
if xlsx_hard_limit_active && request.total_rows.is_some_and(|total| total > XLSX_MAX_DATA_ROWS as u64) {
return Err(XLSX_ROW_LIMIT_ERROR.to_string());
}
let agent_max_rows = if xlsx_hard_limit_active {
XLSX_MAX_DATA_ROWS + 1
} else {
effective_row_limit.unwrap_or(AGENT_UNBOUNDED_ROW_LIMIT)
}
.max(1);
let effective_row_limit = effective_row_limit(request);
let agent_max_rows = effective_row_limit.unwrap_or(AGENT_UNBOUNDED_ROW_LIMIT).max(1);
on_progress(progress(request, 0, ExportStatus::Running, None));
@@ -695,11 +671,6 @@ async fn export_query_result_core_inner(
break;
}
let this_page = remaining.map_or(page_size, |rem| rem.min(page_size)).max(1);
let fetch_limit = if xlsx_hard_limit_active && remaining.is_some_and(|rem| rem <= page_size) {
this_page.saturating_add(1)
} else {
this_page
};
let (sql_to_execute, plan_limit, use_agent_result_session) = if let Some(plan) = keyset_plan.as_ref() {
(
@@ -710,9 +681,9 @@ async fn export_query_result_core_inner(
&request.database_type,
&plan.primary_keys,
&plan.last_pk_values,
fetch_limit,
this_page,
),
fetch_limit,
this_page,
false,
)
} else {
@@ -720,7 +691,7 @@ async fn export_query_result_core_inner(
sql: request.sql.clone(),
query_base_sql: request.query_base_sql.clone(),
database_type: Some(request.database_type),
pagination: QueryPagination { limit: fetch_limit, offset, session_id: session_id.clone() },
pagination: QueryPagination { limit: this_page, offset, session_id: session_id.clone() },
use_agent_cursor: request.use_agent_cursor,
first_page_uses_actual_sql: true,
});
@@ -801,12 +772,6 @@ async fn export_query_result_core_inner(
}
}
let fetched_row_count = result.rows.len();
if xlsx_hard_limit_active {
let remaining_rows = XLSX_MAX_DATA_ROWS.saturating_sub(rows_exported as usize);
if fetched_row_count > remaining_rows {
return Err(XLSX_ROW_LIMIT_ERROR.to_string());
}
}
if result.rows.len() > this_page {
result.rows.truncate(this_page);
}
@@ -964,10 +929,8 @@ async fn try_export_postgres_query_result_stream(
state.touch_pool_activity(&pool_key).await;
let _activity_touch = state.pool_activity_touch(&pool_key);
let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request);
let row_limit = effective_row_limit(format, request);
let stream_row_limit =
if xlsx_hard_limit_active { row_limit.map(|limit| limit.saturating_add(1)) } else { row_limit };
let row_limit = effective_row_limit(request);
let stream_row_limit = row_limit;
let progress_row_interval = request.page_size.max(1) as u64;
let mut columns: Vec<String> = Vec::new();
let mut temporal_column_types: Vec<String> = Vec::new();
@@ -1020,9 +983,6 @@ async fn try_export_postgres_query_result_stream(
}
}
crate::db::postgres::PostgresQueryStreamItem::Row(row) => {
if xlsx_hard_limit_active && rows_exported as usize >= XLSX_MAX_DATA_ROWS {
return Err(XLSX_ROW_LIMIT_ERROR.to_string());
}
let formatted = crate::temporal_format::format_temporal_export_row_with_string_types(
&row,
&temporal_column_types,
@@ -1161,10 +1121,8 @@ async fn try_export_mysql_query_result_stream(
crate::query_execution_sql::check_read_only(&request.sql, &name, database_type)?;
}
let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request);
let row_limit = effective_row_limit(format, request);
let stream_row_limit =
if xlsx_hard_limit_active { row_limit.map(|limit| limit.saturating_add(1)) } else { row_limit };
let row_limit = effective_row_limit(request);
let stream_row_limit = row_limit;
let progress_row_interval = request.page_size.max(1) as u64;
let mut columns: Vec<String> = Vec::new();
let mut temporal_column_types: Vec<String> = Vec::new();
@@ -1266,9 +1224,6 @@ async fn try_export_mysql_query_result_stream(
}
}
crate::db::mysql::MySqlQueryStreamItem::Row(row) => {
if xlsx_hard_limit_active && rows_exported as usize >= XLSX_MAX_DATA_ROWS {
return Err(XLSX_ROW_LIMIT_ERROR.to_string());
}
let formatted = crate::temporal_format::format_temporal_export_row_with_string_types(
&row,
&temporal_column_types,
@@ -1427,10 +1382,8 @@ async fn try_export_clickhouse_query_result_stream(
state.touch_pool_activity(&pool_key).await;
let _activity_touch = state.pool_activity_touch(&pool_key);
let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request);
let row_limit = effective_row_limit(format, request);
let stream_row_limit =
if xlsx_hard_limit_active { row_limit.map(|limit| limit.saturating_add(1)) } else { row_limit };
let row_limit = effective_row_limit(request);
let stream_row_limit = row_limit;
let progress_row_interval = request.page_size.max(1) as u64;
let mut columns: Vec<String> = Vec::new();
let mut temporal_column_types: Vec<String> = Vec::new();
@@ -1484,9 +1437,6 @@ async fn try_export_clickhouse_query_result_stream(
}
}
crate::db::clickhouse_driver::ClickHouseQueryStreamItem::Row(row) => {
if xlsx_hard_limit_active && rows_exported as usize >= XLSX_MAX_DATA_ROWS {
return Err(XLSX_ROW_LIMIT_ERROR.to_string());
}
let formatted = crate::temporal_format::format_temporal_export_row_with_string_types(
&row,
&temporal_column_types,
@@ -1600,10 +1550,8 @@ async fn try_export_sqlserver_query_result_stream(
state.running_queries.set_pool_key(execution_id, pool_key);
}
let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request);
let row_limit = effective_row_limit(format, request);
let stream_row_limit =
if xlsx_hard_limit_active { row_limit.map(|limit| limit.saturating_add(1)) } else { row_limit };
let row_limit = effective_row_limit(request);
let stream_row_limit = row_limit;
let mut columns: Vec<String> = Vec::new();
let mut temporal_column_types: Vec<String> = Vec::new();
let mut rows_exported = 0_u64;
@@ -1667,9 +1615,6 @@ async fn try_export_sqlserver_query_result_stream(
}
}
crate::db::sqlserver::SqlServerStreamItem::Row(row) => {
if xlsx_hard_limit_active && rows_exported as usize >= XLSX_MAX_DATA_ROWS {
return Err(XLSX_ROW_LIMIT_ERROR.to_string());
}
let formatted = crate::temporal_format::format_temporal_export_row_with_string_types(
row,
&temporal_column_types,
@@ -1865,12 +1810,12 @@ mod tests {
#[test]
fn csv_unlimited_export_has_no_effective_row_limit() {
assert_eq!(effective_row_limit("csv", &request("csv", None, None)), None);
assert_eq!(effective_row_limit(&request("csv", None, None)), None);
}
#[test]
fn txt_unlimited_export_has_no_effective_row_limit() {
assert_eq!(effective_row_limit("txt", &request("txt", None, None)), None);
assert_eq!(effective_row_limit(&request("txt", None, None)), None);
}
#[test]
@@ -1879,31 +1824,22 @@ mod tests {
}
#[test]
fn xlsx_unlimited_export_uses_excel_hard_limit() {
assert_eq!(effective_row_limit("xlsx", &request("xlsx", None, None)), Some(XLSX_MAX_DATA_ROWS));
fn xlsx_no_row_limit_has_no_query_layer_cap() {
// Without the old hard limit, XLSX uses the writer's internal splitting.
assert_eq!(effective_row_limit(&request("xlsx", None, None)), None);
}
#[test]
fn xlsx_row_limit_caps_to_excel_hard_limit() {
assert_eq!(
effective_row_limit("xlsx", &request("xlsx", Some(XLSX_MAX_DATA_ROWS + 10), None)),
Some(XLSX_MAX_DATA_ROWS)
);
fn xlsx_user_row_limit_still_respected() {
assert_eq!(effective_row_limit(&request("xlsx", Some(500), None)), Some(500));
}
#[test]
fn xlsx_known_total_above_hard_limit_errors_before_export() {
let req = request("xlsx", None, Some(XLSX_MAX_DATA_ROWS as u64 + 1));
assert!(xlsx_hard_limit_active("xlsx", &req));
assert!(req.total_rows.is_some_and(|total| total > XLSX_MAX_DATA_ROWS as u64));
}
#[test]
fn sql_insert_export_has_no_xlsx_row_cap() {
// SQL format should not be limited by XLSX_MAX_DATA_ROWS.
let req = request("sql", None, None);
assert!(!xlsx_hard_limit_active("sql", &req));
assert_eq!(effective_row_limit("sql", &req), None);
fn xlsx_total_rows_above_sheet_limit_no_longer_errors() {
// total_rows > 1M no longer triggers a pre-check error; the writer splits.
let req = request("xlsx", None, Some(2_000_000));
assert!(effective_row_limit(&req).is_none());
// The function that used to check this (xlsx_hard_limit_active) no longer exists.
}
#[test]
crates/dbx-core/src/xlsx_export.rs
+838 -97
View File
File diff suppressed because it is too large. Load diff