From 4e14ac86769c24f9185e9df41c54058a16567a98 Mon Sep 17 00:00:00 2001 From: QuakeWang <45645138+QuakeWang@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:15:08 +0800 Subject: [PATCH] feat: add strict CSV header validation (#8426) * feat: add strict CSV header validation Signed-off-by: QuakeWang * style: simplify strict header parse error Signed-off-by: QuakeWang --------- Signed-off-by: QuakeWang --- src/common/datasource/src/file_format.rs | 1 + src/common/datasource/src/file_format/csv.rs | 39 +++- src/operator/src/error.rs | 17 ++ src/operator/src/statement/copy_table_from.rs | 178 +++++++++++++++++- tests-integration/src/tests/instance_test.rs | 128 +++++++++++++ 5 files changed, 361 insertions(+), 2 deletions(-) diff --git a/src/common/datasource/src/file_format.rs b/src/common/datasource/src/file_format.rs index 028aeef14d6..61f838899ef 100644 --- a/src/common/datasource/src/file_format.rs +++ b/src/common/datasource/src/file_format.rs @@ -63,6 +63,7 @@ pub const FORMAT_SCHEMA_INFER_MAX_RECORD: &str = "schema_infer_max_record"; pub const FORMAT_HEADERS: &str = "headers"; pub const FORMAT_HAS_HEADER: &str = "has_header"; pub const FORMAT_SKIP_BAD_RECORDS: &str = "skip_bad_records"; +pub const FORMAT_STRICT_HEADERS: &str = "strict_headers"; pub const FORMAT_TYPE: &str = "format"; pub const FILE_PATTERN: &str = "pattern"; pub const TIMESTAMP_FORMAT: &str = "timestamp_format"; diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index 24b61340f5b..d016d4eabe9 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -49,6 +49,7 @@ const SKIP_BAD_RECORDS_BATCH_SIZE: usize = 1; pub struct CsvFormat { pub has_header: bool, pub skip_bad_records: bool, + pub strict_headers: bool, pub delimiter: u8, pub schema_infer_max_record: Option, pub compression_type: CompressionType, @@ -111,6 +112,16 @@ impl TryFrom<&HashMap> for CsvFormat { format.skip_bad_records = parse_bool(file_format::FORMAT_SKIP_BAD_RECORDS, skip_bad_records)?; }; + if let Some(strict_headers) = value.get(file_format::FORMAT_STRICT_HEADERS) { + format.strict_headers = parse_bool(file_format::FORMAT_STRICT_HEADERS, strict_headers)?; + }; + if format.strict_headers && !format.has_header { + return error::ParseFormatSnafu { + key: file_format::FORMAT_STRICT_HEADERS, + value: "strict_headers=true requires headers=true", + } + .fail(); + } if let Some(timestamp_format) = value.get(file_format::TIMESTAMP_FORMAT) { format.timestamp_format = Some(timestamp_format.clone()); } @@ -135,6 +146,7 @@ impl Default for CsvFormat { Self { has_header: true, skip_bad_records: false, + strict_headers: false, delimiter: b',', schema_infer_max_record: Some(file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD), compression_type: CompressionType::Uncompressed, @@ -365,7 +377,8 @@ mod tests { use super::*; use crate::file_format::{ FORMAT_COMPRESSION_TYPE, FORMAT_DELIMITER, FORMAT_HAS_HEADER, FORMAT_HEADERS, - FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FileFormat, file_to_stream, + FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FORMAT_STRICT_HEADERS, FileFormat, + file_to_stream, }; use crate::test_util::{format_schema, test_store}; @@ -492,6 +505,7 @@ mod tests { delimiter: b'\t', has_header: false, skip_bad_records: false, + strict_headers: false, timestamp_format: None, time_format: None, date_format: None @@ -509,6 +523,17 @@ mod tests { } ); + let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "true".to_string())]); + let format = CsvFormat::try_from(&map).unwrap(); + + assert_eq!( + format, + CsvFormat { + strict_headers: true, + ..CsvFormat::default() + } + ); + let map = HashMap::from([(FORMAT_HEADERS.to_string(), "true".to_string())]); let format = CsvFormat::try_from(&map).unwrap(); assert_eq!(format, CsvFormat::default()); @@ -544,6 +569,9 @@ mod tests { let map = HashMap::from([(FORMAT_HEADERS.to_string(), "yes".to_string())]); assert!(CsvFormat::try_from(&map).is_err()); + + let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "yes".to_string())]); + assert!(CsvFormat::try_from(&map).is_err()); } #[test] @@ -555,6 +583,15 @@ mod tests { assert!(CsvFormat::try_from(&map).is_err()); } + #[test] + fn test_try_from_rejects_strict_headerless_csv() { + let map = HashMap::from([ + (FORMAT_HEADERS.to_string(), "false".to_string()), + (FORMAT_STRICT_HEADERS.to_string(), "true".to_string()), + ]); + assert!(CsvFormat::try_from(&map).is_err()); + } + #[tokio::test] async fn test_compressed_csv() { // Create test data diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 13b86140aa1..66fb66d7527 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -589,6 +589,22 @@ pub enum Error { location: Location, }, + #[snafu(display( + "CSV header mismatch in path: {}, unknown columns: {:?}, missing columns: {:?}, duplicate columns: {:?}", + path, + unknown_columns, + missing_columns, + duplicate_columns + ))] + CsvHeaderMismatch { + path: String, + unknown_columns: Vec, + missing_columns: Vec, + duplicate_columns: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to project schema"))] ProjectSchema { #[snafu(source)] @@ -936,6 +952,7 @@ impl ErrorExt for Error { | Error::ColumnNotFound { .. } | Error::BuildRegex { .. } | Error::InvalidSchema { .. } + | Error::CsvHeaderMismatch { .. } | Error::ProjectSchema { .. } | Error::UnsupportedFormat { .. } | Error::ColumnNoneDefaultValue { .. } diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index b64c654923c..014a2d32e9b 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::future::Future; use std::path::Path; use std::pin::Pin; @@ -392,6 +392,7 @@ impl StatementExecutor { .collect_metadata(&object_store, format.clone(), path.to_string()) .await?; + validate_csv_headers_if_required(&file_metadata, &table_schema)?; let schema_mapping = copy_from_schema_mapping(&file_metadata, &table_schema); let projected_file_schema = Arc::new( file_metadata @@ -642,6 +643,62 @@ fn ensure_schema_compatible(from: &SchemaRef, to: &SchemaRef) -> Result<()> { } } +fn validate_csv_headers_if_required(file_metadata: &FileMetadata, table: &SchemaRef) -> Result<()> { + let FileMetadata::Csv { + schema, + format, + path, + } = file_metadata + else { + return Ok(()); + }; + + if !format.strict_headers { + return Ok(()); + } + + let mut seen_file_columns = HashSet::with_capacity(schema.fields().len()); + let duplicate_columns = schema + .fields() + .iter() + .filter_map(|field| { + if seen_file_columns.insert(field.name().clone()) { + None + } else { + Some(field.name().clone()) + } + }) + .collect::>() + .into_iter() + .collect::>(); + let file_columns = seen_file_columns.into_iter().collect::>(); + let table_columns = table + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let unknown_columns = file_columns + .difference(&table_columns) + .cloned() + .collect::>(); + let missing_columns = table_columns + .difference(&file_columns) + .cloned() + .collect::>(); + + ensure!( + unknown_columns.is_empty() && missing_columns.is_empty() && duplicate_columns.is_empty(), + error::CsvHeaderMismatchSnafu { + path, + unknown_columns, + missing_columns, + duplicate_columns, + } + ); + + Ok(()) +} + /// Generates a maybe compatible schema of the file schema. /// /// If there is a field is found in table schema, @@ -998,12 +1055,131 @@ mod tests { } } + fn make_strict_csv_metadata(schema: Arc) -> FileMetadata { + FileMetadata::Csv { + schema, + format: CsvFormat { + strict_headers: true, + ..CsvFormat::default() + }, + path: "test.csv".to_string(), + } + } + fn assert_field(schema: &Schema, idx: usize, name: &str, data_type: &DataType) { let field = schema.field(idx); assert_eq!(field.name(), name); assert_eq!(field.data_type(), data_type); } + #[test] + fn test_strict_csv_headers_allows_reordered_columns() { + let file_schema = make_test_schema(&[ + Field::new("ts", DataType::Utf8, true), + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap(); + } + + #[test] + fn test_strict_csv_headers_rejects_unknown_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + Field::new("extra", DataType::Utf8, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns == vec!["extra".to_string()] + && missing_columns.is_empty() + && duplicate_columns.is_empty() + )); + } + + #[test] + fn test_strict_csv_headers_rejects_missing_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("ts", DataType::Utf8, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns.is_empty() + && missing_columns == vec!["reading_value".to_string()] + && duplicate_columns.is_empty() + )); + } + + #[test] + fn test_strict_csv_headers_rejects_duplicate_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + Field::new("host_id", DataType::UInt16, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns.is_empty() + && missing_columns.is_empty() + && duplicate_columns == vec!["host_id".to_string()] + )); + } + #[test] fn test_headerless_csv_schema_projection_is_positional() { let file_schema = make_test_schema(&[ diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index fa9bad0fac0..628f97de96a 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -1090,6 +1090,134 @@ async fn test_execute_copy_from_headerless_csv(instance: Arc) check_output_stream(output, expect).await; } +#[apply(both_instances_cases)] +async fn test_execute_copy_from_csv_strict_headers(instance: Arc) { + let instance = instance.frontend(); + let tmp_dir = temp_dir::create_temp_dir("test_execute_copy_from_csv_strict_headers"); + let matching_path = tmp_dir.path().join("matching.csv"); + let unknown_path = tmp_dir.path().join("unknown.csv"); + let missing_path = tmp_dir.path().join("missing.csv"); + let duplicate_path = tmp_dir.path().join("duplicate.csv"); + std::fs::write( + &matching_path, + "ts,host_id,reading_value\n2024-01-01T00:00:00,1,10.5\n", + ) + .unwrap(); + std::fs::write( + &unknown_path, + "host_id,reading_value,ts,extra\n1,10.5,2024-01-01T00:00:00,ignored\n", + ) + .unwrap(); + std::fs::write(&missing_path, "host_id,ts\n1,2024-01-01T00:00:00\n").unwrap(); + std::fs::write( + &duplicate_path, + "host_id,reading_value,ts,host_id\n1,10.5,2024-01-01T00:00:00,2\n", + ) + .unwrap(); + + let output = execute_sql( + &instance, + "CREATE TABLE csv_strict_headers(host_id INT, reading_value DOUBLE, ts TIMESTAMP TIME INDEX);", + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(0))); + + let output = execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + matching_path.display() + ), + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(1))); + + let output = execute_sql(&instance, "SELECT * FROM csv_strict_headers;") + .await + .data; + let expect = "\ ++---------+---------------+---------------------+ +| host_id | reading_value | ts | ++---------+---------------+---------------------+ +| 1 | 10.5 | 2024-01-01T00:00:00 | ++---------+---------------+---------------------+"; + check_output_stream(output, expect).await; + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + unknown_path.display() + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns == &vec!["extra".to_string()] + && missing_columns.is_empty() + && duplicate_columns.is_empty() + )); + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + missing_path.display() + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns.is_empty() + && missing_columns == &vec!["reading_value".to_string()] + && duplicate_columns.is_empty() + )); + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + duplicate_path.display() + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns.is_empty() + && missing_columns.is_empty() + && duplicate_columns == &vec!["host_id".to_string()] + )); +} + #[apply(both_instances_cases)] async fn test_execute_query_external_table_json(instance: Arc) { unsafe {