From c8b4224587f6133c8f6d4d73c0516f6ba48c09ad Mon Sep 17 00:00:00 2001 From: DengQingNian <117324885+DengQingNian@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:47:00 +0800 Subject: [PATCH] perf(export): stream large SQL INSERT exports in background --- apps/desktop/src/components/grid/DataGrid.vue | 8 +- .../src/components/layout/ContentArea.vue | 4 +- .../src/composables/useDataGridExport.ts | 79 +++- apps/desktop/src/lib/backend/tauri.ts | 4 +- apps/desktop/src/stores/queryStore.ts | 6 +- crates/dbx-core/src/query_result_export.rs | 380 +++++++++++++++++- .../live_clickhouse_query_result_export.rs | 2 + crates/dbx-core/tests/live_mysql57.rs | 4 + .../live_postgres_query_result_export.rs | 10 + .../tests/live_sqlserver_completion.rs | 2 + .../live_sqlserver_query_result_export.rs | 2 + ...esult-sql-export-background-task-design.md | 380 ++++++++++++++++++ packages/app-tests/useDataGridExport.test.ts | 29 +- src-tauri/src/commands/query_result_export.rs | 175 +++++++- 14 files changed, 1047 insertions(+), 38 deletions(-) create mode 100644 docs/superpowers/specs/2026-07-25-query-result-sql-export-background-task-design.md diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 7fe430aff6..be05730195 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -306,7 +306,7 @@ interface DataGridProps { exportSql?: string; onExecuteSql?: (sql: string) => Promise; fullExportResult?: (onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => Promise; - queryResultExportRequest?: (options: { exportId: string; filePath: string; format: "csv" | "xlsx" | "txt"; includeSqlSheet?: boolean }) => Promise; + queryResultExportRequest?: (options: { exportId: string; filePath: string; format: "csv" | "xlsx" | "txt" | "sql"; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array }) => Promise; allExportResults?: Array<{ sheetName: string; result: QueryResult; sql?: string }>; exportFileBaseName?: string; customSaveHandler?: import("@/composables/useDataGridEditor").CustomSaveHandler; @@ -1760,8 +1760,8 @@ const tableColumnTypesByName = computed(() => { } return map; }); -const visibleColumnTypes = computed(() => - visibleColumnIndexes.value.map((index) => +const allColumnTypes = computed(() => + props.result.columns.map((_, index) => resolveResultColumnType({ resultColumnType: props.result.column_types?.[index], resultColumnName: props.result.columns[index]?.toLocaleLowerCase(), @@ -1770,6 +1770,7 @@ const visibleColumnTypes = computed(() => }), ), ); +const visibleColumnTypes = computed(() => visibleColumnIndexes.value.map((index) => allColumnTypes.value[index])); const visibleColumnCount = computed(() => visibleColumnIndexes.value.length); const numericColumnRightAlign = computed(() => (settingsStore.editorSettings.numericColumnRightAlign ?? true) && !showTranspose.value); @@ -5182,6 +5183,7 @@ const { sourceColumns: visibleSourceColumns, mongoDocuments: computed(() => props.result.mongo_copy_documents ?? props.result.mongo_documents), columnTypes: visibleColumnTypes, + allColumnTypes, whereInput: computed(() => currentWhereInput()), orderBy: computed(() => currentOrderBy()), exportBatchSize: computed(() => settingsStore.editorSettings.exportBatchSize), diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 5609195b7d..adaf7c88fd 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -1353,7 +1353,9 @@ defineExpose({ focusSearch, refreshData, refreshQueryEditorCompletionCache, hand :total-row-count-loading="activeTab.resultTotalRowCountLoading" :on-execute-sql="async (sql: string) => emit('executeSql', sql)" :full-export-result="(onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => queryStore.fetchTabResultForExport(activeTab.id, onProgress)" - :query-result-export-request="(options: { exportId: string; filePath: string; format: 'csv' | 'xlsx' | 'txt'; includeSqlSheet?: boolean }) => queryStore.buildQueryResultExportRequest(activeTab.id, options)" + :query-result-export-request=" + (options: { exportId: string; filePath: string; format: 'csv' | 'xlsx' | 'txt' | 'sql'; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array }) => queryStore.buildQueryResultExportRequest(activeTab.id, options) + " :all-export-results="allResultExportSheets" :export-file-base-name="activeTab.title" @update:order-by-input="(v: string) => (activeTab.orderByInput = v)" diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index 3d88126872..e970fb7c68 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -6,6 +6,7 @@ import * as api from "@/lib/backend/api"; import { type CellSelectionMatrix, type CellSelectionRange, type SelectionData } from "@/lib/dataGrid/gridSelection"; import type { DataGridExtractorOptions } from "@/lib/dataGrid/dataGridCopyExtractor"; import { useToast } from "@/composables/useToast"; +import { useExportTracker } from "@/composables/useExportTracker"; import { displayCellValue, type CellValue } from "@/lib/dataGrid/cellValue"; import { tryStartExclusiveActivation, type ActionActivationGuard } from "@/lib/connection/actionActivation"; import { copyToClipboard } from "@/lib/common/clipboard"; @@ -78,6 +79,7 @@ export interface UseDataGridExportOptions { sourceColumns: ComputedRef | undefined>; mongoDocuments?: ComputedRef; columnTypes: ComputedRef | undefined>; + allColumnTypes?: ComputedRef | undefined>; whereInput: ComputedRef; orderBy: ComputedRef; exportBatchSize: ComputedRef; @@ -92,7 +94,7 @@ export interface UseDataGridExportOptions { selectedRowIds: Ref> | ComputedRef>; hasRowSelection: ComputedRef; fullExportResult?: (onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => Promise; - queryResultExportRequest?: (options: { exportId: string; filePath: string; format: "csv" | "xlsx" | "txt"; includeSqlSheet?: boolean }) => Promise; + queryResultExportRequest?: (options: { exportId: string; filePath: string; format: "csv" | "xlsx" | "txt" | "sql"; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array }) => Promise; /** * True when the in-memory result already holds the complete result set — * i.e. the query ran without server-side pagination, was not truncated, and @@ -138,6 +140,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { const { t } = useI18n(); const { toast } = useToast(); const exportGuard: ActionActivationGuard = {}; + const { addTask, updateTableExportTask, registerTaskCancelHandler, unregisterTaskCancelHandler, removeTask } = useExportTracker(); const { columns, @@ -159,6 +162,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { whereInput, orderBy, columnTypes, + allColumnTypes: allColumnTypesOption, exportBatchSize, hasCellSelection, hasColumnSelection: hasColumnSelectionOption, @@ -185,6 +189,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { const allColumns = allColumnsOption ?? columns; const allDisplayItems = allDisplayItemsOption ?? displayItems; const allSourceColumns = allSourceColumnsOption ?? sourceColumns; + const allColumnTypes = computed(() => allColumnTypesOption?.value ?? columnTypes.value); const visibleColumnIndexes = visibleColumnIndexesOption ?? computed(() => columns.value.map((_, index) => index)); const hasColumnSelection = hasColumnSelectionOption ?? computed(() => false); @@ -1018,11 +1023,83 @@ export function useDataGridExport(options: UseDataGridExportOptions) { return true; } + async function exportQueryResultSqlViaBackend(rowIds?: number[]): Promise { + // Guard: only for query-result context without complete local result, desktop only + if (rowIds !== undefined || context.value !== "results" || !queryResultExportRequest) return false; + if (hasCompleteLocalResult?.value) return false; + if (!isTauriRuntime()) return false; // Web → local export fallback + + // 1. Save dialog FIRST (immediate user feedback) + let outputPath = exportFileName("query-result", "sql"); + const { save } = await import("@tauri-apps/plugin-dialog"); + const path = await save({ + defaultPath: outputPath, + filters: [{ name: "SQL", extensions: ["sql"] }], + }); + if (!path) return true; + outputPath = path as string; + + // 2. Register background task FIRST so its exportId drives the request. + // addTask generates its own ID — capture it so progress callbacks match. + const task = addTask(tableMeta.value?.tableName || "Query Result", "sql", outputPath); + const taskExportId = task.exportId; + + let request: api.QueryResultExportRequest; + try { + const built = await queryResultExportRequest({ + exportId: taskExportId, + filePath: outputPath, + format: "sql", + exportTableName: tableMeta.value?.tableName, + exportColumnTypes: allColumnTypes.value?.map((t) => t ?? null) as Array | undefined, + }); + if (!built) { + // builder declined — clean up the task we just registered + removeTask(taskExportId); + return false; + } + request = built; + } catch { + removeTask(taskExportId); + return false; + } + + registerTaskCancelHandler(taskExportId, () => api.cancelQueryResultExport(taskExportId, request.executionId)); + + try { + await api.startQueryResultExport(request, (progress) => { + updateTableExportTask(taskExportId, progress); + if (progress.status === "Done") { + toast(t("grid.exported")); + } + }); + } catch (e) { + // 4. Startup rejection → mark task Error (fixes stuck-Running bug) + updateTableExportTask(taskExportId, { + exportId: taskExportId, + tableName: tableMeta.value?.tableName || "Query Result", + rowsExported: 0, + totalRows: null, + status: "Error" as const, + errorMessage: (e as Error)?.message || String(e), + }); + throw e; + } finally { + unregisterTaskCancelHandler(taskExportId); + } + return true; + } + async function exportSql(rowIds?: number[]) { await runExclusiveExport(async () => { try { + // Step 1: table-data context — existing backend table export if (await exportFullTableDataViaBackend("sql", rowIds)) return; + // Step 2: query-result context — NEW backend streaming with background task + if (await exportQueryResultSqlViaBackend(rowIds)) return; + + // Step 3: fallback — local export (Web and edge-case scenarios) const result = await resultToExport(rowIds, undefined, true, false); const exportData = sqlInsertExportData(result); const content = await formatSqlInsert({ diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 011fc66f3d..3e95464bc5 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -2747,7 +2747,7 @@ export interface QueryResultExportRequest { databaseType: DatabaseType; useAgentCursor: boolean; filePath: string; - format: "csv" | "xlsx" | "txt"; + format: "csv" | "xlsx" | "txt" | "sql"; includeSqlSheet?: boolean; pageSize: number; rowLimit?: number | null; @@ -2757,6 +2757,8 @@ export interface QueryResultExportRequest { clientSessionId?: string; executionId?: string; dateTimeFormat?: string; + exportTableName?: string; + exportColumnTypes?: Array; numericColumnRightAlign?: boolean; } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 8bbb0c725c..d2e9bcccc4 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -84,8 +84,10 @@ function cloneTabDraft(value: T): T { interface BuildQueryResultExportRequestOptions { exportId: string; filePath: string; - format: "csv" | "xlsx" | "txt"; + format: "csv" | "xlsx" | "txt" | "sql"; includeSqlSheet?: boolean; + exportTableName?: string; + exportColumnTypes?: Array; } type DroppedTableObjectType = "TABLE" | "VIEW" | "MATERIALIZED_VIEW"; @@ -4705,6 +4707,8 @@ export const useQueryStore = defineStore("query", () => { keysetOptimizationEnabled: settings.queryExportKeysetOptimizationEnabled, clientSessionId, executionId: uuid(), + exportTableName: options.exportTableName, + exportColumnTypes: options.exportColumnTypes, numericColumnRightAlign: settings.numericColumnRightAlign, }; } diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index bbb7720187..09dbd4dc2f 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -1,6 +1,7 @@ use serde::{Deserialize, Serialize}; use std::fs::File; use std::io::{BufWriter, Seek, Write}; +use std::mem; use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, @@ -9,8 +10,8 @@ use std::time::{Duration, Instant}; use crate::connection::{AppState, PoolKind}; use crate::csv_export::{format_query_result_csv, format_query_result_csv_rows, format_tsv, format_tsv_rows}; -use crate::database_export::is_export_cancelled; pub use crate::database_export::ExportStatus; +use crate::database_export::{build_export_insert_statements, is_export_cancelled, BuildExportInsertStatementsOptions}; use crate::models::connection::DatabaseType; use crate::query::{ await_stream_with_progress_timeout, canceled_error, close_query_session, execute_sql_statement_with_options, @@ -39,6 +40,7 @@ const STREAMING_PAGINATION_UNSUPPORTED_ERROR: &str = "当前查询暂不支持 const AGENT_SESSION_MISSING_ERROR: &str = "查询结果流式导出需要驱动返回结果集会话，但当前驱动未返回 session_id。"; const STREAM_PROGRESS_TIME_INTERVAL: Duration = Duration::from_secs(1); const EXCEL_CELL_CHARACTER_LIMIT: usize = 32_767; +const SQL_INSERT_BATCH_SIZE: usize = 100; async fn disconnect_with_timeout( connection: C, @@ -86,6 +88,14 @@ pub struct QueryResultExportRequest { pub execution_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub date_time_format: Option, + // -- new fields for SQL INSERT export -- + #[serde(default, skip_serializing_if = "Option::is_none")] + pub export_table_name: Option, + /// Column type overrides for SQL INSERT export. Each entry may be `null` + /// (meaning "infer from the query result") so the inner element is `Option`. + /// Frontend sends these in original full-query column order. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub export_column_types: Option>>, #[serde(default)] pub numeric_column_right_align: bool, } @@ -189,6 +199,162 @@ fn progress( } } +fn stream_export_was_cancelled(error: &str, token_cancelled: bool, export_cancelled: bool) -> bool { + error == QUERY_CANCELED || token_cancelled || export_cancelled +} + +/// Map the request's export_column_types (Web export may omit them) to +/// the Vec> expected by build_export_insert_statements. +/// +/// `column_types` are the types returned by the executed query (original column +/// order). The request's overrides are expected to align 1:1 in the same order. +/// If the request provides fewer overrides than the result has columns the +/// extra columns are left untyped. Overrides that are `None` or empty are +/// treated as "infer from the query result". +fn sql_insert_column_types(request: &QueryResultExportRequest, column_types: &[String]) -> Vec> { + match request.export_column_types.as_ref() { + Some(overrides) => { + let mut result: Vec> = overrides + .iter() + .map(|t| match t { + Some(s) if !s.is_empty() => Some(s.clone()), + _ => None, + }) + .collect(); + // Pad with None if fewer overrides than result columns + result.resize(column_types.len(), None); + result + } + None => vec![None; column_types.len()], + } +} + +/// Bounded SQL INSERT writer with temp-file + atomic-rename safety. +/// +/// Rows are buffered and flushed to a temp file every [`SQL_INSERT_BATCH_SIZE`] +/// rows, so memory stays bounded regardless of the query page size. The temp file +/// lives alongside the user-chosen target (`.dbx-export-tmp`) and is +/// renamed onto the target only when [`SqlInsertWriter::finish`] succeeds. If the +/// writer is dropped without finishing (export error or cancellation), the temp +/// file is removed and the user's target file is never truncated or deleted. +struct SqlInsertWriter { + file: Option>, + temp_path: std::path::PathBuf, + target_path: std::path::PathBuf, + pending_rows: Vec>, + columns: Vec, + column_types: Vec>, + database_type: DatabaseType, + schema: Option, + table_name: String, + finished: bool, +} + +impl SqlInsertWriter { + /// Create the writer and open the temp file. Column metadata is supplied later + /// via [`SqlInsertWriter::set_columns`] once the executed result is known. + fn create(request: &QueryResultExportRequest) -> Result { + let target_path = std::path::PathBuf::from(&request.file_path); + let mut temp_path = target_path.clone(); + let mut temp_name = target_path + .file_name() + .map(|name| name.to_os_string()) + .unwrap_or_else(|| std::ffi::OsString::from("export.sql")); + temp_name.push(".dbx-export-tmp"); + temp_path.set_file_name(temp_name); + let file = BufWriter::new( + File::create(&temp_path).map_err(|e| format!("Failed to create SQL export temp file: {e}"))?, + ); + let table_name = request + .export_table_name + .as_deref() + .filter(|name| !name.trim().is_empty()) + .unwrap_or("query_result") + .to_string(); + Ok(Self { + file: Some(file), + temp_path, + target_path, + pending_rows: Vec::new(), + columns: Vec::new(), + column_types: Vec::new(), + database_type: request.database_type, + schema: request.schema.clone(), + table_name, + finished: false, + }) + } + + /// Provide result metadata once columns are known. `result_column_types` is the + /// column-type list returned by the executed query (original column order); the + /// request's export_column_types may override it when present. + fn set_columns( + &mut self, + columns: Vec, + result_column_types: &[String], + request: &QueryResultExportRequest, + ) { + self.column_types = sql_insert_column_types(request, result_column_types); + self.columns = columns; + } + + fn write_row(&mut self, row: Vec) -> Result<(), String> { + self.pending_rows.push(row); + if self.pending_rows.len() >= SQL_INSERT_BATCH_SIZE { + self.flush_batch()?; + } + Ok(()) + } + + fn flush_batch(&mut self) -> Result<(), String> { + if self.pending_rows.is_empty() { + return Ok(()); + } + let stmts = build_export_insert_statements(BuildExportInsertStatementsOptions { + database_type: Some(self.database_type), + schema: self.schema.clone(), + table_name: Some(self.table_name.clone()), + qualified_table_name: None, + columns: self.columns.clone(), + column_types: self.column_types.clone(), + column_extras: Vec::new(), + rows: mem::take(&mut self.pending_rows), + batch_size: Some(SQL_INSERT_BATCH_SIZE), + })?; + let file = self.file.as_mut().ok_or_else(|| "SQL export file already closed".to_string())?; + for stmt in &stmts { + writeln!(file, "{stmt}").map_err(|e| format!("Failed to write SQL: {e}"))?; + } + Ok(()) + } + + /// Flush remaining rows, close the temp file, and atomically rename it onto the + /// target path. After this succeeds the writer no longer removes the file on drop. + fn finish(mut self) -> Result<(), String> { + self.flush_batch()?; + if let Some(file) = self.file.as_mut() { + file.flush().map_err(|e| format!("Failed to flush SQL file: {e}"))?; + } + // Close the file handle before rename (required on Windows). + self.file.take(); + self.finished = true; + std::fs::rename(&self.temp_path, &self.target_path) + .map_err(|e| format!("Failed to finalize SQL export file: {e}"))?; + Ok(()) + } +} + +impl Drop for SqlInsertWriter { + fn drop(&mut self) { + if !self.finished { + // Best-effort: close the handle then remove the temp file so the user's + // chosen target is never left truncated or deleted by a failed export. + self.file.take(); + let _ = std::fs::remove_file(&self.temp_path); + } + } +} + fn effective_row_limit(format: &str, request: &QueryResultExportRequest) -> Option { if format == "xlsx" { Some(request.row_limit.map_or(XLSX_MAX_DATA_ROWS, |limit| limit.min(XLSX_MAX_DATA_ROWS))) @@ -415,7 +581,7 @@ async fn export_query_result_core_inner( session_id: &mut Option, ) -> Result<(), String> { let format = request.format.to_lowercase(); - if format != "csv" && format != "xlsx" && format != "txt" { + if format != "csv" && format != "xlsx" && format != "txt" && format != "sql" { return Err(format!("Unsupported streaming query-result export format: {format}")); } @@ -475,6 +641,9 @@ async fn export_query_result_core_inner( return Err(STREAMING_PAGINATION_UNSUPPORTED_ERROR.to_string()); } + let mut sql_writer: Option = + if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; + loop { if cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) || is_export_cancelled(&request.export_id).await @@ -594,6 +763,9 @@ async fn export_query_result_core_inner( if columns.is_empty() { columns = result.columns.clone(); column_types = result.column_types.clone(); + if let Some(writer) = sql_writer.as_mut() { + writer.set_columns(columns.clone(), &column_types, request); + } } let fetched_row_count = result.rows.len(); if xlsx_hard_limit_active { @@ -627,6 +799,11 @@ async fn export_query_result_core_inner( write!(file, "\n{rows}").map_err(|e| format!("Failed to write export rows: {e}"))?; } } + } else if format == "sql" { + let writer = sql_writer.as_mut().ok_or_else(|| "SQL export writer missing".to_string())?; + for row in formatted_rows { + writer.write_row(row)?; + } } else { if xlsx.is_none() { let xlsx_file = @@ -691,6 +868,10 @@ async fn export_query_result_core_inner( if let Some(file) = text_file.as_mut() { file.flush().map_err(|e| format!("Failed to flush text export file: {e}"))?; } + } else if format == "sql" { + if let Some(writer) = sql_writer.take() { + writer.finish()?; + } } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; @@ -769,17 +950,19 @@ async fn try_export_postgres_query_result_stream( None }; let mut xlsx = None; + let mut sql_writer: Option = + if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let budget = operation_budget_for_pool_key(state, &pool_key, query_export_timeout(request.timeout_secs)).await; let cancel_context = state.get_postgres_cancel_context(&pool_key).await; let setup_sql = safe_postgres_temp_setup_sql(&request.setup_sql).unwrap_or_default(); - crate::db::postgres::stream_select_query_with_cancel( + let stream_result = crate::db::postgres::stream_select_query_with_cancel( &pool, request.schema.as_deref(), &setup_sql, &request.sql, stream_row_limit, - cancel_token, + cancel_token.clone(), budget, cancel_context, |item| { @@ -787,7 +970,9 @@ async fn try_export_postgres_query_result_stream( crate::db::postgres::PostgresQueryStreamItem::Columns { columns: stream_columns, column_types } => { columns = stream_columns; temporal_column_types = column_types.clone(); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.set_columns(columns.clone(), &column_types, request); + } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; } else { @@ -810,7 +995,9 @@ async fn try_export_postgres_query_result_stream( &temporal_column_types, request.date_time_format.as_deref(), ); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.write_row(formatted)?; + } else if let Some(file) = text_file.as_mut() { let rows = format_text_export_rows(format, std::slice::from_ref(&formatted)); write!(file, "\n{rows}").map_err(|e| format!("Failed to write export rows: {e}"))?; } else if let Some(writer) = xlsx.as_mut() { @@ -841,7 +1028,25 @@ async fn try_export_postgres_query_result_stream( Ok(()) }, ) - .await?; + .await; + + if let Err(error) = stream_result { + let export_cancelled = is_export_cancelled(&request.export_id).await; + if stream_export_was_cancelled( + &error, + cancel_token.as_ref().is_some_and(|token| token.is_cancelled()), + export_cancelled, + ) { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(true); + } + return Err(error); + } if rows_exported != last_progress_rows { on_progress(progress(request, rows_exported, ExportStatus::Running, None)); @@ -850,7 +1055,9 @@ async fn try_export_postgres_query_result_stream( if let Some(file) = text_file.as_mut() { file.flush().map_err(|e| format!("Failed to flush text export file: {e}"))?; } - if let Some(writer) = xlsx { + if let Some(writer) = sql_writer.take() { + writer.finish()?; + } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; @@ -940,6 +1147,8 @@ async fn try_export_mysql_query_result_stream( None }; let mut xlsx = None; + let mut sql_writer: Option = + if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let query_timeout = query_export_timeout(request.timeout_secs); let operation_budget = operation_budget_for_pool_key(state, &pool_key, query_timeout).await; let mut conn = crate::db::mysql::get_conn_with_health_check_with_cancel( @@ -1007,7 +1216,9 @@ async fn try_export_mysql_query_result_stream( crate::db::mysql::MySqlQueryStreamItem::Columns { columns: stream_columns, column_types } => { columns = stream_columns; temporal_column_types = column_types.clone(); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.set_columns(columns.clone(), &column_types, request); + } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; } else { @@ -1030,7 +1241,9 @@ async fn try_export_mysql_query_result_stream( &temporal_column_types, request.date_time_format.as_deref(), ); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.write_row(formatted)?; + } else if let Some(file) = text_file.as_mut() { let rows = format_text_export_rows(format, std::slice::from_ref(&formatted)); write!(file, "\n{rows}").map_err(|e| format!("Failed to write export rows: {e}"))?; } else if let Some(writer) = xlsx.as_mut() { @@ -1120,7 +1333,9 @@ async fn try_export_mysql_query_result_stream( if let Some(file) = text_file.as_mut() { file.flush().map_err(|e| format!("Failed to flush text export file: {e}"))?; } - if let Some(writer) = xlsx { + if let Some(writer) = sql_writer.take() { + writer.finish()?; + } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; @@ -1198,6 +1413,8 @@ async fn try_export_clickhouse_query_result_stream( None }; let mut xlsx = None; + let mut sql_writer: Option = + if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let query_timeout = query_export_timeout(request.timeout_secs); let clickhouse_database = if database.is_empty() { "default" } else { database }; @@ -1217,7 +1434,9 @@ async fn try_export_clickhouse_query_result_stream( } => { columns = stream_columns; temporal_column_types = column_types.clone(); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.set_columns(columns.clone(), &column_types, request); + } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; } else { @@ -1240,7 +1459,9 @@ async fn try_export_clickhouse_query_result_stream( &temporal_column_types, request.date_time_format.as_deref(), ); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.write_row(formatted)?; + } else if let Some(file) = text_file.as_mut() { let rows = format_text_export_rows(format, std::slice::from_ref(&formatted)); write!(file, "\n{rows}").map_err(|e| format!("Failed to write export rows: {e}"))?; } else if let Some(writer) = xlsx.as_mut() { @@ -1304,7 +1525,9 @@ async fn try_export_clickhouse_query_result_stream( if let Some(file) = text_file.as_mut() { file.flush().map_err(|e| format!("Failed to flush text export file: {e}"))?; } - if let Some(writer) = xlsx { + if let Some(writer) = sql_writer.take() { + writer.finish()?; + } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; @@ -1363,13 +1586,23 @@ async fn try_export_sqlserver_query_result_stream( None }; let mut xlsx = None; + let mut sql_writer: Option = + if format == "sql" { Some(SqlInsertWriter::create(request)?) } else { None }; let query_timeout = query_export_timeout(request.timeout_secs); let mut client = match cancel_token.as_ref() { Some(token) => { tokio::select! { biased; - _ = token.cancelled() => return Err(canceled_error()), + _ = token.cancelled() => { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(true); + }, guard = client.lock() => guard, } } @@ -1388,7 +1621,9 @@ async fn try_export_sqlserver_query_result_stream( crate::db::sqlserver::SqlServerStreamItem::Columns { columns: stream_columns, column_types } => { columns = stream_columns.to_vec(); temporal_column_types = column_types.to_vec(); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.set_columns(columns.clone(), &temporal_column_types, request); + } else if let Some(file) = text_file.as_mut() { let header = format_text_export_header(format, &columns); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write export header: {e}"))?; } else { @@ -1407,7 +1642,9 @@ async fn try_export_sqlserver_query_result_stream( &temporal_column_types, request.date_time_format.as_deref(), ); - if let Some(file) = text_file.as_mut() { + if let Some(writer) = sql_writer.as_mut() { + writer.write_row(formatted)?; + } else if let Some(file) = text_file.as_mut() { let rows = format_text_export_rows(format, std::slice::from_ref(&formatted)); write!(file, "\n{rows}").map_err(|e| format!("Failed to write export rows: {e}"))?; } else if let Some(writer) = xlsx.as_mut() { @@ -1441,16 +1678,34 @@ async fn try_export_sqlserver_query_result_stream( Ok(()) }, ); - await_stream_with_progress_timeout( + let stream_result = await_stream_with_progress_timeout( stream_future, query_timeout, progress_clock, cancel_token.as_ref(), format!("Query timed out after {} seconds", query_timeout.map_or(0, |timeout| timeout.as_secs())), ) - .await?; + .await; drop(client); + if let Err(error) = stream_result { + let export_cancelled = is_export_cancelled(&request.export_id).await; + if stream_export_was_cancelled( + &error, + cancel_token.as_ref().is_some_and(|token| token.is_cancelled()), + export_cancelled, + ) { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(true); + } + return Err(error); + } + if rows_exported != last_progress_rows { on_progress(progress(request, rows_exported, ExportStatus::Running, None)); } @@ -1458,7 +1713,9 @@ async fn try_export_sqlserver_query_result_stream( if let Some(file) = text_file.as_mut() { file.flush().map_err(|e| format!("Failed to flush text export file: {e}"))?; } - if let Some(writer) = xlsx { + if let Some(writer) = sql_writer.take() { + writer.finish()?; + } else if let Some(writer) = xlsx { let mut buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; @@ -1471,6 +1728,14 @@ async fn try_export_sqlserver_query_result_stream( mod tests { use super::*; + #[test] + fn stream_cancel_detection_covers_driver_token_and_export_flags() { + assert!(stream_export_was_cancelled(QUERY_CANCELED, false, false)); + assert!(stream_export_was_cancelled("driver closed", true, false)); + assert!(stream_export_was_cancelled("driver closed", false, true)); + assert!(!stream_export_was_cancelled("network failure", false, false)); + } + #[test] fn postgres_temp_setup_accepts_only_session_local_table_operations() { let safe = vec![ @@ -1522,6 +1787,8 @@ mod tests { client_session_id: None, execution_id: None, date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, } } @@ -1561,6 +1828,79 @@ mod tests { 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); + } + + #[test] + fn sql_insert_column_types_maps_request_types_to_option_vec() { + let req = request("sql", None, None); + // No export_column_types set → every column becomes None + let result = sql_insert_column_types(&req, &["int4".into(), "text".into()]); + assert_eq!(result, vec![None, None]); + + // Export_column_types provided → null becomes None, Some becomes Some + let mut req = req; + req.export_column_types = Some(vec![Some("int4".into()), None, Some("jsonb".into())]); + let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into()]); + assert_eq!(result, vec![Some("int4".into()), None, Some("jsonb".into())]); + + // Empty string in an override is treated as None + req.export_column_types = Some(vec![Some("".into())]); + let result = sql_insert_column_types(&req, &["int4".into()]); + assert_eq!(result, vec![None]); + } + + #[test] + fn sql_insert_column_types_handles_partial_overrides_gracefully() { + let req = request("sql", None, None); + // Fewer overrides than result columns → extra columns become None + let mut req = req; + req.export_column_types = Some(vec![Some("int4".into()), None]); + let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into(), "bool".into()]); + assert_eq!(result, vec![Some("int4".into()), None, None, None]); + + // More overrides than result columns → extra overrides are ignored + req.export_column_types = Some(vec![ + Some("int4".into()), + Some("text".into()), + Some("json".into()), + Some("bool".into()), + Some("numeric".into()), + ]); + let result = sql_insert_column_types(&req, &["int4".into(), "text".into()]); + assert_eq!(result, vec![Some("int4".into()), Some("text".into())]); + } + + #[test] + fn sql_insert_column_types_handles_all_none_and_all_some() { + let req = request("sql", None, None); + // All None + let mut req = req; + req.export_column_types = Some(vec![None, None, None]); + let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into()]); + assert_eq!(result, vec![None, None, None]); + + // All Some + req.export_column_types = Some(vec![Some("int4".into()), Some("text".into()), Some("json".into())]); + let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into()]); + assert_eq!(result, vec![Some("int4".into()), Some("text".into()), Some("json".into())]); + } + + #[test] + fn sql_export_sql_file_is_initialized_only_for_sql_format() { + // The sql_file variable is initialized at declaration only for "sql" format. + // This is tested by verifying the helper functions produce correct defaults. + let req = request("sql", None, None); + assert_eq!(req.format, "sql"); + assert!(req.export_table_name.is_none()); + assert!(req.export_column_types.is_none()); + } + #[test] fn xlsx_sql_sheet_uses_the_effective_export_sql_and_splits_long_cells() { let mut req = request("xlsx", None, None); diff --git a/crates/dbx-core/tests/live_clickhouse_query_result_export.rs b/crates/dbx-core/tests/live_clickhouse_query_result_export.rs index cb0b94250d..d4aa22e12c 100644 --- a/crates/dbx-core/tests/live_clickhouse_query_result_export.rs +++ b/crates/dbx-core/tests/live_clickhouse_query_result_export.rs @@ -100,6 +100,8 @@ async fn live_clickhouse_query_result_export_xlsx_streams_random_order_query_onc client_session_id: None, execution_id: Some(format!("live-clickhouse-query-export-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let done_seen = AtomicBool::new(false); diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 67961a3f3d..52beda31d1 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -168,6 +168,8 @@ async fn live_mysql_query_result_export_xlsx_streams_single_query_without_duplic client_session_id: None, execution_id: Some(format!("live-mysql-query-export-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let done_seen = AtomicBool::new(false); @@ -254,6 +256,8 @@ async fn live_mysql_xlsx_export_can_outlive_query_timeout_while_rows_keep_arrivi client_session_id: None, execution_id: Some(format!("live-mysql-query-export-timeout-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let rows_exported = AtomicU64::new(0); diff --git a/crates/dbx-core/tests/live_postgres_query_result_export.rs b/crates/dbx-core/tests/live_postgres_query_result_export.rs index 19a1ff6d6d..c0ddcb75aa 100644 --- a/crates/dbx-core/tests/live_postgres_query_result_export.rs +++ b/crates/dbx-core/tests/live_postgres_query_result_export.rs @@ -135,6 +135,8 @@ async fn live_postgres_query_result_export_uses_single_streamed_query() { client_session_id: None, execution_id: Some(format!("live-postgres-query-export-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let done_seen = AtomicBool::new(false); @@ -211,6 +213,8 @@ async fn live_postgres_query_result_xlsx_preserves_temporal_cell_types() { client_session_id: Some(format!("live-postgres-xlsx-temporal-{suffix}")), execution_id: Some(format!("live-postgres-xlsx-temporal-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; @@ -280,6 +284,8 @@ async fn live_postgres_truncated_batch_result_export_replays_safe_temp_setup() { client_session_id: Some(format!("temp-export-csv-{short_suffix}")), execution_id: Some(format!("temp-export-csv-{short_suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let csv_rows = AtomicU64::new(0); @@ -352,6 +358,8 @@ async fn live_postgres_xlsx_export_can_outlive_query_timeout_while_rows_keep_arr client_session_id: None, execution_id: Some(format!("live-postgres-query-export-timeout-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let rows_exported = AtomicU64::new(0); @@ -415,6 +423,8 @@ async fn live_postgres_stream_still_times_out_without_progress_and_recovers() { client_session_id: None, execution_id: Some(format!("live-postgres-query-export-stall-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let started_at = Instant::now(); diff --git a/crates/dbx-core/tests/live_sqlserver_completion.rs b/crates/dbx-core/tests/live_sqlserver_completion.rs index b146a5fbcc..17889fa974 100644 --- a/crates/dbx-core/tests/live_sqlserver_completion.rs +++ b/crates/dbx-core/tests/live_sqlserver_completion.rs @@ -450,6 +450,8 @@ async fn live_sqlserver_query_result_export_streams_cte_query_to_csv() { client_session_id: None, execution_id: Some(format!("live-sqlserver-export-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let done_seen = AtomicBool::new(false); diff --git a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs index 12fa2e6ef5..2d37d65ea9 100644 --- a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs +++ b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs @@ -110,6 +110,8 @@ async fn live_sqlserver_xlsx_export_can_outlive_query_timeout_while_rows_keep_ar client_session_id: None, execution_id: Some(format!("live-sqlserver-xlsx-{suffix}")), date_time_format: None, + export_table_name: None, + export_column_types: None, numeric_column_right_align: false, }; let rows_exported = AtomicU64::new(0); diff --git a/docs/superpowers/specs/2026-07-25-query-result-sql-export-background-task-design.md b/docs/superpowers/specs/2026-07-25-query-result-sql-export-background-task-design.md new file mode 100644 index 0000000000..6648f2014f --- /dev/null +++ b/docs/superpowers/specs/2026-07-25-query-result-sql-export-background-task-design.md @@ -0,0 +1,380 @@ +# Query Result SQL INSERT Async Background Export + +Date: 2026-07-25 +Author: dbx team +Status: Draft + +## Problem + +点击"导出当前结果集全部数据 SQL INSERT"后，UI 会冻结数十秒甚至一分钟（4 万行 ~20s，10 万行 ~1min+），然后才弹出保存对话框。用户以为功能坏了。 + +## Root Cause + +当前 `exportSql()` 对 query-result 场景走 `resultToExport()` → 全量数据先拉到前端内存 → `formatSqlInsert()` 通过 Tauri IPC 发给 Rust 生成 INSERT → 再回前端 → `saveTextFile()` 写入。整个过程阻塞 UI 线程，保存对话框出现在最后。 + +**对比 CSV/XLSX/TXT 的后台流式导出**（右键表名导出、工具栏导出 query result）： +1. 先弹保存对话框（立即响应） +2. 注册后台任务（进度显示在 toolbar popover） +3. Rust 端流式分页读取 DB → 直接写入文件 +4. 完成后 toast 通知 + +## Design + +### 方案 + +在 `export_query_result_core_inner()` 的通用 JDBC 分页循环中添加 `"sql"` 格式支持。复用 `build_export_insert_statements()`，按 100 行一批 flush 到文件。 + +**不碰** Postgres/MySQL/ClickHouse/SQL Server 四个原生流式路径——SQL 格式只走分页循环（原生路径检测到 `format == "sql"` 返回 `false` 后回退）。 + +前端 `exportSql()` 增加 Step 2：先尝试 `exportQueryResultSqlViaBackend()` 走后台流式导出，失败/Web 才回退到当前本地导出。 + +### 修复的 Blocker + +1. **复用 `buildQueryResultExportRequest`**（`queryResultExportRequest` callback），不手撸 request 对象。保留 `rowLimit`、`totalRows`、`timeoutSecs`、`useAgentCursor`、`setupSql`、`clientSessionId`、`executionId`、`keysetOptimizationEnabled` 等字段。 +2. **`export_table_name` 兜底**：Rust 侧当 `export_table_name` 为空时回退到 `"query_result"`，纯 `SELECT` 查询不再报错。 +3. **启动失败处理**：`startQueryResultExport()` rejecting 时将后台任务标记为 `Error`，清理 cancel handler，防止任务卡在 `Running` 状态。 + +## Files Changed + +| File | What | Lines | +|------|------|-------| +| `crates/dbx-core/src/query_result_export.rs` | 格式守卫放松 + 分页循环 SQL 分支 + 新字段 + helper | ~80 | +| `apps/desktop/src/composables/useDataGridExport.ts` | 新增 `exportQueryResultSqlViaBackend`，`exportSql()` 增加 Step 2 | ~50 | +| `apps/desktop/src/components/grid/DataGrid.vue` | `queryResultExportRequest` prop 类型加 `"sql"` 和两个 new fields | ~5 | +| `apps/desktop/src/components/layout/ContentArea.vue` | 同 DataGrid.vue 类型更新 | ~3 | +| `apps/desktop/src/stores/queryStore.ts` | `BuildQueryResultExportRequestOptions` 加 `exportTableName`、`exportColumnTypes` | ~5 | +| `packages/app-tests/useDataGridExport.test.ts` | 测试适配（TXT/SQL Web fallback） | ~5 | + +**Total: ~148 lines, zero new files, zero new dependencies.** + +## Detailed Design + +### 1. Rust: `QueryResultExportRequest` — new optional fields + +File: `crates/dbx-core/src/query_result_export.rs` + +```rust + #[serde(default, skip_serializing_if = "Option::is_none")] + pub date_time_format: Option, + // -- new fields -- + #[serde(default, skip_serializing_if = "Option::is_none")] + pub export_table_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub export_column_types: Option>, +} +``` + +### 2. Rust: format guard + +Relax the existing guard to accept `"sql"`: + +```rust +if format != "csv" && format != "xlsx" && format != "txt" && format != "sql" { + return Err(format!("Unsupported streaming query-result export format: {format}")); +} +``` + +### 3. Rust: effective_row_limit + +SQL format has no XLSX row cap: + +```rust +fn effective_row_limit(format: &str, request: &QueryResultExportRequest) -> Option { + if format == "xlsx" { + Some(request.row_limit.map_or(XLSX_MAX_DATA_ROWS, |limit| limit.min(XLSX_MAX_DATA_ROWS))) + } else { + request.row_limit + } +} +``` + +### 4. Rust: native stream skip + +原生流式路径（Postgres/MySQL/ClickHouse/SQL Server）只支持 CSV/TXT/XLSX。在 `export_query_result_core_inner` 中跳过它们对 `"sql"` 的调用，改为直接走通用分页循环： + +```rust +// Before: 对每种格式都尝试原生流 +// After: SQL 格式跳过原生流路径 +if format != "sql" { + if try_export_postgres_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? { return Ok(()); } + if try_export_sqlserver_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? { return Ok(()); } + if try_export_mysql_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? { return Ok(()); } + if try_export_clickhouse_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? { return Ok(()); } +} +``` + +1 行变更（在原有 `if` 外包装一个 `format != "sql"` 条件），避免了修改 4 个原生函数的 60+ 行。 + +### 5. Rust: pagination loop — `"sql"` branch + +Inserted alongside existing CSV/TXT branch. Uses a `pending_rows` buffer that flushes via `build_export_insert_statements` every 100 rows. + +**New imports:** + +```rust +use std::mem; +use crate::database_export::{build_export_insert_statements, BuildExportInsertStatementsOptions}; +``` + +**New constant:** + +```rust +const SQL_INSERT_BATCH_SIZE: usize = 100; +``` + +**New helper:** + +```rust +fn sql_insert_column_types(request: &QueryResultExportRequest, column_types: &[String]) -> Vec> { + request.export_column_types + .as_ref() + .map(|types| types.iter().map(|t| if t.is_empty() { None } else { Some(t.clone()) }).collect()) + .unwrap_or_else(|| vec![None; column_types.len()]) +} +``` + +**New variables** (alongside existing `columns`/`column_types`): + +```rust +let mut sql_file: Option> = None; +let mut pending_rows: Vec> = Vec::new(); +let mut sql_insert_col_types: Vec> = Vec::new(); +``` + +**Initialize types** (inside existing `if columns.is_empty()` block): + +```rust +if columns.is_empty() { + columns = result.columns.clone(); + column_types = result.column_types.clone(); + sql_insert_col_types = sql_insert_column_types(request, &column_types); +} +``` + +**Write branch** (after the `if format == "csv" || format == "txt"` block): + +```rust +} else if format == "sql" { + if sql_file.is_none() { + sql_file = Some(BufWriter::new( + File::create(&request.file_path) + .map_err(|e| format!("Failed to create SQL file: {e}"))?, + )); + } + pending_rows.extend(formatted_rows); + if pending_rows.len() >= SQL_INSERT_BATCH_SIZE { + let col_types = sql_insert_col_types.clone(); + let table_name = request.export_table_name + .as_deref() + .filter(|n| !n.trim().is_empty()) + .unwrap_or("query_result"); + let stmts = build_export_insert_statements(BuildExportInsertStatementsOptions { + database_type: Some(request.database_type), + schema: request.schema.clone(), + table_name: Some(table_name.to_string()), + qualified_table_name: None, + columns: columns.clone(), + column_types: col_types, + column_extras: Vec::new(), + rows: mem::take(&mut pending_rows), + batch_size: Some(SQL_INSERT_BATCH_SIZE), + })?; + let file = sql_file.as_mut().unwrap(); + for stmt in &stmts { + writeln!(file, "{stmt}").map_err(|e| format!("Failed to write SQL: {e}"))?; + } + } +} +``` + +**Post-loop flush** (alongside CSV/TXT flush): + +```rust +} else if format == "sql" { + if !pending_rows.is_empty() { + let col_types = sql_insert_col_types.clone(); + let table_name = request.export_table_name + .as_deref() + .filter(|n| !n.trim().is_empty()) + .unwrap_or("query_result"); + let stmts = build_export_insert_statements(BuildExportInsertStatementsOptions { + database_type: Some(request.database_type), + schema: request.schema.clone(), + table_name: Some(table_name.to_string()), + qualified_table_name: None, + columns: columns.clone(), + column_types: col_types, + column_extras: Vec::new(), + rows: mem::take(&mut pending_rows), + batch_size: Some(SQL_INSERT_BATCH_SIZE), + })?; + for stmt in &stmts { + writeln!(sql_file.as_mut().unwrap(), "{stmt}") + .map_err(|e| format!("Failed to write SQL: {e}"))?; + } + } + if let Some(file) = sql_file.as_mut() { + file.flush().map_err(|e| format!("Failed to flush SQL file: {e}"))?; + } +} +``` + +**File open output**: the `File::create` overwrites by design — same behavior as all other export formats. + +### 6. Frontend: Types + +**queryStore.ts:** + +```typescript +interface BuildQueryResultExportRequestOptions { + exportId: string; + filePath: string; + format: "csv" | "xlsx" | "txt" | "sql"; // ← "sql" 新增 + includeSqlSheet?: boolean; + exportTableName?: string; // ← new + exportColumnTypes?: Array; // ← new +} +``` + +**DataGrid.vue prop type:** + +```typescript +queryResultExportRequest?: (options: { + exportId: string; + filePath: string; + format: "csv" | "xlsx" | "txt" | "sql"; // ← "sql" 新增 + includeSqlSheet?: boolean; + exportTableName?: string; // ← new + exportColumnTypes?: Array; // ← new +}) => Promise; +``` + +### 6. Frontend: `exportQueryResultSqlViaBackend` + +New function in `useDataGridExport.ts`. + +**新增 import** (文件头部): + +```typescript +import { useExportTracker } from "@/composables/useExportTracker"; +``` + +**新增 destructure** (在 `useDataGridExport` 函数体开头): + +```typescript +const { addTask, updateTableExportTask, registerTaskCancelHandler, unregisterTaskCancelHandler } = useExportTracker(); +``` + +```typescript +async function exportQueryResultSqlViaBackend(rowIds?: number[]): Promise { + // Guard: 仅用于 query-result、无完整结果、桌面端 + if (rowIds !== undefined || context.value !== "results" || !queryResultExportRequest) return false; + if (hasCompleteLocalResult?.value) return false; + if (!isTauriRuntime()) return false; // Web fall through → local export + + // 1. 保存对话框 FIRST（立即响应） + let outputPath = exportFileName("query-result", "sql"); + const { save } = await import("@tauri-apps/plugin-dialog"); + const path = await save({ + defaultPath: outputPath, + filters: [{ name: "SQL", extensions: ["sql"] }], + }); + if (!path) return true; + outputPath = path as string; + + // 2. 通过 buildQueryResultExportRequest 构建完整 request + const exportId = uuid(); + const request = await queryResultExportRequest({ + exportId, + filePath: outputPath, + format: "sql", + exportTableName: tableMeta.value?.tableName, + exportColumnTypes: columnTypes.value?.map( + t => t ?? null + ) as Array | undefined, + }); + if (!request) throw new Error("Unable to build query result export request"); + + // 3. 注册后台任务 + cancel handler + registerTaskCancelHandler(exportId, () => api.cancelQueryResultExport(exportId, request.executionId)); + addTask(tableMeta.value?.tableName || "Query Result", "sql", outputPath, exportId); + + try { + await api.startQueryResultExport(request, (progress) => { + updateTableExportTask(exportId, progress); + if (progress.status === "Done") { + toast(t("grid.exported")); + } + }); + } catch (e) { + // 4. 启动失败 → 标记 Error（修复 Blocker 3） + updateTableExportTask(exportId, { + exportId, + tableName: tableMeta.value?.tableName || "Query Result", + rowsExported: 0, + totalRows: null, + status: "Error" as const, + errorMessage: e?.message || String(e), + }); + throw e; + } finally { + unregisterTaskCancelHandler(exportId); + } + return true; +} +``` + +### 7. Frontend: `exportSql()` — new middle step + +```typescript +async function exportSql(rowIds?: number[]) { + await runExclusiveExport(async () => { + try { + // Step 1: table-data context — 已有后台表导出（不变） + if (await exportFullTableDataViaBackend("sql", rowIds)) return; + + // Step 2: query-result context — NEW 后台流式导出 + if (await exportQueryResultSqlViaBackend(rowIds)) return; + + // Step 3: fallback — 本地导出（Web 和异常场景） + const result = await resultToExport(rowIds, undefined, true, false); + const exportData = sqlInsertExportData(result); + const content = await formatSqlInsert({ ... }); + await saveTextFile(content, exportFileName(..., "sql", ...), "SQL", "sql"); + toast(t("grid.exported")); + } catch (e: any) { + toast(t("grid.exportFailed", { message: e?.message || String(e) }), 5000); + } + }); +} +``` + +## User Experience Flow + +``` +1. 用户点击 "导出当前结果集全部数据 SQL INSERT" +2. 系统保存对话框立即弹出 ← 无冻结 +3. 用户选择路径 → 点击保存 +4. 后台任务指示器在 toolbar 出现（带进度条） +5. Export 分页流式执行：DB → 分批读取 → build INSERT → 写入文件 +6. 完成 → toast "已导出" + 或 + 失败 → 后台任务显示 Error 状态 +``` + +## Backward Compatibility + +- 新增字段均为 `Option` + `skip_serializing_if`，现有 caller 无损 +- SQL INSERT 输出格式与现有 `build_export_insert_statements` 完全一致（batch 从 1→100，语义等价） +- 纯查询兜底表名 `"query_result"`，导出文件内容与之前不同（之前失败），属于 bugfix 而非 regression +- Web 端保持不变，仍走本地导出（Blob download） + +## Error Handling + +| Scenario | Behavior | +|----------|----------| +| `startQueryResultExport` 拒绝启动 | catch 中标记 task Error + 清理 cancel handler | +| 文件创建失败 | 错误传播到前端，task 标记 Error | +| 磁盘满 | `std::io::Error` → 游标中断 → task Error | +| 用户取消 | CancellationToken → OnProgress(Cancelled) → 清理文件 | +| SQL 表单 export_table_name 为空 | Rust 回退到 `"query_result"` | diff --git a/packages/app-tests/useDataGridExport.test.ts b/packages/app-tests/useDataGridExport.test.ts index 2ac12bbb61..0b85be8775 100644 --- a/packages/app-tests/useDataGridExport.test.ts +++ b/packages/app-tests/useDataGridExport.test.ts @@ -23,10 +23,13 @@ const apiMock = vi.hoisted(() => ({ const clipboardMock = vi.hoisted(() => ({ copyToClipboard: vi.fn(), })); +const runtimeMock = vi.hoisted(() => ({ isTauri: false })); +const dialogMock = vi.hoisted(() => ({ save: vi.fn() })); vi.mock("@/lib/backend/api", () => apiMock); vi.mock("@/lib/common/clipboard", () => clipboardMock); -vi.mock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); +vi.mock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => runtimeMock.isTauri })); +vi.mock("@tauri-apps/plugin-dialog", () => ({ save: dialogMock.save })); vi.mock("@/composables/useToast", () => ({ useToast: () => ({ toast: vi.fn() }) })); vi.mock("vue-i18n", () => ({ useI18n: () => ({ t: (key: string) => key }) })); @@ -79,6 +82,7 @@ function buildExportHarness( exportFileBaseName?: string; columns?: string[]; columnTypes?: Array; + allColumnTypes?: Array; rows?: QueryResult["rows"]; allExportResults?: Array<{ sheetName: string; result: QueryResult; sql?: string }>; completeLocalResult?: QueryResult; @@ -109,7 +113,7 @@ function buildExportHarness( const fullExportResult = vi.fn(async () => { throw new Error("fullExportResult should not be called for streaming CSV/XLSX query exports"); }); - const queryResultExportRequest = vi.fn(async (options: { exportId: string; filePath: string; format: "csv" | "xlsx" | "txt"; includeSqlSheet?: boolean }) => ({ + const queryResultExportRequest = vi.fn(async (options: { exportId: string; filePath: string; format: "csv" | "xlsx" | "txt" | "sql"; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array }) => ({ exportId: options.exportId, connectionId: "conn-1", database: "db", @@ -121,6 +125,8 @@ function buildExportHarness( filePath: options.filePath, format: options.format, includeSqlSheet: options.includeSqlSheet, + exportTableName: options.exportTableName, + exportColumnTypes: options.exportColumnTypes, pageSize: 1000, rowLimit: 100000, totalRows: 2, @@ -142,6 +148,7 @@ function buildExportHarness( context: computed(() => options.context ?? "results"), sourceColumns: computed(() => options.sourceColumns), columnTypes: computed(() => options.columnTypes), + allColumnTypes: computed(() => options.allColumnTypes), whereInput: computed(() => undefined), orderBy: computed(() => undefined), exportBatchSize: computed(() => 1000), @@ -236,6 +243,8 @@ function buildTableDataExportHarness() { beforeEach(() => { setActivePinia(createPinia()); vi.clearAllMocks(); + runtimeMock.isTauri = false; + dialogMock.save.mockResolvedValue(null); clipboardMock.copyToClipboard.mockResolvedValue(undefined); apiMock.startQueryResultExport.mockImplementation(async (_request, onProgress) => { onProgress({ exportId: _request.exportId, tableName: "", rowsExported: 2, totalRows: 2, status: "Done" }); @@ -430,6 +439,22 @@ test("MySQL joined query SQL export keeps result aliases instead of source colum } }); +test("background SQL export keeps full result column types when visible columns are reordered", async () => { + runtimeMock.isTauri = true; + dialogMock.save.mockResolvedValue("/tmp/query-result.sql"); + const { composable, queryResultExportRequest } = buildExportHarness({ + columns: ["payload", "created_at"], + columnTypes: ["jsonb", "timestamp"], + allColumnTypes: ["bytea", "jsonb", "timestamp"], + tableMeta: { tableName: "events", primaryKeys: ["id"] }, + }); + + await composable.exportSql(); + + assert.deepEqual(queryResultExportRequest.mock.calls[0][0].exportColumnTypes, ["bytea", "jsonb", "timestamp"]); + assert.deepEqual(apiMock.startQueryResultExport.mock.calls[0][0].exportColumnTypes, ["bytea", "jsonb", "timestamp"]); +}); + test("table data SQL export keeps source column names", async () => { const download = installTextDownloadCapture(); apiMock.buildExportSqlInsert.mockResolvedValueOnce("INSERT INTO `users` (`id`, `name`) VALUES (1, 'Ada');"); diff --git a/src-tauri/src/commands/query_result_export.rs b/src-tauri/src/commands/query_result_export.rs index 15eb46418e..3f61e0e866 100644 --- a/src-tauri/src/commands/query_result_export.rs +++ b/src-tauri/src/commands/query_result_export.rs @@ -1,6 +1,7 @@ +use std::path::PathBuf; use std::sync::{ atomic::{AtomicBool, Ordering}, - Arc, + Arc, Mutex, }; use tauri::{AppHandle, Emitter, State}; @@ -16,15 +17,137 @@ fn emit_progress(app: &AppHandle, progress: TableExportProgress) { let _ = app.emit("query-result-export-progress", progress); } +fn route_core_progress( + progress: TableExportProgress, + deferred_done: &Mutex>, + cancelled: &AtomicBool, + emit: impl FnOnce(TableExportProgress), +) { + match progress.status { + ExportStatus::Done => { + *deferred_done.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(progress); + } + ExportStatus::Cancelled => { + cancelled.store(true, Ordering::SeqCst); + emit(progress); + } + _ => emit(progress), + } +} + +/// Build a temp-file path alongside the target path. +/// The temp file has a `.dbx-export-tmp` suffix so it's never confused with the +/// user's chosen file. On success the temp file is renamed atomically onto the +/// target; on error/cancel only the temp file is cleaned up, leaving the user's +/// chosen path untouched. +fn temp_file_path(target: &str) -> (PathBuf, PathBuf) { + let target_path = PathBuf::from(target); + let mut temp_path = target_path.clone(); + let mut temp_name = + target_path.file_name().map(|name| name.to_os_string()).unwrap_or_else(|| std::ffi::OsString::from("export")); + temp_name.push(".dbx-export-tmp"); + temp_path.set_file_name(temp_name); + (target_path, temp_path) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn temp_file_path_adds_suffix_to_simple_filename() { + let (target, temp) = temp_file_path("C:\\exports\\backup.sql"); + assert_eq!(target, PathBuf::from("C:\\exports\\backup.sql")); + assert_eq!(temp, PathBuf::from("C:\\exports\\backup.sql.dbx-export-tmp")); + } + + #[test] + fn temp_file_path_adds_suffix_to_filename_with_multiple_extensions() { + let (target, temp) = temp_file_path("/home/user/query-result.xlsx"); + assert_eq!(target, PathBuf::from("/home/user/query-result.xlsx")); + assert_eq!(temp, PathBuf::from("/home/user/query-result.xlsx.dbx-export-tmp")); + } + + #[test] + fn temp_file_path_handles_filename_without_extension() { + let (target, temp) = temp_file_path("/tmp/export"); + assert_eq!(target, PathBuf::from("/tmp/export")); + assert_eq!(temp, PathBuf::from("/tmp/export.dbx-export-tmp")); + } + + #[test] + fn temp_file_path_uses_export_fallback_for_empty_target() { + let (target, temp) = temp_file_path(""); + assert_eq!(target, PathBuf::from("")); + // file_name() on a path without a filename returns None, so we fall + // back to "export" as the base name. + assert!(temp.to_string_lossy().ends_with("export.dbx-export-tmp")); + } + + #[test] + fn route_core_progress_defers_done_until_finalization() { + let deferred_done = Mutex::new(None); + let cancelled = AtomicBool::new(false); + let emitted = Mutex::new(Vec::new()); + let done = TableExportProgress { + export_id: "export-1".to_string(), + table_name: String::new(), + rows_exported: 42, + total_rows: Some(42), + status: ExportStatus::Done, + error_message: None, + }; + + route_core_progress(done, &deferred_done, &cancelled, |progress| { + emitted.lock().unwrap().push(progress); + }); + + assert!(emitted.lock().unwrap().is_empty()); + assert_eq!(deferred_done.lock().unwrap().as_ref().map(|progress| progress.rows_exported), Some(42)); + assert!(!cancelled.load(Ordering::SeqCst)); + } + + #[test] + fn route_core_progress_keeps_cancelled_terminal_state() { + let deferred_done = Mutex::new(None); + let cancelled = AtomicBool::new(false); + let emitted = Mutex::new(Vec::new()); + let progress = TableExportProgress { + export_id: "export-1".to_string(), + table_name: String::new(), + rows_exported: 7, + total_rows: None, + status: ExportStatus::Cancelled, + error_message: Some("Export cancelled".to_string()), + }; + + route_core_progress(progress, &deferred_done, &cancelled, |progress| { + emitted.lock().unwrap().push(progress); + }); + + assert!(deferred_done.lock().unwrap().is_none()); + assert!(cancelled.load(Ordering::SeqCst)); + assert!(matches!( + emitted.lock().unwrap().as_slice(), + [TableExportProgress { status: ExportStatus::Cancelled, .. }] + )); + } +} + #[tauri::command] pub async fn start_query_result_export( app: AppHandle, state: State<'_, Arc>, - request: QueryResultExportRequest, + mut request: QueryResultExportRequest, ) -> Result<(), String> { let state = state.inner().clone(); let export_id = request.export_id.clone(); - let file_path = request.file_path.clone(); + + // Redirect file I/O to a temp file so the user's chosen path is never + // truncated before the query completes. On success the temp file is + // renamed onto the target; on error/cancel only the temp file is removed. + let (target_path, temp_path) = temp_file_path(&request.file_path); + request.file_path = temp_path.to_string_lossy().to_string(); tokio::spawn(async move { let execution_id = request.execution_id.clone().filter(|id| !id.trim().is_empty()); @@ -41,18 +164,20 @@ pub async fn start_query_result_export( let cancel_token = registered_query.as_ref().map(|query| query.token()); let cancelled = Arc::new(AtomicBool::new(false)); let cancelled_progress = cancelled.clone(); + let deferred_done = Arc::new(Mutex::new(None)); + let deferred_done_progress = deferred_done.clone(); let result = dbx_core::query_result_export::export_query_result_core(&state, &request, cancel_token, |progress| { - if matches!(progress.status, ExportStatus::Cancelled) { - cancelled_progress.store(true, Ordering::SeqCst); - } - emit_progress(&app, progress); + route_core_progress(progress, &deferred_done_progress, &cancelled_progress, |progress| { + emit_progress(&app, progress); + }); }) .await; drop(registered_query); + let completed_progress = deferred_done.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).take(); if let Err(e) = result { - let _ = tokio::fs::remove_file(&file_path).await; + let _ = tokio::fs::remove_file(&request.file_path).await; emit_progress( &app, TableExportProgress { @@ -65,7 +190,39 @@ pub async fn start_query_result_export( }, ); } else if cancelled.load(Ordering::SeqCst) { - let _ = tokio::fs::remove_file(&file_path).await; + let _ = tokio::fs::remove_file(&request.file_path).await; + } else if let Some(progress) = completed_progress { + // Success: atomically rename temp → target so the user's chosen + // file is only created/replaced when the export fully completes. + if let Err(e) = std::fs::rename(&request.file_path, &target_path) { + let _ = tokio::fs::remove_file(&request.file_path).await; + emit_progress( + &app, + TableExportProgress { + export_id: export_id.clone(), + table_name: String::new(), + rows_exported: 0, + total_rows: None, + status: ExportStatus::Error, + error_message: Some(format!("Failed to finalize export file: {e}")), + }, + ); + } else { + emit_progress(&app, progress); + } + } else { + let _ = tokio::fs::remove_file(&request.file_path).await; + emit_progress( + &app, + TableExportProgress { + export_id: export_id.clone(), + table_name: String::new(), + rows_exported: 0, + total_rows: None, + status: ExportStatus::Error, + error_message: Some("Export finished without a completion status".to_string()), + }, + ); } dbx_core::database_export::clear_export_cancelled(&export_id).await;