rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-19 02:12:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add strict CSV header validation (#8426)

Browse Source 
* feat: add strict CSV header validation

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>

* style: simplify strict header parse error

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>

---------

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
This commit is contained in:
QuakeWang
2026-07-10 02:15:08 +00:00
committed by GitHub
parent e6472fd12a
commit 4e14ac8676
5 changed files with 361 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/datasource/src/file_format.rs
+1
View File
@@ -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";
src/common/datasource/src/file_format/csv.rs
+38 -1
View File
@@ -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<usize>,
pub compression_type: CompressionType,
@@ -111,6 +112,16 @@ impl TryFrom<&HashMap<String, String>> 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
src/operator/src/error.rs
+17
View File
@@ -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<String>,
missing_columns: Vec<String>,
duplicate_columns: Vec<String>,
#[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 { .. }
src/operator/src/statement/copy_table_from.rs
+177 -1
View File
@@ -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::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let file_columns = seen_file_columns.into_iter().collect::<BTreeSet<_>>();
let table_columns = table
.fields()
.iter()
.map(|field| field.name().clone())
.collect::<BTreeSet<_>>();
let unknown_columns = file_columns
.difference(&table_columns)
.cloned()
.collect::<Vec<_>>();
let missing_columns = table_columns
.difference(&file_columns)
.cloned()
.collect::<Vec<_>>();
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<Schema>) -> 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(&[
tests-integration/src/tests/instance_test.rs
+128
View File
@@ -1090,6 +1090,134 @@ async fn test_execute_copy_from_headerless_csv(instance: Arc<dyn MockInstance>)
check_output_stream(output, expect).await;
}
#[apply(both_instances_cases)]
async fn test_execute_copy_from_csv_strict_headers(instance: Arc<dyn MockInstance>) {
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<dyn MockInstance>) {
unsafe {