diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 036e9c350b..17e7cef1ef 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -525,6 +525,7 @@ interface DataGridProps { includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; + exportColumnExtras?: Array; insertMode?: SqlInsertMode; }) => Promise; allExportResults?: Array<{ diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index bff81de661..bc9d47f450 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -2116,8 +2116,16 @@ defineExpose({ :on-execute-sql="async (sql: string) => emit('executeSql', activeTab.id, 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' | 'json' | 'txt' | 'sql'; includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; insertMode?: SqlInsertMode }) => - queryStore.buildQueryResultExportRequest(activeTab.id, options) + (options: { + exportId: string; + filePath: string; + format: 'csv' | 'xlsx' | 'json' | 'txt' | 'sql'; + includeSqlSheet?: boolean; + exportTableName?: string; + exportColumnTypes?: Array; + exportColumnExtras?: Array; + insertMode?: SqlInsertMode; + }) => queryStore.buildQueryResultExportRequest(activeTab.id, options) " :all-export-results="allResultExportSheets" :export-file-base-name="activeTab.title" diff --git a/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts index 1baabcb313..006778085b 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridExport.sqlProgress.spec.ts @@ -140,6 +140,25 @@ describe("query result SQL export progress", () => { expect(api.startQueryResultExport).toHaveBeenCalledWith(expect.objectContaining({ format: "sql", exportTableName: "users", exportColumnTypes: ["int4", "text"] }), expect.any(Function)); }); + it("forwards result column EXTRA metadata so identity INSERT exports can be replayed", async () => { + const state = useDataGridExport( + createOptions({ + tableMeta: computed(() => ({ + tableName: "users", + primaryKeys: ["id"], + columns: [ + { name: "id", data_type: "int", is_nullable: false, is_primary_key: true, extra: "identity(1,1)" }, + { name: "name", data_type: "text", is_nullable: true }, + ], + })), + }), + ); + + await state.exportSql(); + + expect(api.startQueryResultExport).toHaveBeenCalledWith(expect.objectContaining({ format: "sql", exportColumnTypes: ["int4", "text"], exportColumnExtras: ["identity(1,1)", null] }), expect.any(Function)); + }); + it("routes incomplete query-result JSON exports through the streaming backend", async () => { const state = useDataGridExport(createOptions()); diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index 7491b9e5a5..fcdefbcdba 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -123,6 +123,7 @@ export interface UseDataGridExportOptions { includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; + exportColumnExtras?: Array; insertMode?: SqlInsertMode; }) => Promise; /** @@ -1277,6 +1278,18 @@ export function useDataGridExport(options: UseDataGridExportOptions) { return autoGenerated.length > 0 ? autoGenerated : undefined; } + /** + * 导出列对应的原表 EXTRA 元数据(SQL Server/Dameng 的 identity 等)。后端据此 + * 给导出的 INSERT 包上 `SET IDENTITY_INSERT`,否则回放时报 SQL Server 544。 + * 拿不到表元数据时返回 undefined,后端保持“未知”语义,不额外查一次元数据。 + */ + function sqlExportColumnExtras(columnNames: string[]): Array | undefined { + const metaColumns = tableMeta.value?.columns; + if (!metaColumns?.length) return undefined; + const extras = columnNames.map((column) => metaColumns.find((meta) => normalizeColumnName(meta.name) === normalizeColumnName(column))?.extra ?? null); + return extras.some((extra) => !!extra) ? extras : undefined; + } + /** 供导出请求使用:把“不含主键”设置转换成后端请求字段。 */ function sqlExportPrimaryKeyOptions(): { excludePrimaryKeys?: boolean; primaryKeys?: string[] } { const excludeColumns = sqlExportExcludedColumns(); @@ -1347,6 +1360,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { csvQuoteMode: editorSettings.csvQuoteMode, columns: columns.value, columnTypes: columnTypes.value, + ...(format === "sql" ? { columnExtras: sqlExportColumnExtras(columns.value) } : {}), columnComments: format === "xlsx" ? buildXlsxHeaderOverrides(columns.value, visibleXlsxColumnComments.value, headerMode) : undefined, primaryKeys: meta.primaryKeys, ...sqlExportPrimaryKeyOptions(), @@ -1416,6 +1430,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { includeSqlSheet, exportTableName: format === "sql" ? tableMeta.value?.tableName : undefined, exportColumnTypes: format === "sql" ? allColumnTypes.value?.map((type) => type ?? null) : undefined, + exportColumnExtras: format === "sql" ? sqlExportColumnExtras(allColumns.value) : undefined, ...(format === "sql" && insertMode ? { insertMode } : {}), }); const columnComments = format === "xlsx" ? buildXlsxHeaderOverrides(allColumns.value, allXlsxColumnComments.value, headerMode) : undefined; @@ -1586,6 +1601,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { tableName: tableMeta.value?.tableName || "table_name", columns: exportData.columns, columnTypes: exportData.columnTypes, + columnExtras: exportData.columnExtras, spatialColumns: exportData.spatialColumns, spatialValues: exportData.spatialValues, rows: exportData.rows, @@ -1628,6 +1644,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { tableName: tableMeta.value?.tableName || "table_name", columns: exportData.columns, columnTypes: exportData.columnTypes, + columnExtras: exportData.columnExtras, spatialColumns: exportData.spatialColumns, spatialValues: exportData.spatialValues, rows: exportData.rows, @@ -1650,6 +1667,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { function sqlInsertExportData(result: { columns: string[]; rows: CellValue[][]; spatialColumns?: QueryResult["spatial_columns"]; spatialValues?: QueryResult["spatial_values"] }): { columns: string[]; columnTypes?: Array; + columnExtras?: Array; spatialColumns?: QueryResult["spatial_columns"]; spatialValues?: QueryResult["spatial_values"]; rows: CellValue[][]; @@ -1657,6 +1675,8 @@ export function useDataGridExport(options: UseDataGridExportOptions) { const exportColumns = context.value === "table-data" && tableMeta.value ? effectiveColumns(sourceColumns.value, result.columns) : result.columns; const columnIndexes = exportColumns.map((column, index) => ({ column, index })).filter((item): item is { column: string; index: number } => !!item.column); const exportColumnTypes = columnTypes.value?.length === result.columns.length ? columnTypes.value : undefined; + const metaColumns = tableMeta.value?.columns; + const exportColumnExtras = metaColumns?.length ? exportColumns.map((column) => (column ? (metaColumns.find((meta) => normalizeColumnName(meta.name) === normalizeColumnName(column))?.extra ?? null) : null)) : undefined; const indexBySource = new Map(columnIndexes.map((item, index) => [item.index, index])); const spatialColumns = result.spatialColumns?.flatMap((column) => { const columnIndex = indexBySource.get(column.column_index); @@ -1665,6 +1685,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { return { columns: columnIndexes.map((item) => item.column), columnTypes: exportColumnTypes ? columnIndexes.map((item) => exportColumnTypes[item.index]) : undefined, + columnExtras: exportColumnExtras?.some((extra) => !!extra) ? columnIndexes.map((item) => exportColumnExtras[item.index] ?? null) : undefined, ...(spatialColumns?.length ? { spatialColumns } : {}), ...(result.spatialValues?.length ? { spatialValues: result.spatialValues.map((row) => columnIndexes.map((item) => row[item.index] ?? null)) } : {}), rows: result.rows.map((row) => columnIndexes.map((item) => row[item.index] ?? null)), diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index c679cc3a82..0a688c053d 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -5469,6 +5469,8 @@ export interface TableExportRequest { csvQuoteMode?: CsvQuoteMode; columns?: string[]; columnTypes?: Array; + /** 与 `columns` 对齐的列 EXTRA 元数据(identity 等),用于 SQL INSERT 导出的 `SET IDENTITY_INSERT`。 */ + columnExtras?: Array; columnComments?: Array | null; primaryKeys?: string[]; /** 导出 SQL 时是否排除主键列(对应数据提取设置里的“排除主键”)。 */ @@ -5531,6 +5533,11 @@ export interface QueryResultExportRequest { dateTimeFormat?: string; exportTableName?: string; exportColumnTypes?: Array; + /** + * 结果列对应的原表 EXTRA 元数据(identity 等)。后端据此为 SQL INSERT 导出 + * 补上 `SET IDENTITY_INSERT` 包裹,缺省表示未知。 + */ + exportColumnExtras?: Array; numericColumnRightAlign?: boolean; columnComments?: Array | null; autoFilter?: boolean; diff --git a/apps/desktop/src/lib/export/exportFormats.ts b/apps/desktop/src/lib/export/exportFormats.ts index 404605a6d5..800ff986f7 100644 --- a/apps/desktop/src/lib/export/exportFormats.ts +++ b/apps/desktop/src/lib/export/exportFormats.ts @@ -36,6 +36,8 @@ export interface FormatSqlInsertOptions { qualifiedTableName?: string; columns: string[]; columnTypes?: Array; + /** 与 `columns` 对齐的列 EXTRA 元数据,用于 identity 列的 `SET IDENTITY_INSERT` 包裹。 */ + columnExtras?: Array; spatialColumns?: QueryResult["spatial_columns"]; spatialValues?: QueryResult["spatial_values"]; rows: ExportCellValue[][]; diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 319e547e5e..1d4d72016e 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -168,6 +168,7 @@ interface BuildQueryResultExportRequestOptions { includeSqlSheet?: boolean; exportTableName?: string; exportColumnTypes?: Array; + exportColumnExtras?: Array; insertMode?: SqlInsertMode; } @@ -8768,6 +8769,7 @@ export const useQueryStore = defineStore("query", () => { executionId: uuid(), exportTableName: options.exportTableName, exportColumnTypes: options.exportColumnTypes, + exportColumnExtras: options.exportColumnExtras, numericColumnRightAlign: settings.numericColumnRightAlign, identifierQuote: connStore.connectionIdentifierQuote(location.connectionId), }; diff --git a/crates/dbx-core/src/data/query_result_export.rs b/crates/dbx-core/src/data/query_result_export.rs index fe08ec3096..9df4879796 100644 --- a/crates/dbx-core/src/data/query_result_export.rs +++ b/crates/dbx-core/src/data/query_result_export.rs @@ -112,6 +112,13 @@ pub struct QueryResultExportRequest { /// Frontend sends these in original full-query column order. #[serde(default, skip_serializing_if = "Option::is_none")] pub export_column_types: Option>>, + /// Column EXTRA metadata (for example SQL Server/Dameng `identity`) for SQL + /// INSERT export, sent by the frontend from the result's table metadata. + /// Without it the exported INSERT writes identity values without + /// `SET IDENTITY_INSERT` and fails with SQL Server error 544 on replay. + /// Entries align with `export_column_types` (original full-query order). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub export_column_extras: Option>>, #[serde(default)] pub numeric_column_right_align: bool, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -324,6 +331,7 @@ struct SqlInsertWriter { insert_mode: SqlInsertMode, columns: Vec, column_types: Vec>, + column_extras: Vec>, spatial_columns: Vec, database_type: DatabaseType, schema: Option, @@ -412,6 +420,7 @@ impl SqlInsertWriter { insert_mode: request.insert_mode, columns: Vec::new(), column_types: Vec::new(), + column_extras: Vec::new(), spatial_columns: Vec::new(), database_type: request.database_type, schema: request.schema.clone(), @@ -432,6 +441,9 @@ impl SqlInsertWriter { request: &QueryResultExportRequest, ) { self.column_types = sql_insert_column_types(request, result_column_types); + // Column extras come from the request and align with the result columns the + // same way `export_column_types` does; missing entries simply mean "unknown". + self.column_extras = request.export_column_extras.clone().unwrap_or_default(); self.spatial_columns = spatial_columns.to_vec(); self.columns = columns; } @@ -458,7 +470,7 @@ impl SqlInsertWriter { qualified_table_name: None, columns: self.columns.clone(), column_types: self.column_types.clone(), - column_extras: Vec::new(), + column_extras: self.column_extras.clone(), spatial_columns: self.spatial_columns.clone(), spatial_values: mem::take(&mut self.pending_spatial_values), rows: mem::take(&mut self.pending_rows), @@ -2091,6 +2103,91 @@ mod tests { assert!(!stream_export_was_cancelled("network failure", false, false)); } + fn sqlserver_identity_export_request( + file_path: &std::path::Path, + export_column_extras: Option>>, + ) -> QueryResultExportRequest { + QueryResultExportRequest { + export_id: "export-identity".to_string(), + connection_id: "conn-1".to_string(), + database: "dbx_test".to_string(), + schema: Some("dbo".to_string()), + catalog: None, + sql: "SELECT * FROM [dbo].[gen_table]".to_string(), + query_base_sql: "SELECT * FROM [dbo].[gen_table]".to_string(), + setup_sql: Vec::new(), + database_type: DatabaseType::SqlServer, + 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("gen_table".to_string()), + export_column_types: Some(vec![Some("int".to_string()), Some("nvarchar(200)".to_string())]), + export_column_extras, + numeric_column_right_align: false, + column_comments: None, + auto_filter: None, + identifier_quote: None, + insert_mode: SqlInsertMode::Batch, + csv_quote_mode: Default::default(), + exclude_primary_keys: false, + primary_keys: Vec::new(), + } + } + + #[test] + fn sql_insert_writer_wraps_sqlserver_identity_columns_with_identity_insert() { + let dir = tempfile::tempdir().expect("temp dir"); + let file_path = dir.path().join("gen_table.sql"); + let request = + sqlserver_identity_export_request(&file_path, Some(vec![Some("identity(1,1)".to_string()), None])); + + let mut writer = SqlInsertWriter::create(&request).expect("create sql insert writer"); + writer.set_columns( + vec!["table_id".to_string(), "table_name".to_string()], + &["int".to_string(), "nvarchar(200)".to_string()], + &[], + &request, + ); + writer.write_row(vec![json!(23), json!("t_destype")], None).expect("write export row"); + writer.finish().expect("finish export"); + + assert_eq!( + std::fs::read_to_string(&file_path).expect("read sql export"), + "SET IDENTITY_INSERT [dbo].[gen_table] ON;\nINSERT INTO [dbo].[gen_table] ([table_id], [table_name]) VALUES (23, N't_destype');\nSET IDENTITY_INSERT [dbo].[gen_table] OFF;\n" + ); + } + + #[test] + fn sql_insert_writer_skips_identity_wrapper_without_column_extras() { + let dir = tempfile::tempdir().expect("temp dir"); + let file_path = dir.path().join("gen_table.sql"); + let request = sqlserver_identity_export_request(&file_path, None); + + let mut writer = SqlInsertWriter::create(&request).expect("create sql insert writer"); + writer.set_columns( + vec!["table_id".to_string(), "table_name".to_string()], + &["int".to_string(), "nvarchar(200)".to_string()], + &[], + &request, + ); + writer.write_row(vec![json!(23), json!("t_destype")], 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 [dbo].[gen_table] ([table_id], [table_name]) VALUES (23, N't_destype');\n" + ); + } + #[test] fn sql_insert_writer_omits_excluded_primary_key_columns() { let dir = tempfile::tempdir().expect("temp dir"); @@ -2119,6 +2216,7 @@ mod tests { date_time_format: None, export_table_name: Some("users".to_string()), export_column_types: None, + export_column_extras: None, numeric_column_right_align: false, column_comments: None, auto_filter: None, @@ -2201,6 +2299,7 @@ mod tests { csv_quote_mode: CsvQuoteMode::All, export_table_name: None, export_column_types: None, + export_column_extras: None, numeric_column_right_align: false, column_comments: None, auto_filter: None, diff --git a/crates/dbx-core/src/data/table_export.rs b/crates/dbx-core/src/data/table_export.rs index 8113dd2ba4..396b948cfd 100644 --- a/crates/dbx-core/src/data/table_export.rs +++ b/crates/dbx-core/src/data/table_export.rs @@ -89,6 +89,12 @@ pub struct TableExportRequest { pub columns: Option>, #[serde(default)] pub column_types: Option>>, + /// Column EXTRA metadata for `columns`, supplied by the data grid. SQL + /// INSERT exports need it for dialect rules such as SQL Server/Dameng + /// `SET IDENTITY_INSERT` (error 544) and MySQL generated-column skipping; + /// without it the export cannot know those rules without a metadata query. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub column_extras: Option>>, #[serde(default)] pub primary_keys: Option>, /// 导出 SQL 时是否排除主键列(对应前端数据提取设置里的“排除主键”)。 @@ -229,6 +235,24 @@ fn requested_mysql_sql_export_needs_column_metadata(database_type: DatabaseType, database_type == DatabaseType::Mysql && format.eq_ignore_ascii_case("sql") } +/// Column EXTRA values supplied by the data grid, filtered and aligned exactly +/// like [`resolve_requested_export_columns`] keeps name/type metadata aligned. +fn resolve_requested_column_extras_by_position( + database_type: DatabaseType, + columns: &[String], + column_extras: Option<&[Option]>, +) -> Vec> { + let Some(column_extras) = column_extras else { + return Vec::new(); + }; + columns + .iter() + .enumerate() + .filter(|(_, column)| !is_internal_export_column(Some(database_type), column)) + .map(|(index, _)| column_extras.get(index).cloned().flatten()) + .collect() +} + fn resolve_requested_export_column_types( requested_columns: &[String], requested_column_types: &[Option], @@ -1540,7 +1564,14 @@ async fn export_table_data_core_inner( resolve_requested_export_column_extras(&col_names, &table_columns), ) } else { - (requested_column_types, Vec::new()) + ( + requested_column_types, + resolve_requested_column_extras_by_position( + db_type, + requested_columns, + request.column_extras.as_deref(), + ), + ) }; (col_names, column_types, column_extras, primary_keys) } else { @@ -2386,6 +2417,7 @@ mod tests { insert_mode: Default::default(), columns: Some(vec!["id".to_string(), "name".to_string()]), column_types: Some(vec![Some("INTEGER".to_string()), Some("VARCHAR".to_string())]), + column_extras: None, primary_keys: Some(vec!["id".to_string()]), exclude_primary_keys: false, where_input: None, @@ -2552,6 +2584,7 @@ mod tests { exclude_primary_keys: false, columns: None, column_types: None, + column_extras: None, primary_keys: None, where_input: None, order_by: None, @@ -2717,6 +2750,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: Some("WHERE temperature > 1".to_string()), @@ -2777,6 +2811,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -2815,6 +2850,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -2848,6 +2884,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -2889,6 +2926,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -2942,6 +2980,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: Some("WHERE status = 'active'".to_string()), @@ -2989,6 +3028,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -3064,6 +3104,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -3121,6 +3162,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -3180,6 +3222,7 @@ mod tests { insert_mode: Default::default(), columns: None, column_types: None, + column_extras: None, primary_keys: None, exclude_primary_keys: false, where_input: None, @@ -3265,6 +3308,26 @@ mod tests { ); } + #[test] + fn requested_column_extras_stay_aligned_with_filtered_columns() { + let columns = vec!["table_id".to_string(), "table_name".to_string()]; + let extras = vec![Some("identity(1,1)".to_string()), None]; + assert_eq!( + resolve_requested_column_extras_by_position(DatabaseType::SqlServer, &columns, Some(&extras)), + vec![Some("identity(1,1)".to_string()), None] + ); + assert!(resolve_requested_column_extras_by_position(DatabaseType::SqlServer, &columns, None).is_empty()); + + // The synthetic ROWID column never reaches the export, so its EXTRA entry + // must be dropped with it instead of shifting the remaining columns. + let oracle_columns = vec!["__DBX_ROWID".to_string(), "id".to_string()]; + let oracle_extras = vec![Some("synthetic".to_string()), Some("identity".to_string())]; + assert_eq!( + resolve_requested_column_extras_by_position(DatabaseType::Oracle, &oracle_columns, Some(&oracle_extras)), + vec![Some("identity".to_string())] + ); + } + #[test] fn agent_table_read_unsupported_detects_old_agent_errors() { assert!(is_agent_table_read_unsupported("Agent RPC error (-1): unknown method: start_table_read")); 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 8d4c35fea3..9f9a757cc5 100644 --- a/crates/dbx-core/tests/live_clickhouse_query_result_export.rs +++ b/crates/dbx-core/tests/live_clickhouse_query_result_export.rs @@ -105,6 +105,7 @@ async fn live_clickhouse_query_result_export_xlsx_streams_random_order_query_onc csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 289af7aa5d..61d1e8d16a 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -480,6 +480,7 @@ async fn live_mysql_query_result_export_xlsx_streams_single_query_without_duplic csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -573,6 +574,7 @@ async fn live_mysql_csv_temporal_export_round_trip_preserves_dbx_force_text_valu csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -690,6 +692,7 @@ async fn live_mysql_xlsx_export_can_outlive_query_timeout_while_rows_keep_arrivi csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, 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 6a2697586a..31cb3b0299 100644 --- a/crates/dbx-core/tests/live_postgres_query_result_export.rs +++ b/crates/dbx-core/tests/live_postgres_query_result_export.rs @@ -152,6 +152,7 @@ async fn live_postgres_query_result_export_uses_single_streamed_query() { csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -238,6 +239,7 @@ async fn live_postgres_query_result_xlsx_preserves_temporal_cell_types() { csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -313,6 +315,7 @@ async fn live_postgres_numeric_xlsx_ignores_fractional_trailing_zeros() { csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -393,6 +396,7 @@ async fn live_postgres_truncated_batch_result_export_replays_safe_temp_setup() { csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -475,6 +479,7 @@ async fn live_postgres_xlsx_export_can_outlive_query_timeout_while_rows_keep_arr csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, @@ -548,6 +553,7 @@ async fn live_postgres_stream_still_times_out_without_progress_and_recovers() { csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, diff --git a/crates/dbx-core/tests/live_sqlserver_completion.rs b/crates/dbx-core/tests/live_sqlserver_completion.rs index ccb0b6b085..fbcffa18cc 100644 --- a/crates/dbx-core/tests/live_sqlserver_completion.rs +++ b/crates/dbx-core/tests/live_sqlserver_completion.rs @@ -1214,6 +1214,7 @@ async fn live_sqlserver_query_result_export_streams_cte_query_to_csv() { csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, 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 388f94de26..d8d511857b 100644 --- a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs +++ b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs @@ -131,6 +131,7 @@ async fn live_sqlserver_xlsx_export_can_outlive_query_timeout_while_rows_keep_ar csv_quote_mode: Default::default(), export_table_name: None, export_column_types: None, + export_column_extras: None, column_comments: None, auto_filter: None, identifier_quote: None, diff --git a/crates/dbx-sql/src/data_grid_sql.rs b/crates/dbx-sql/src/data_grid_sql.rs index 7d2a6688d5..5f2298ebed 100644 --- a/crates/dbx-sql/src/data_grid_sql.rs +++ b/crates/dbx-sql/src/data_grid_sql.rs @@ -606,22 +606,37 @@ pub fn build_data_grid_copy_insert_statement(options: DataGridCopyInsertStatemen ) }) .collect::>(); - if options.insert_mode == DataGridCopyInsertMode::RowByRow + let statements = if options.insert_mode == DataGridCopyInsertMode::RowByRow || options.database_type.is_some_and(uses_single_row_insert_statements) { + value_rows.iter().map(|values| format!("INSERT INTO {table} ({columns}) VALUES {values};")).collect::>() + } else { + vec![format!( + "INSERT INTO {table} ({columns}) VALUES{}{};", + if value_rows.len() == 1 { " " } else { "\n" }, + value_rows.join(",\n") + )] + }; + // SQL Server and Dameng reject explicit values for identity columns unless the + // statement runs between `SET IDENTITY_INSERT ON` and `OFF` (SQL Server + // error 544), so a copied INSERT that carries the identity column must ship the + // wrapper. Each statement is wrapped on its own so row-by-row copies stay + // individually executable, matching the SQL export path. + let needs_identity_insert_wrapper = + matches!(options.database_type, Some(DatabaseType::SqlServer | DatabaseType::Dameng)) + && insert_columns.iter().any(|(_, _, info)| info.as_ref().is_some_and(is_auto_generated_column)); + if needs_identity_insert_wrapper { return Some( - value_rows + statements .iter() - .map(|values| format!("INSERT INTO {table} ({columns}) VALUES {values};")) + .map(|statement| { + format!("SET IDENTITY_INSERT {table} ON;\n{statement}\nSET IDENTITY_INSERT {table} OFF;") + }) .collect::>() .join("\n"), ); } - Some(format!( - "INSERT INTO {table} ({columns}) VALUES{}{};", - if value_rows.len() == 1 { " " } else { "\n" }, - value_rows.join(",\n") - )) + Some(statements.join("\n")) } pub fn build_data_grid_context_filter_condition(options: DataGridContextFilterConditionOptions) -> Option { @@ -4629,6 +4644,94 @@ mod tests { ); } + fn sqlserver_identity_copy_insert_options( + rows: Vec>, + insert_mode: DataGridCopyInsertMode, + ) -> DataGridCopyInsertStatementOptions { + DataGridCopyInsertStatementOptions { + database_type: Some(DatabaseType::SqlServer), + identifier_quote: None, + table_meta: Some(DataGridTableMeta { + catalog: None, + database: Some("dbx_test".to_string()), + schema: Some("dbo".to_string()), + table_name: "gen_table".to_string(), + primary_keys: vec!["table_id".to_string()], + columns: Some(vec![ + column("table_id", "int", false, Some("identity")), + column("table_name", "nvarchar(200)", false, None), + ]), + }), + columns: vec!["table_id".to_string(), "table_name".to_string()], + column_types: None, + source_columns: None, + rows, + exclude_primary_keys: false, + include_computed_columns: false, + include_database_name: true, + insert_mode, + } + } + + #[test] + fn sqlserver_copy_insert_wraps_identity_columns_with_identity_insert() { + let statement = build_data_grid_copy_insert_statement(sqlserver_identity_copy_insert_options( + vec![vec![json!(1), json!("t_destype")]], + DataGridCopyInsertMode::Merged, + )); + assert_eq!( + statement.as_deref(), + Some( + "SET IDENTITY_INSERT [dbo].[gen_table] ON;\nINSERT INTO [dbo].[gen_table] ([table_id], [table_name]) VALUES (1, N't_destype');\nSET IDENTITY_INSERT [dbo].[gen_table] OFF;" + ) + ); + } + + #[test] + fn sqlserver_row_by_row_copy_insert_wraps_every_statement() { + let statement = build_data_grid_copy_insert_statement(sqlserver_identity_copy_insert_options( + vec![vec![json!(1), json!("t_destype")], vec![json!(2), json!("t_user")]], + DataGridCopyInsertMode::RowByRow, + )); + assert_eq!( + statement.as_deref(), + Some( + "SET IDENTITY_INSERT [dbo].[gen_table] ON;\nINSERT INTO [dbo].[gen_table] ([table_id], [table_name]) VALUES (1, N't_destype');\nSET IDENTITY_INSERT [dbo].[gen_table] OFF;\nSET IDENTITY_INSERT [dbo].[gen_table] ON;\nINSERT INTO [dbo].[gen_table] ([table_id], [table_name]) VALUES (2, N't_user');\nSET IDENTITY_INSERT [dbo].[gen_table] OFF;" + ) + ); + } + + #[test] + fn sqlserver_copy_insert_omits_identity_wrapper_without_identity_columns() { + let mut options = sqlserver_identity_copy_insert_options( + vec![vec![json!(1), json!("t_destype")]], + DataGridCopyInsertMode::Merged, + ); + options.table_meta.as_mut().expect("table meta").columns = + Some(vec![column("table_id", "int", false, None), column("table_name", "nvarchar(200)", false, None)]); + let statement = build_data_grid_copy_insert_statement(options); + assert_eq!( + statement.as_deref(), + Some("INSERT INTO [dbo].[gen_table] ([table_id], [table_name]) VALUES (1, N't_destype');") + ); + } + + #[test] + fn dameng_copy_insert_wraps_identity_columns_with_identity_insert() { + let mut options = sqlserver_identity_copy_insert_options( + vec![vec![json!(1), json!("t_destype")]], + DataGridCopyInsertMode::Merged, + ); + options.database_type = Some(DatabaseType::Dameng); + let statement = build_data_grid_copy_insert_statement(options); + assert_eq!( + statement.as_deref(), + Some( + "SET IDENTITY_INSERT \"dbo\".\"gen_table\" ON;\nINSERT INTO \"dbo\".\"gen_table\" (\"table_id\", \"table_name\") VALUES (1, 't_destype');\nSET IDENTITY_INSERT \"dbo\".\"gen_table\" OFF;" + ) + ); + } + #[test] fn mysql_copy_statements_preserve_blob_hex_literals() { let table_meta = DataGridTableMeta { diff --git a/docs/screenshots/issue-9734-after-copy-sql.png b/docs/screenshots/issue-9734-after-copy-sql.png new file mode 100644 index 0000000000..bc0958a318 Binary files /dev/null and b/docs/screenshots/issue-9734-after-copy-sql.png differ diff --git a/docs/screenshots/issue-9734-after-replay-success.png b/docs/screenshots/issue-9734-after-replay-success.png new file mode 100644 index 0000000000..c3720430ac Binary files /dev/null and b/docs/screenshots/issue-9734-after-replay-success.png differ diff --git a/docs/screenshots/issue-9734-before-error-544.png b/docs/screenshots/issue-9734-before-error-544.png new file mode 100644 index 0000000000..b321b1bfdf Binary files /dev/null and b/docs/screenshots/issue-9734-before-error-544.png differ diff --git a/docs/screenshots/issue-9734-query-export-done.png b/docs/screenshots/issue-9734-query-export-done.png new file mode 100644 index 0000000000..2211689af9 Binary files /dev/null and b/docs/screenshots/issue-9734-query-export-done.png differ diff --git a/docs/screenshots/issue-9734-table-export-done.png b/docs/screenshots/issue-9734-table-export-done.png new file mode 100644 index 0000000000..f32a539810 Binary files /dev/null and b/docs/screenshots/issue-9734-table-export-done.png differ