From e3d5d141832c7c6fd7808ee8fe5406bdb9ae43e3 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 7 Sep 2026 18:24:37 +0800 Subject: [PATCH] fix(grid): keep manual primary keys in exclude-PK insert copies Closes #8293 --- .../__tests__/useDataGridExport.spec.ts | 28 ++- .../src/composables/useDataGridExport.ts | 7 +- .../dbx-core/src/data_grid_extractors/sql.rs | 8 +- .../src/data_grid_extractors/tests.rs | 66 +++++++- crates/dbx-core/src/data_grid_sql.rs | 159 +++++++++++++++++- 5 files changed, 261 insertions(+), 7 deletions(-) diff --git a/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts index 3c3eeaf2bc..48c46d5cf8 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridExport.spec.ts @@ -923,17 +923,43 @@ describe("useDataGridExport prepared row statements", () => { }); it("disables INSERT when primary-key exclusion removes every selected column", () => { + const autoIncrementTable: DataGridTableMeta = { + tableName: "users", + primaryKeys: ["id"], + columns: [{ name: "id", data_type: "int", is_nullable: false, is_primary_key: true, extra: "auto_increment" }], + }; const matrix: CellSelectionMatrix = { rowIndexes: [0], columnIndexes: [0], columns: ["id"], rows: [[1]] }; const excludePrimaryKeys = { ...DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS, sql: { ...DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS.sql, skipGeneratedColumns: false, excludePrimaryKeysFromInsert: true }, }; - const state = createExportState(editableTable, ["id"], matrix, [1], undefined, undefined, [], excludePrimaryKeys); + const state = createExportState(autoIncrementTable, ["id"], matrix, [1], undefined, undefined, [], excludePrimaryKeys); expect(state.canCopyWithExtractor("sql-inserts")).toBe(false); }); + it("keeps a manually-assigned primary key insertable under primary-key exclusion", () => { + const compositeKeyTable: DataGridTableMeta = { + tableName: "daily_stats", + primaryKeys: ["id", "stat_date"], + columns: [ + { name: "id", data_type: "bigint", is_nullable: false, is_primary_key: true, extra: "auto_increment" }, + { name: "stat_date", data_type: "date", is_nullable: false, is_primary_key: true }, + { name: "name", data_type: "varchar", is_nullable: false }, + ], + }; + const matrix: CellSelectionMatrix = { rowIndexes: [0], columnIndexes: [1], columns: ["stat_date"], rows: [["2026-08-18"]] }; + const excludePrimaryKeys = { + ...DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS, + sql: { ...DEFAULT_DATA_GRID_EXTRACTOR_OPTIONS.sql, skipGeneratedColumns: false, excludePrimaryKeysFromInsert: true }, + }; + + const state = createExportState(compositeKeyTable, ["id", "stat_date", "name"], matrix, [1, "2026-08-18", "Ada"], undefined, undefined, [], excludePrimaryKeys); + + expect(state.canCopyWithExtractor("sql-inserts")).toBe(true); + }); + it("keeps an auto-increment primary key when copying only the primary key column as INSERT", () => { const autoIncrementTable: DataGridTableMeta = { tableName: "users", diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index 8664d5498e..9876ef62dc 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -582,7 +582,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { function insertableCopyColumnCount(excludePrimaryKeys: boolean, copyColumns = effectiveColumns(sourceColumns.value, columns.value), extractorOptions?: DataGridExtractorOptions): number { const primaryKeySet = new Set((tableMeta.value?.primaryKeys ?? []).map(normalizeColumnName)); - return copyColumns.filter((column): column is string => !!column && !isCopyInsertOmittedColumn(databaseType.value, column, tableMeta.value, extractorOptions) && (!excludePrimaryKeys || !primaryKeySet.has(normalizeColumnName(column)))).length; + return copyColumns.filter((column): column is string => !!column && !isCopyInsertOmittedColumn(databaseType.value, column, tableMeta.value, extractorOptions) && (!excludePrimaryKeys || !primaryKeySet.has(normalizeColumnName(column)) || !isAutoGeneratedColumn(column, tableMeta.value))).length; } async function buildMongoExtractorInsert(extractorOptions: DataGridExtractorOptions, rowLimit?: number): Promise { @@ -1502,6 +1502,11 @@ function effectiveColumns(sourceColumns: Array | undefined, return sourceColumns; } +function isAutoGeneratedColumn(column: string, tableMeta: DataGridTableMeta | undefined): boolean { + const columnInfo = tableMeta?.columns?.find((item) => normalizeColumnName(item.name) === normalizeColumnName(column)); + return /\b(auto_increment|autoincrement|identity)\b/.test(columnInfo?.extra?.toLowerCase() ?? ""); +} + function isCopyInsertOmittedColumn(databaseType: DatabaseType | undefined, column: string, tableMeta: DataGridTableMeta | undefined, extractorOptions?: DataGridExtractorOptions): boolean { if (usesSyntheticRowIdKey(databaseType, [column])) return true; const columnInfo = tableMeta?.columns?.find((item) => normalizeColumnName(item.name) === normalizeColumnName(column)); diff --git a/crates/dbx-core/src/data_grid_extractors/sql.rs b/crates/dbx-core/src/data_grid_extractors/sql.rs index 6a24447ef1..b0674a7dc0 100644 --- a/crates/dbx-core/src/data_grid_extractors/sql.rs +++ b/crates/dbx-core/src/data_grid_extractors/sql.rs @@ -257,7 +257,13 @@ fn sql_selected_data(context: &ExtractContext<'_>, for_update: bool) -> Result)> = insertable_columns .iter() - .filter(|(column, _, _)| { - !options.exclude_primary_keys || !primary_key_set.contains(&normalize_column_name(column)) + .filter(|(column, _, info)| { + !options.exclude_primary_keys + || !primary_key_set.contains(&normalize_column_name(column)) + || !info.as_ref().is_some_and(is_auto_generated_column) }) .cloned() .collect(); @@ -3792,7 +3798,32 @@ mod tests { schema: None, table_name: "users".to_string(), primary_keys: vec!["id".to_string()], - columns: None, + columns: Some(vec![ + DataGridColumnInfo { + name: "id".to_string(), + data_type: "bigint".to_string(), + is_nullable: false, + is_primary_key: true, + column_default: None, + extra: Some("auto_increment".to_string()), + }, + DataGridColumnInfo { + name: "login_name".to_string(), + data_type: "varchar(64)".to_string(), + is_nullable: false, + is_primary_key: false, + column_default: None, + extra: None, + }, + DataGridColumnInfo { + name: "display_name".to_string(), + data_type: "varchar(64)".to_string(), + is_nullable: true, + is_primary_key: false, + column_default: None, + extra: None, + }, + ]), }), columns: vec!["id".to_string(), "login_name".to_string(), "display_name".to_string()], column_types: None, @@ -3808,6 +3839,128 @@ mod tests { ); } + #[test] + fn copy_insert_primary_key_exclusion_keeps_manual_composite_key_members() { + let statement = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions { + database_type: Some(DatabaseType::Mysql), + identifier_quote: None, + table_meta: Some(DataGridTableMeta { + catalog: None, + database: None, + schema: None, + table_name: "move_statistic_product_daily".to_string(), + primary_keys: vec!["id".to_string(), "stat_date".to_string()], + columns: Some(vec![ + DataGridColumnInfo { + name: "id".to_string(), + data_type: "bigint unsigned".to_string(), + is_nullable: false, + is_primary_key: true, + column_default: None, + extra: Some("auto_increment".to_string()), + }, + DataGridColumnInfo { + name: "stat_date".to_string(), + data_type: "date".to_string(), + is_nullable: false, + is_primary_key: true, + column_default: None, + extra: None, + }, + DataGridColumnInfo { + name: "product_name".to_string(), + data_type: "varchar(255)".to_string(), + is_nullable: false, + is_primary_key: false, + column_default: Some("".to_string()), + extra: None, + }, + ]), + }), + columns: vec!["id".to_string(), "stat_date".to_string(), "product_name".to_string()], + column_types: None, + source_columns: None, + rows: vec![vec![json!(1), json!("2026-08-18"), json!("sweater")]], + exclude_primary_keys: true, + include_computed_columns: false, + insert_mode: DataGridCopyInsertMode::Merged, + }); + assert_eq!( + statement.as_deref(), + Some("INSERT INTO `move_statistic_product_daily` (`stat_date`, `product_name`) VALUES ('2026-08-18', 'sweater');") + ); + } + + #[test] + fn copy_insert_primary_key_exclusion_keeps_manual_primary_keys() { + let statement = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions { + database_type: Some(DatabaseType::Postgres), + identifier_quote: None, + table_meta: Some(DataGridTableMeta { + catalog: None, + database: None, + schema: Some("public".to_string()), + table_name: "countries".to_string(), + primary_keys: vec!["code".to_string()], + columns: Some(vec![ + DataGridColumnInfo { + name: "code".to_string(), + data_type: "varchar(2)".to_string(), + is_nullable: false, + is_primary_key: true, + column_default: None, + extra: None, + }, + DataGridColumnInfo { + name: "label".to_string(), + data_type: "text".to_string(), + is_nullable: false, + is_primary_key: false, + column_default: None, + extra: None, + }, + ]), + }), + columns: vec!["code".to_string(), "label".to_string()], + column_types: None, + source_columns: None, + rows: vec![vec![json!("AD"), json!("Andorra")]], + exclude_primary_keys: true, + include_computed_columns: false, + insert_mode: DataGridCopyInsertMode::Merged, + }); + assert_eq!( + statement.as_deref(), + Some("INSERT INTO \"public\".\"countries\" (\"code\", \"label\") VALUES ('AD', 'Andorra');") + ); + } + + #[test] + fn copy_insert_primary_key_exclusion_keeps_unknown_metadata_primary_keys() { + // Without column metadata we cannot prove the key is auto-generated; + // keep it rather than silently dropping NOT NULL data. + let statement = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions { + database_type: Some(DatabaseType::Mysql), + identifier_quote: None, + table_meta: Some(DataGridTableMeta { + catalog: None, + database: None, + schema: None, + table_name: "users".to_string(), + primary_keys: vec!["id".to_string()], + columns: None, + }), + columns: vec!["id".to_string(), "login_name".to_string()], + column_types: None, + source_columns: None, + rows: vec![vec![json!(1), json!("ada")]], + exclude_primary_keys: true, + include_computed_columns: false, + insert_mode: DataGridCopyInsertMode::Merged, + }); + assert_eq!(statement.as_deref(), Some("INSERT INTO `users` (`id`, `login_name`) VALUES (1, 'ada');")); + } + #[test] fn copy_insert_keeps_json_cells_as_single_json_literals() { let statement = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions {