diff --git a/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts b/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts index 90d88eff98..17faa997d3 100644 --- a/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts @@ -35,6 +35,7 @@ vi.mock("vue-i18n", () => ({ })); import { useSidebarTreeExportRuntime } from "@/composables/useSidebarTreeExportRuntime"; +import { DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS } from "@/lib/dataGrid/dataGridCopyExtractor"; import { showStructurePreviewDialog, structurePreviewSql, structurePreviewTitle } from "@/components/sidebar/sidebarTreeDialogState"; function functionBody(name: string): string { @@ -78,6 +79,7 @@ function exportSettings() { exportBatchSize: 128, exportRowLimit: 500, exportRowLimitEnabled: true, + dataGridExtractorOptions: DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS, }, }; } diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index c479edc52f..d2b144adcf 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -1170,6 +1170,24 @@ export function useDataGridExport(options: UseDataGridExportOptions) { await exportAllResultsXlsxResult(true); } + /** + * SQL 导出需要排除的主键列：跟随数据提取设置里的“排除主键”，与 + * copy-as-INSERT (e3d5d1418) 一致只剔除自增/identity 主键，手动赋值的 + * 主键保留，否则回放 INSERT 会缺值；只有拿得到表元数据时才生效。 + */ + function sqlExportExcludedColumns(): string[] | undefined { + if (!extractorOptionsOption?.value?.sql.excludePrimaryKeysFromInsert) return undefined; + const primaryKeys = tableMeta.value?.primaryKeys ?? []; + const autoGenerated = primaryKeys.filter((column) => isAutoGeneratedColumn(column, tableMeta.value)); + return autoGenerated.length > 0 ? autoGenerated : undefined; + } + + /** 供导出请求使用：把“不含主键”设置转换成后端请求字段。 */ + function sqlExportPrimaryKeyOptions(): { excludePrimaryKeys?: boolean; primaryKeys?: string[] } { + const excludeColumns = sqlExportExcludedColumns(); + return excludeColumns ? { excludePrimaryKeys: true, primaryKeys: excludeColumns } : {}; + } + async function exportFullTableDataViaBackend(format: "csv" | "xlsx" | "json" | "markdown" | "sql" | "txt", rowIds?: number[], headerMode: XlsxHeaderMode = "name", autoFilter = true, sqlExportOptions?: SqlExportOptions): Promise { const meta = tableMeta.value; // The backend table exporter currently builds two-part table names. External @@ -1236,6 +1254,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { columnTypes: columnTypes.value, columnComments: format === "xlsx" ? buildXlsxHeaderOverrides(columns.value, visibleXlsxColumnComments.value, headerMode) : undefined, primaryKeys: meta.primaryKeys, + ...sqlExportPrimaryKeyOptions(), whereInput: whereInput.value, orderBy: orderBy.value, skipCount: false, @@ -1309,6 +1328,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { ? { ...baseRequest, csvQuoteMode: useSettingsStore().editorSettings.csvQuoteMode, + ...sqlExportPrimaryKeyOptions(), dateTimeFormat: useSettingsStore().editorSettings.globalDateTimeExportFormat || undefined, numericColumnRightAlign: useSettingsStore().editorSettings.numericColumnRightAlign ?? true, columnComments, @@ -1475,6 +1495,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { spatialValues: exportData.spatialValues, rows: exportData.rows, insertMode, + excludeColumns: sqlExportExcludedColumns(), }); logExportStage( "sql-build-done", @@ -1516,6 +1537,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { spatialValues: exportData.spatialValues, rows: exportData.rows, insertMode, + excludeColumns: sqlExportExcludedColumns(), }); await saveTextFile(content, exportFileName("export-page", "sql", { page: true }), "SQL", "sql"); toast(t("grid.exported")); diff --git a/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts b/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts index 720eeec31a..17c7cae0a5 100644 --- a/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts +++ b/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts @@ -445,6 +445,7 @@ export function useSidebarTreeExportRuntime(options: SidebarTreeExportRuntimeOpt columnComments, autoFilter: format === "xlsx" ? autoFilter : undefined, primaryKeys, + excludePrimaryKeys: settingsStore.editorSettings.dataGridExtractorOptions.sql.excludePrimaryKeysFromInsert === true, batchSize: target.batchSize, skipCount: format === "sql", rowLimit: target.rowLimit, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 589797c735..ec7851316d 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -5365,6 +5365,8 @@ export interface TableExportRequest { columnTypes?: Array; columnComments?: Array | null; primaryKeys?: string[]; + /** 导出 SQL 时是否排除主键列（对应数据提取设置里的“排除主键”）。 */ + excludePrimaryKeys?: boolean; whereInput?: string; orderBy?: string; skipCount?: boolean; @@ -5427,6 +5429,10 @@ export interface QueryResultExportRequest { columnComments?: Array | null; autoFilter?: boolean; identifierQuote?: string; + /** 导出 SQL 时是否排除主键列（对应数据提取设置里的“排除主键”）。 */ + excludePrimaryKeys?: boolean; + /** 结果集对应的原表主键列名，由前端从表元数据带入。 */ + primaryKeys?: string[]; } export async function startTableExport(request: TableExportRequest, onProgress: (progress: TableExportProgress) => void): Promise { diff --git a/apps/desktop/src/lib/export/databaseExport.ts b/apps/desktop/src/lib/export/databaseExport.ts index c503b8810b..db43f586bd 100644 --- a/apps/desktop/src/lib/export/databaseExport.ts +++ b/apps/desktop/src/lib/export/databaseExport.ts @@ -47,6 +47,8 @@ export interface BuildExportInsertStatementsOptions { spatialColumns?: QueryResult["spatial_columns"]; spatialValues?: QueryResult["spatial_values"]; rows: QueryResult["rows"]; + /** 生成 INSERT 时需要排除的列名（例如导出时不带主键），忽略大小写匹配。 */ + excludeColumns?: string[]; batchSize?: number; } diff --git a/apps/desktop/src/lib/export/exportFormats.ts b/apps/desktop/src/lib/export/exportFormats.ts index 7c66b24899..404605a6d5 100644 --- a/apps/desktop/src/lib/export/exportFormats.ts +++ b/apps/desktop/src/lib/export/exportFormats.ts @@ -40,6 +40,7 @@ export interface FormatSqlInsertOptions { spatialValues?: QueryResult["spatial_values"]; rows: ExportCellValue[][]; insertMode?: SqlInsertMode; + excludeColumns?: string[]; } export function formatSqlInsert({ insertMode = "batch", ...options }: FormatSqlInsertOptions): Promise { diff --git a/crates/dbx-core/src/data_grid_sql.rs b/crates/dbx-core/src/data_grid_sql.rs index 2cf24f35d1..5a08101ede 100644 --- a/crates/dbx-core/src/data_grid_sql.rs +++ b/crates/dbx-core/src/data_grid_sql.rs @@ -3250,16 +3250,17 @@ pub(crate) fn is_neo4j_element_id(database_type: Option, name: Opt database_type == Some(DatabaseType::Neo4j) && name == Some(DBX_NEO4J_ELEMENT_ID_COLUMN) } -pub(crate) fn is_auto_generated_column(column: &DataGridColumnInfo) -> bool { - column - .extra - .as_deref() - .unwrap_or("") +pub(crate) fn extra_is_auto_generated(extra: &str) -> bool { + extra .to_ascii_lowercase() .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_') .any(|part| matches!(part, "auto_increment" | "autoincrement" | "identity")) } +pub(crate) fn is_auto_generated_column(column: &DataGridColumnInfo) -> bool { + extra_is_auto_generated(column.extra.as_deref().unwrap_or("")) +} + fn grid_value_is_empty(value: &Value) -> bool { value.is_null() || value.as_str().is_some_and(str::is_empty) } diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 59e4140385..33ff108ee4 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -473,6 +473,9 @@ pub struct BuildExportInsertStatementsOptions { pub struct BuildExportSqlInsertOptions { #[serde(flatten)] pub insert: BuildExportInsertStatementsOptions, + /// 生成 INSERT 时需要排除的列名（例如导出时不带主键），忽略大小写匹配。 + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub exclude_columns: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -1154,10 +1157,24 @@ fn is_export_numeric_literal(text: &str) -> bool { } pub fn build_export_insert_statements(options: BuildExportInsertStatementsOptions) -> Result, String> { + build_export_insert_statements_excluding(options, &[]) +} + +/// 与 [`build_export_insert_statements`] 行为一致，但可以额外按列名排除若干列 +/// （典型场景：导出 SQL 时不带主键）。列名比较忽略大小写。 +/// +/// 行值与空间列仍然按列在 `columns` 中的原始下标取值，所以调用方不需要自己裁剪 +/// 行数据，也不会出现“列删了、值没删”导致的错位。 +pub fn build_export_insert_statements_excluding( + options: BuildExportInsertStatementsOptions, + exclude_columns: &[String], +) -> Result, String> { if options.columns.is_empty() || options.rows.is_empty() { return Ok(Vec::new()); } + let excluded_names: HashSet = + exclude_columns.iter().map(|column| column.trim().to_ascii_uppercase()).collect(); let table = export_qualified_table_name( options.database_type, options.schema.as_deref(), @@ -1172,6 +1189,26 @@ pub fn build_export_insert_statements(options: BuildExportInsertStatementsOption .iter() .enumerate() .filter_map(|(index, column)| { + let column_type = export_column_type(&options.column_types, index, options.database_type, &spatial_columns); + let excluded = excluded_names.contains(&column.trim().to_ascii_uppercase()); + (is_export_insert_column( + options.database_type, + column, + column_type, + options.column_extras.get(index).and_then(|value| value.as_deref()), + ) && !excluded) + .then(|| { + let sqlserver_unicode_string = options.database_type == Some(DatabaseType::SqlServer) + && column_type.is_some_and(is_sqlserver_unicode_export_type); + (index, column, sqlserver_unicode_string) + }) + }) + .collect::>(); + if insert_columns.is_empty() { + // Only fail when the exclusion itself removed the last insertable column; + // emptiness caused by other omission rules (e.g. generated columns) keeps + // the silent empty result. + let had_insertable_without_exclusion = options.columns.iter().enumerate().any(|(index, column)| { let column_type = export_column_type(&options.column_types, index, options.database_type, &spatial_columns); is_export_insert_column( options.database_type, @@ -1179,14 +1216,10 @@ pub fn build_export_insert_statements(options: BuildExportInsertStatementsOption column_type, options.column_extras.get(index).and_then(|value| value.as_deref()), ) - .then(|| { - let sqlserver_unicode_string = options.database_type == Some(DatabaseType::SqlServer) - && column_type.is_some_and(is_sqlserver_unicode_export_type); - (index, column, sqlserver_unicode_string) - }) - }) - .collect::>(); - if insert_columns.is_empty() { + }); + if had_insertable_without_exclusion { + return Err("No insertable columns remain after excluding columns from the export.".to_string()); + } return Ok(Vec::new()); } let batch_size = if options.database_type.is_some_and(uses_single_row_insert_statements) { @@ -1354,7 +1387,8 @@ fn is_postgres_bytea_export_column(database_type: Option, column_t } pub fn build_export_sql_insert(options: BuildExportSqlInsertOptions) -> Result { - build_export_insert_statements(options.insert).map(|statements| statements.join("\n")) + build_export_insert_statements_excluding(options.insert, &options.exclude_columns) + .map(|statements| statements.join("\n")) } pub fn build_database_sql_export(options: BuildDatabaseSqlExportOptions) -> Result { @@ -3837,18 +3871,20 @@ mod tests { }; use super::{ build_database_export_object_source_sql, build_database_sql_export, build_export_insert_statements, - create_database_export_writer, database_export_query_options_for_timeout, database_export_select_sql, - database_export_total_objects, drop_table_if_exists_sql, ensure_export_destination_dir, - export_destination_identity_mismatch, filter_export_table_infos, format_export_sql_literal, - format_export_table_ddl, format_mysql_spatial_export_literal, format_xugu_spatial_export_literal, - generate_postgres_extension_ddl, generate_postgres_sequence_create_ddl, generate_postgres_sequence_owner_ddl, + build_export_insert_statements_excluding, build_export_sql_insert, create_database_export_writer, + database_export_query_options_for_timeout, database_export_select_sql, database_export_total_objects, + drop_table_if_exists_sql, ensure_export_destination_dir, export_destination_identity_mismatch, + filter_export_table_infos, format_export_sql_literal, format_export_table_ddl, + format_mysql_spatial_export_literal, format_xugu_spatial_export_literal, generate_postgres_extension_ddl, + generate_postgres_sequence_create_ddl, generate_postgres_sequence_owner_ddl, generate_postgres_sequence_setval_sql, is_postgres_extension_member_routine, mysql_database_export_preamble, mysql_view_dependencies_from_rows, mysql_view_dependencies_sql, normalize_export_table_ddl, record_export_destination_identity, record_export_error, replace_database_export_select_list, sort_export_views_by_dependencies, split_postgres_export_table_triggers, write_database_export_rows, - BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, DatabaseExportObjectCounts, - DatabaseExportRequest, DatabaseExportWriter, DdlNormalizeOptions, ExportedTableSql, PostgresExportExtension, - PostgresExportSequence, PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT, + BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, BuildExportSqlInsertOptions, + DatabaseExportObjectCounts, DatabaseExportRequest, DatabaseExportWriter, DdlNormalizeOptions, ExportedTableSql, + PostgresExportExtension, PostgresExportSequence, PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, + DATABASE_EXPORT_ROW_LIMIT, }; use super::{ExportProgress, LenientExportErrors}; use crate::connection::AppState; @@ -5450,6 +5486,78 @@ mod tests { assert_eq!(statements, vec!["INSERT INTO \"public\".\"articles\" (\"id\", \"title\") VALUES (1, 'Hello');"]); } + #[test] + fn sql_insert_export_omits_excluded_columns_and_keeps_values_aligned() { + let statements = build_export_insert_statements_excluding( + BuildExportInsertStatementsOptions { + database_type: Some(DatabaseType::Postgres), + identifier_quote: None, + schema: Some("public".to_string()), + table_name: Some("users".to_string()), + qualified_table_name: None, + columns: vec!["id".to_string(), "name".to_string(), "email".to_string()], + column_types: vec![Some("integer".to_string()), Some("text".to_string()), Some("text".to_string())], + column_extras: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows: vec![vec![json!(1), json!("Ada"), json!("ada@example.com")]], + batch_size: Some(10), + }, + &["id".to_string()], + ) + .unwrap(); + + assert_eq!( + statements, + vec!["INSERT INTO \"public\".\"users\" (\"name\", \"email\") VALUES ('Ada', 'ada@example.com');"] + ); + } + + #[test] + fn sql_insert_export_reports_error_when_excluding_leaves_no_columns() { + let result = build_export_insert_statements_excluding( + BuildExportInsertStatementsOptions { + database_type: Some(DatabaseType::Postgres), + identifier_quote: None, + schema: Some("public".to_string()), + table_name: Some("user_roles".to_string()), + qualified_table_name: None, + columns: vec!["user_id".to_string()], + column_types: vec![Some("integer".to_string())], + column_extras: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows: vec![vec![json!(1)]], + batch_size: Some(10), + }, + &["user_id".to_string()], + ); + + assert_eq!( + result.expect_err("excluding every column must fail instead of writing an empty export"), + "No insertable columns remain after excluding columns from the export." + ); + } + + #[test] + fn build_export_sql_insert_honors_exclude_columns_from_payload() { + let options: BuildExportSqlInsertOptions = serde_json::from_value(json!({ + "databaseType": "postgres", + "schema": "public", + "tableName": "users", + "columns": ["id", "name"], + "columnTypes": ["integer", "text"], + "rows": [[1, "Ada"]], + "excludeColumns": ["id"], + "batchSize": 10 + })) + .expect("deserialize export insert payload"); + + let sql = build_export_sql_insert(options).expect("build export sql insert"); + + assert_eq!(sql, "INSERT INTO \"public\".\"users\" (\"name\") VALUES ('Ada');"); + } + #[test] fn mysql_generated_columns_are_omitted_from_sql_inserts_but_kept_in_ddl() { let ddl = "CREATE TABLE `orders` (`id` bigint AUTO_INCREMENT, `quantity` int, `unit_price` decimal(10,2), `virtual_total` decimal(10,2) GENERATED ALWAYS AS ((`quantity` * `unit_price`)) VIRTUAL, `stored_total` decimal(10,2) GENERATED ALWAYS AS ((`quantity` * `unit_price`)) STORED);"; diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 8b0b0a07d0..201bc5e9c7 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -16,7 +16,7 @@ use crate::csv_export::{ }; pub use crate::database_export::ExportStatus; use crate::database_export::{ - build_export_insert_statements, is_export_cancelled, BuildExportInsertStatementsOptions, SqlInsertMode, + build_export_insert_statements_excluding, is_export_cancelled, BuildExportInsertStatementsOptions, SqlInsertMode, }; use crate::models::connection::DatabaseType; use crate::query::{ @@ -120,6 +120,13 @@ pub struct QueryResultExportRequest { pub auto_filter: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub identifier_quote: Option, + /// 导出 SQL 时是否排除主键列（对应前端数据提取设置里的“排除主键”）。 + #[serde(default)] + pub exclude_primary_keys: bool, + /// 结果集对应的原表主键列名。查询结果导出由前端从表元数据带过来； + /// 没有表元数据（例如自由 SQL 查询）时为空，此时无法排除主键。 + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub primary_keys: Vec, } pub struct StagedExportTarget { @@ -322,6 +329,8 @@ struct SqlInsertWriter { schema: Option, table_name: String, identifier_quote: Option, + /// 生成 INSERT 时需要排除的列名，来自请求里的“不含主键”设置。 + exclude_columns: Vec, } /// Streams query-result JSON rows without retaining the complete result set. @@ -408,6 +417,7 @@ impl SqlInsertWriter { schema: request.schema.clone(), table_name, identifier_quote: request.identifier_quote.clone(), + exclude_columns: if request.exclude_primary_keys { request.primary_keys.clone() } else { Vec::new() }, }) } @@ -439,20 +449,23 @@ impl SqlInsertWriter { if self.pending_rows.is_empty() { return Ok(()); } - let stmts = build_export_insert_statements(BuildExportInsertStatementsOptions { - database_type: Some(self.database_type), - identifier_quote: self.identifier_quote.clone(), - 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(), - spatial_columns: self.spatial_columns.clone(), - spatial_values: mem::take(&mut self.pending_spatial_values), - rows: mem::take(&mut self.pending_rows), - batch_size: Some(self.insert_mode.batch_size(SQL_INSERT_BATCH_SIZE)), - })?; + let stmts = build_export_insert_statements_excluding( + BuildExportInsertStatementsOptions { + database_type: Some(self.database_type), + identifier_quote: self.identifier_quote.clone(), + 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(), + spatial_columns: self.spatial_columns.clone(), + spatial_values: mem::take(&mut self.pending_spatial_values), + rows: mem::take(&mut self.pending_rows), + batch_size: Some(self.insert_mode.batch_size(SQL_INSERT_BATCH_SIZE)), + }, + &self.exclude_columns, + )?; 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}"))?; @@ -1965,6 +1978,7 @@ async fn try_export_sqlserver_query_result_stream( #[cfg(test)] mod tests { use super::*; + use serde_json::json; #[test] fn staged_export_target_preserves_existing_destination_on_discard_and_replace_failure() { @@ -2010,6 +2024,60 @@ mod tests { assert!(!stream_export_was_cancelled("network failure", false, false)); } + #[test] + fn sql_insert_writer_omits_excluded_primary_key_columns() { + let dir = tempfile::tempdir().expect("temp dir"); + let file_path = dir.path().join("users.sql"); + let request = QueryResultExportRequest { + export_id: "export-1".to_string(), + connection_id: "conn-1".to_string(), + database: "db".to_string(), + schema: Some("public".to_string()), + catalog: None, + sql: "SELECT * FROM users".to_string(), + query_base_sql: "SELECT * FROM users".to_string(), + setup_sql: Vec::new(), + database_type: DatabaseType::Postgres, + use_agent_cursor: false, + file_path: file_path.to_string_lossy().to_string(), + format: "sql".to_string(), + include_sql_sheet: false, + page_size: 1000, + row_limit: None, + total_rows: None, + timeout_secs: None, + keyset_optimization_enabled: false, + client_session_id: None, + execution_id: None, + date_time_format: None, + export_table_name: Some("users".to_string()), + export_column_types: None, + numeric_column_right_align: false, + column_comments: None, + auto_filter: None, + identifier_quote: None, + insert_mode: Default::default(), + csv_quote_mode: Default::default(), + exclude_primary_keys: true, + primary_keys: vec!["id".to_string()], + }; + + let mut writer = SqlInsertWriter::create(&request).expect("create sql insert writer"); + writer.set_columns( + vec!["id".to_string(), "name".to_string()], + &["integer".to_string(), "text".to_string()], + &[], + &request, + ); + writer.write_row(vec![json!(1), json!("Ada")], None).expect("write export row"); + writer.finish().expect("finish export"); + + assert_eq!( + std::fs::read_to_string(&file_path).expect("read sql export"), + "INSERT INTO \"public\".\"users\" (\"name\") VALUES ('Ada');\n" + ); + } + #[test] fn postgres_temp_setup_accepts_only_session_local_table_operations() { let safe = vec![ @@ -2070,6 +2138,8 @@ mod tests { column_comments: None, auto_filter: None, identifier_quote: None, + exclude_primary_keys: false, + primary_keys: Vec::new(), } } diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index c49418e157..c405bbb23b 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -18,10 +18,11 @@ use crate::csv_export::{ format_csv_with_quote_mode, format_tsv, push_table_csv_row, push_table_csv_row_with_quote_mode, push_tsv_row, CsvQuoteMode, }; +use crate::data_grid_sql::extra_is_auto_generated; pub use crate::database_export::ExportStatus; use crate::database_export::{ - build_export_insert_statements, is_export_cancelled, is_internal_export_column, BuildExportInsertStatementsOptions, - SqlInsertMode, + build_export_insert_statements_excluding, is_export_cancelled, is_internal_export_column, + BuildExportInsertStatementsOptions, SqlInsertMode, }; use crate::db::agent_driver::AgentTableReadStartParams; use crate::models::connection::DatabaseType; @@ -42,6 +43,31 @@ pub fn table_export_client_session_id(export_id: &str) -> String { task_client_session_id("table-export", export_id) } +/// SQL 导出时需要排除的列名：用户选择了“不含主键”时，与 copy-as-INSERT +/// 一致只剔除自增/identity 主键；手动赋值的主键保留，否则回放 INSERT 缺值。 +fn sql_export_excluded_columns( + request: &TableExportRequest, + primary_keys: &[String], + col_names: &[String], + column_extras: &[Option], +) -> Vec { + if !request.exclude_primary_keys { + return Vec::new(); + } + primary_keys + .iter() + .filter(|pk| { + col_names + .iter() + .position(|name| name.eq_ignore_ascii_case(pk)) + .and_then(|index| column_extras.get(index)) + .and_then(|extra| extra.as_deref()) + .is_some_and(extra_is_auto_generated) + }) + .cloned() + .collect() +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TableExportRequest { @@ -65,6 +91,9 @@ pub struct TableExportRequest { pub column_types: Option>>, #[serde(default)] pub primary_keys: Option>, + /// 导出 SQL 时是否排除主键列（对应前端数据提取设置里的“排除主键”）。 + #[serde(default)] + pub exclude_primary_keys: bool, #[serde(default, skip_serializing_if = "Option::is_none")] pub where_input: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -1314,20 +1343,23 @@ async fn try_export_native_table_stream( if pending_rows.is_empty() { return Ok(()); } - let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { - database_type: Some(*db_type), - identifier_quote: request.identifier_quote.clone(), - schema: request.schema.clone(), - table_name: Some(request.table_name.clone()), - qualified_table_name: None, - columns: col_names.to_vec(), - column_types: column_types.to_vec(), - column_extras: column_extras.to_vec(), - spatial_columns: Vec::new(), - spatial_values: Vec::new(), - rows: std::mem::take(pending_rows), - batch_size: Some(request.insert_mode.batch_size(SQL_INSERT_BATCH_SIZE)), - })?; + let statements = build_export_insert_statements_excluding( + BuildExportInsertStatementsOptions { + database_type: Some(*db_type), + identifier_quote: request.identifier_quote.clone(), + schema: request.schema.clone(), + table_name: Some(request.table_name.clone()), + qualified_table_name: None, + columns: col_names.to_vec(), + column_types: column_types.to_vec(), + column_extras: column_extras.to_vec(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows: std::mem::take(pending_rows), + batch_size: Some(request.insert_mode.batch_size(SQL_INSERT_BATCH_SIZE)), + }, + &sql_export_excluded_columns(request, primary_keys, col_names, column_extras), + )?; if !statements.is_empty() { if wrote_statements { file.write_all(b"\n").map_err(|e| format!("Failed to write SQL: {e}"))?; @@ -2152,20 +2184,23 @@ async fn export_table_data_core_inner( break; } - let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { - database_type: Some(db_type), - identifier_quote: request.identifier_quote.clone(), - schema: request.schema.clone(), - table_name: Some(request.table_name.clone()), - qualified_table_name: None, - columns: col_names.clone(), - column_types: column_types.clone(), - column_extras: column_extras.clone(), - spatial_columns: result.spatial_columns.clone(), - spatial_values: result.spatial_values.clone(), - rows: result.rows.clone(), - batch_size: Some(request.insert_mode.batch_size(SQL_INSERT_BATCH_SIZE)), - })?; + let statements = build_export_insert_statements_excluding( + BuildExportInsertStatementsOptions { + database_type: Some(db_type), + identifier_quote: request.identifier_quote.clone(), + schema: request.schema.clone(), + table_name: Some(request.table_name.clone()), + qualified_table_name: None, + columns: col_names.clone(), + column_types: column_types.clone(), + column_extras: column_extras.clone(), + spatial_columns: result.spatial_columns.clone(), + spatial_values: result.spatial_values.clone(), + rows: result.rows.clone(), + batch_size: Some(request.insert_mode.batch_size(SQL_INSERT_BATCH_SIZE)), + }, + &sql_export_excluded_columns(request, &primary_keys, &col_names, &column_extras), + )?; if !statements.is_empty() { if wrote_statements { file.write_all(b"\n").map_err(|e| format!("Failed to write SQL: {e}"))?; @@ -2226,7 +2261,7 @@ async fn export_table_data_core_inner( #[cfg(test)] mod tests { use super::*; - use crate::database_export::{clear_export_cancelled, set_export_cancelled}; + use crate::database_export::{build_export_insert_statements, clear_export_cancelled, set_export_cancelled}; use crate::models::connection::ConnectionConfig; #[cfg(unix)] use crate::plugins::{ @@ -2351,6 +2386,7 @@ mod tests { columns: Some(vec!["id".to_string(), "name".to_string()]), column_types: Some(vec![Some("INTEGER".to_string()), Some("VARCHAR".to_string())]), primary_keys: Some(vec!["id".to_string()]), + exclude_primary_keys: false, where_input: None, order_by: None, skip_count, @@ -2512,6 +2548,7 @@ mod tests { format: "csv".to_string(), insert_mode: Default::default(), csv_quote_mode: CsvQuoteMode::All, + exclude_primary_keys: false, columns: None, column_types: None, primary_keys: None, @@ -2680,6 +2717,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: Some("WHERE temperature > 1".to_string()), order_by: Some("Time DESC".to_string()), skip_count: true, @@ -2739,6 +2777,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: true, @@ -2776,6 +2815,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: true, @@ -2808,6 +2848,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: true, @@ -2848,6 +2889,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: true, @@ -2900,6 +2942,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: Some("WHERE status = 'active'".to_string()), order_by: None, skip_count: false, @@ -2946,6 +2989,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: false, @@ -3020,6 +3064,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: false, @@ -3076,6 +3121,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: true, @@ -3134,6 +3180,7 @@ mod tests { columns: None, column_types: None, primary_keys: None, + exclude_primary_keys: false, where_input: None, order_by: None, skip_count: false, 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 9d75245687..8d4c35fea3 100644 --- a/crates/dbx-core/tests/live_clickhouse_query_result_export.rs +++ b/crates/dbx-core/tests/live_clickhouse_query_result_export.rs @@ -109,6 +109,8 @@ async fn live_clickhouse_query_result_export_xlsx_streams_random_order_query_onc auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let done_seen = AtomicBool::new(false); let result = export_query_result_core(&state, &request, None, |progress| { diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index bdff3caa61..1b4ea91266 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -484,6 +484,8 @@ async fn live_mysql_query_result_export_xlsx_streams_single_query_without_duplic auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let done_seen = AtomicBool::new(false); let result = export_query_result_core(&state, &request, None, |progress| { @@ -575,6 +577,8 @@ async fn live_mysql_csv_temporal_export_round_trip_preserves_dbx_force_text_valu auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; export_query_result_core(&state, &request, None, |_| {}).await.expect("export fixture"); @@ -690,6 +694,8 @@ async fn live_mysql_xlsx_export_can_outlive_query_timeout_while_rows_keep_arrivi auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let rows_exported = AtomicU64::new(0); let done_seen = AtomicBool::new(false); 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 68bc292e61..6a2697586a 100644 --- a/crates/dbx-core/tests/live_postgres_query_result_export.rs +++ b/crates/dbx-core/tests/live_postgres_query_result_export.rs @@ -156,6 +156,8 @@ async fn live_postgres_query_result_export_uses_single_streamed_query() { auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let done_seen = AtomicBool::new(false); let result = export_query_result_core(&state, &request, None, |progress| { @@ -240,6 +242,8 @@ async fn live_postgres_query_result_xlsx_preserves_temporal_cell_types() { auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; export_query_result_core(&state, &request, None, |_| {}).await.expect("export temporal XLSX"); @@ -313,6 +317,8 @@ async fn live_postgres_numeric_xlsx_ignores_fractional_trailing_zeros() { auto_filter: None, identifier_quote: None, numeric_column_right_align: true, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; export_query_result_core(&state, &request, None, |_| {}).await.expect("export numeric XLSX"); @@ -391,6 +397,8 @@ async fn live_postgres_truncated_batch_result_export_replays_safe_temp_setup() { auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let csv_rows = AtomicU64::new(0); export_query_result_core(&state, &request, None, |progress| { @@ -471,6 +479,8 @@ async fn live_postgres_xlsx_export_can_outlive_query_timeout_while_rows_keep_arr auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let rows_exported = AtomicU64::new(0); let done_seen = AtomicBool::new(false); @@ -542,6 +552,8 @@ async fn live_postgres_stream_still_times_out_without_progress_and_recovers() { auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let started_at = Instant::now(); let result = export_query_result_core(&state, &request, None, |_| {}).await; diff --git a/crates/dbx-core/tests/live_sqlserver_completion.rs b/crates/dbx-core/tests/live_sqlserver_completion.rs index c46d89f0b7..b81e519eac 100644 --- a/crates/dbx-core/tests/live_sqlserver_completion.rs +++ b/crates/dbx-core/tests/live_sqlserver_completion.rs @@ -1218,6 +1218,8 @@ async fn live_sqlserver_query_result_export_streams_cte_query_to_csv() { auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let done_seen = AtomicBool::new(false); let result = export_query_result_core(&state, &request, None, |progress| { 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 e19d40a81c..388f94de26 100644 --- a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs +++ b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs @@ -135,6 +135,8 @@ async fn live_sqlserver_xlsx_export_can_outlive_query_timeout_while_rows_keep_ar auto_filter: None, identifier_quote: None, numeric_column_right_align: false, + exclude_primary_keys: false, + primary_keys: Vec::new(), }; let rows_exported = AtomicU64::new(0); let done_seen = AtomicBool::new(false); diff --git a/packages/app-tests/useDataGridExport.test.ts b/packages/app-tests/useDataGridExport.test.ts index bc70daed49..95ae0dc290 100644 --- a/packages/app-tests/useDataGridExport.test.ts +++ b/packages/app-tests/useDataGridExport.test.ts @@ -4,6 +4,10 @@ import { beforeEach, test, vi } from "vitest"; import { createPinia, setActivePinia } from "pinia"; import { useSettingsStore } from "../../apps/desktop/src/stores/settingsStore.ts"; import type { DataGridTableMeta } from "../../apps/desktop/src/lib/dataGrid/dataGridSql.ts"; +import { + DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS, + type DataGridExtractorOptions, +} from "../../apps/desktop/src/lib/dataGrid/dataGridCopyExtractor.ts"; import type { DatabaseType, QueryResult } from "../../apps/desktop/src/types/database.ts"; const apiMock = vi.hoisted(() => ({ @@ -109,6 +113,7 @@ function buildExportHarness( sourceColumns?: Array; databaseType?: DatabaseType; context?: "results" | "table-data"; + extractorOptions?: DataGridExtractorOptions; } = {}, ) { const exportColumns = options.columns ?? ["id", "name"]; @@ -166,6 +171,7 @@ function buildExportHarness( connectionId: computed(() => "conn-1"), database: computed(() => "db"), context: computed(() => options.context ?? "results"), + extractorOptions: computed(() => options.extractorOptions), sourceColumns: computed(() => options.sourceColumns), columnTypes: computed(() => options.columnTypes), allColumnTypes: computed(() => options.allColumnTypes), @@ -634,6 +640,112 @@ test("table data SQL export keeps source column names", async () => { } }); +function sqlExclusionOptions(): DataGridExtractorOptions { + return { + ...DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS, + sql: { ...DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS.sql, excludePrimaryKeysFromInsert: true }, + }; +} + +test("table data SQL export passes primary key exclusion to the backend", async () => { + runtimeMock.isTauri = true; + dialogMock.save.mockResolvedValue("/tmp/users.sql"); + const { composable } = buildExportHarness({ + context: "table-data", + tableMeta: { + tableName: "users", + primaryKeys: ["id"], + columns: [ + { name: "id", data_type: "integer", is_nullable: false, is_primary_key: true, extra: "auto_increment" }, + { name: "name", data_type: "text", is_nullable: true }, + ], + }, + extractorOptions: sqlExclusionOptions(), + }); + + await composable.exportSql(); + + assert.equal(apiMock.startTableExport.mock.calls[0][0].excludePrimaryKeys, true); +}); + +test("query result SQL export passes primary key exclusion and keys to the backend", async () => { + runtimeMock.isTauri = true; + dialogMock.save.mockResolvedValue("/tmp/query-result.sql"); + const { composable } = buildExportHarness({ + context: "results", + tableMeta: { + tableName: "users", + primaryKeys: ["id"], + columns: [ + { name: "id", data_type: "integer", is_nullable: false, is_primary_key: true, extra: "auto_increment" }, + { name: "name", data_type: "text", is_nullable: true }, + ], + }, + extractorOptions: sqlExclusionOptions(), + }); + + await composable.exportSql(); + + assert.equal(apiMock.startQueryResultExport.mock.calls[0][0].excludePrimaryKeys, true); + assert.deepEqual(apiMock.startQueryResultExport.mock.calls[0][0].primaryKeys, ["id"]); +}); + +test("table data SQL export keeps manually-assigned primary keys", async () => { + runtimeMock.isTauri = true; + dialogMock.save.mockResolvedValue("/tmp/users-manual-pk.sql"); + const { composable } = buildExportHarness({ + context: "table-data", + tableMeta: { + tableName: "users", + primaryKeys: ["user_code"], + columns: [{ name: "user_code", data_type: "varchar", is_nullable: false, is_primary_key: true }], + }, + extractorOptions: sqlExclusionOptions(), + }); + + await composable.exportSql(); + + assert.equal(apiMock.startTableExport.mock.calls[0][0].excludePrimaryKeys, undefined); +}); + +test("local SQL export asks the backend to drop primary key columns", async () => { + const download = installTextDownloadCapture(); + const completeLocalResult: QueryResult = { + columns: ["id", "name"], + column_types: ["int4", "text"], + rows: [[1, "Ada"]], + affected_rows: 0, + execution_time_ms: 1, + truncated: false, + has_more: false, + }; + apiMock.buildExportSqlInsert.mockResolvedValueOnce("INSERT INTO \"users\" (\"name\") VALUES ('Ada');"); + + try { + const { composable } = buildExportHarness({ + columns: completeLocalResult.columns, + rows: completeLocalResult.rows, + completeLocalResult, + tableMeta: { + tableName: "users", + primaryKeys: ["id"], + columns: [ + { name: "id", data_type: "int4", is_nullable: false, is_primary_key: true, extra: "auto_increment" }, + { name: "name", data_type: "text", is_nullable: true }, + ], + }, + extractorOptions: sqlExclusionOptions(), + }); + + await composable.exportSql(); + + assert.deepEqual(apiMock.buildExportSqlInsert.mock.calls[0][0].excludeColumns, ["id"]); + assert.deepEqual(apiMock.buildExportSqlInsert.mock.calls[0][0].columns, ["id", "name"]); + } finally { + download.restore(); + } +}); + test("complete local query result export removes only internal hidden columns", async () => { const completeLocalResult: QueryResult = { columns: ["id", "name", "__DBX_PK_id"],