From 1a4224690958cee832cd62fe3185fc5144c580cd Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 4 Aug 2026 00:10:22 +0800 Subject: [PATCH] fix(import): disambiguate duplicate source headers Closes #5258 --- crates/dbx-core/src/table_import.rs | 212 ++++++++++++++++++++++------ 1 file changed, 168 insertions(+), 44 deletions(-) diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index fc52a56ec..7774d2c2b 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -389,6 +389,24 @@ pub fn normalize_header(value: &str, index: usize) -> String { } } +fn unique_import_headers(headers: impl IntoIterator) -> Vec { + let mut seen = HashSet::new(); + let mut next_suffix = HashMap::::new(); + headers + .into_iter() + .map(|header| { + let suffix = next_suffix.entry(header.to_lowercase()).or_default(); + loop { + let candidate = if *suffix == 0 { header.clone() } else { format!("{header}_{suffix}") }; + *suffix += 1; + if seen.insert(candidate.to_lowercase()) { + break candidate; + } + } + }) + .collect() +} + #[derive(Debug, Clone, Copy)] pub struct DelimitedParseConfig { pub delimiter: u8, @@ -895,11 +913,12 @@ fn parse_csv_reader_inner( index += 1; let row_number = index; if config.row_range.title_row == Some(row_number) { - columns = record - .iter() - .enumerate() - .map(|(index, header)| normalize_header(header.trim_start_matches('\u{feff}'), index)) - .collect(); + columns = unique_import_headers( + record + .iter() + .enumerate() + .map(|(index, header)| normalize_header(header.trim_start_matches('\u{feff}'), index)), + ); continue; } if row_number < config.row_range.data_start_row { @@ -2072,18 +2091,14 @@ fn parse_xlsx_preview_file_with_options( let column_count = end_column.saturating_sub(start_column).saturating_add(1); let mut columns = if let Some(title_row) = row_range.title_row { let absolute_title_row = start_row.saturating_add(title_row.saturating_sub(1)); - (0..column_count) - .map(|index| { - let column = start_column + index; - let value = raw_cells - .get(&(absolute_title_row, column)) - .map(|cell| { - xlsx_preview_cell_value(cell, &shared_strings, &styles, date_1904, empty_string_as_null) - }) - .unwrap_or(serde_json::Value::Null); - normalize_header(&xlsx_preview_cell_label(&value), index) - }) - .collect::>() + unique_import_headers((0..column_count).map(|index| { + let column = start_column + index; + let value = raw_cells + .get(&(absolute_title_row, column)) + .map(|cell| xlsx_preview_cell_value(cell, &shared_strings, &styles, date_1904, empty_string_as_null)) + .unwrap_or(serde_json::Value::Null); + normalize_header(&xlsx_preview_cell_label(&value), index) + })) } else { Vec::new() }; @@ -2157,11 +2172,9 @@ where for (index, source_row) in range.rows().enumerate() { let row_number = index + 1; if row_range.title_row == Some(row_number) { - return source_row - .iter() - .enumerate() - .map(|(index, cell)| normalize_header(&cell_label(cell, None), index)) - .collect(); + return unique_import_headers( + source_row.iter().enumerate().map(|(index, cell)| normalize_header(&cell_label(cell, None), index)), + ); } let row_is_within_range = match row_range.last_data_row { Some(last) => row_number <= last, @@ -2442,11 +2455,12 @@ impl XlsxStreamRowsState { if self.columns.is_empty() { let column_count = self.declared_column_count.unwrap_or(values.len()).max(values.len()); values.resize(column_count, serde_json::Value::Null); - self.columns = values - .iter() - .enumerate() - .map(|(index, value)| normalize_header(&xlsx_preview_cell_label(value), index)) - .collect(); + self.columns = unique_import_headers( + values + .iter() + .enumerate() + .map(|(index, value)| normalize_header(&xlsx_preview_cell_label(value), index)), + ); } return Ok(()); } @@ -2831,18 +2845,14 @@ where for (index, source_row) in range.rows().enumerate() { let row_number = index + 1; if row_range.title_row == Some(row_number) { - columns = source_row - .iter() - .enumerate() - .map(|(index, cell)| { - // Calamine rows are relative to the used range, while XLSX style coordinates are worksheet-absolute. - let cell_position = (range_start_row + row_number, range_start_column + index + 1); - normalize_header( - &cell_label(cell, cell_styles.get(&cell_position).and_then(|style| style.temporal_kind)), - index, - ) - }) - .collect(); + columns = unique_import_headers(source_row.iter().enumerate().map(|(index, cell)| { + // Calamine rows are relative to the used range, while XLSX style coordinates are worksheet-absolute. + let cell_position = (range_start_row + row_number, range_start_column + index + 1); + normalize_header( + &cell_label(cell, cell_styles.get(&cell_position).and_then(|style| style.temporal_kind)), + index, + ) + })); continue; } if row_number < row_range.data_start_row { @@ -4386,11 +4396,12 @@ fn delimited_columns_and_first_record( let record = record.map_err(|e| e.to_string())?; let row_number = index + 1; if config.row_range.title_row == Some(row_number) { - columns = record - .iter() - .enumerate() - .map(|(index, header)| normalize_header(header.trim_start_matches('\u{feff}'), index)) - .collect(); + columns = unique_import_headers( + record + .iter() + .enumerate() + .map(|(index, header)| normalize_header(header.trim_start_matches('\u{feff}'), index)), + ); continue; } if row_number < config.row_range.data_start_row { @@ -6355,6 +6366,21 @@ mod tests { assert_eq!(plan.column_types, vec![Some("BIGINT".to_string()), Some("VARCHAR(255)".to_string())]); } + #[test] + fn duplicate_import_headers_receive_stable_case_insensitive_suffixes() { + assert_eq!( + unique_import_headers(["Name", "name", "name", "name_1"].into_iter().map(str::to_string)), + vec!["Name", "name_1", "name_2", "name_1_1"] + ); + assert_eq!( + unique_import_headers(["id", "display_name"].into_iter().map(str::to_string)), + vec!["id", "display_name"] + ); + let many = unique_import_headers(std::iter::repeat_n("value".to_string(), 10_000)); + assert_eq!(many.len(), 10_000); + assert_eq!(many.last().map(String::as_str), Some("value_9999")); + } + #[test] fn prepared_import_source_is_reused_only_while_fingerprint_matches() { let legacy_prepared: TableImportPreparedSource = serde_json::from_value(serde_json::json!({ @@ -6630,6 +6656,47 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn xlsx_duplicate_headers_are_unique_in_preview_parse_and_streaming() { + let path = + std::env::temp_dir().join(format!("dbx-table-import-duplicate-headers-{}.xlsx", uuid::Uuid::new_v4())); + let sheet_xml = r#" + + + + + name + name + + + Ada + Lovelace + + +"#; + std::fs::write(&path, build_preview_test_xlsx(sheet_xml, None)).unwrap(); + let options = TableImportParseOptions::default(); + + let parsed = parse_xlsx_file_with_options(&path.to_string_lossy(), &options, 10).unwrap(); + let (preview, _) = parse_xlsx_preview_file_with_options(&path.to_string_lossy(), &options, 10).unwrap(); + let (sender, mut receiver) = tokio::sync::mpsc::channel(16); + stream_xlsx_rows_to_channel(&path.to_string_lossy(), &options, 500, None, HashSet::new(), false, sender) + .unwrap(); + + let mut streamed_columns = Vec::new(); + while let Some(message) = receiver.blocking_recv() { + if let XlsxStreamMessage::Header(columns) = message.unwrap() { + streamed_columns = columns; + } + } + + assert_eq!(parsed.columns, vec!["name", "name_1"]); + assert_eq!(preview.columns, parsed.columns); + assert_eq!(streamed_columns, parsed.columns); + assert_eq!(parsed.rows, vec![vec![serde_json::json!("Ada"), serde_json::json!("Lovelace")]]); + let _ = std::fs::remove_file(path); + } + #[test] fn xlsx_preserves_explicit_empty_strings_when_configured() { let options = @@ -6837,6 +6904,41 @@ mod tests { ); } + #[test] + fn duplicate_csv_headers_map_to_distinct_source_indexes() { + let parsed = parse_csv_bytes(b"name,name\nAda,Lovelace\n", 10).unwrap(); + let mappings = vec![ + TableImportColumnMapping { + source_column: "name".to_string(), + target_column: "first_name".to_string(), + target_data_type: None, + }, + TableImportColumnMapping { + source_column: "name_1".to_string(), + target_column: "last_name".to_string(), + target_data_type: None, + }, + ]; + + assert_eq!(parsed.columns, vec!["name", "name_1"]); + let batch = build_import_insert_batch_from_rows( + &parsed.rows, + &parsed.columns, + &mappings, + &[], + "people", + "public", + &DatabaseType::Postgres, + ) + .unwrap() + .unwrap(); + + assert_eq!( + batch.sql, + "INSERT INTO \"public\".\"people\" (\"first_name\", \"last_name\") VALUES\n('Ada', 'Lovelace')" + ); + } + #[test] fn auto_detects_and_parses_gbk_csv() { let (bytes, _, had_errors) = encoding_rs::GBK.encode("id,name\n1,中文\n2,上海\n"); @@ -7107,6 +7209,28 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn streaming_csv_uses_the_same_unique_duplicate_headers_as_preview() { + let path = std::env::temp_dir().join(format!("dbx-table-import-duplicate-stream-{}.csv", uuid::Uuid::new_v4())); + std::fs::write(&path, b"name,name\nAda,Lovelace\n").unwrap(); + let options = TableImportParseOptions::default(); + let preview = parse_csv_bytes(b"name,name\nAda,Lovelace\n", 10).unwrap(); + let (sender, mut receiver) = tokio::sync::mpsc::channel(16); + + stream_delimited_rows_to_channel(&path.to_string_lossy(), TableImportSourceFormat::Csv, &options, 500, sender) + .unwrap(); + + let messages = + std::iter::from_fn(|| receiver.blocking_recv()).map(|message| message.unwrap()).collect::>(); + assert!( + matches!(messages.first(), Some(DelimitedStreamMessage::Header(columns)) if columns == &preview.columns) + ); + assert!(messages.iter().any(|message| { + matches!(message, DelimitedStreamMessage::Rows { rows, .. } if rows == &vec![vec![serde_json::json!("Ada"), serde_json::json!("Lovelace")]]) + })); + let _ = std::fs::remove_file(path); + } + #[test] fn legacy_xls_uses_a_stricter_non_streaming_file_limit() { let path = std::env::temp_dir().join(format!("dbx-table-import-limit-{}.xls", uuid::Uuid::new_v4()));