From e9728192e2e6af647aee9860b115be1be5052ffa Mon Sep 17 00:00:00 2001 From: GatewayJ <835269233@qq.com> Date: Tue, 11 Aug 2026 23:59:38 +0800 Subject: [PATCH] fix(select): enforce typed S3 Select error semantics (#5942) * fix(select): enforce typed S3 Select error semantics * fix(select): classify function argument planner errors --------- Co-authored-by: overtrue --- crates/s3select-api/src/lib.rs | 323 +++++++- crates/s3select-api/src/object_store.rs | 125 +++- crates/s3select-api/src/query/session.rs | 20 +- .../s3select-query/src/dispatcher/manager.rs | 93 ++- .../src/dispatcher/parquet_table.rs | 101 ++- crates/s3select-query/src/sql/planner.rs | 38 +- .../src/test/error_handling_test.rs | 65 +- rustfs/src/app/select_object.rs | 693 ++++++++++++------ 8 files changed, 1159 insertions(+), 299 deletions(-) diff --git a/crates/s3select-api/src/lib.rs b/crates/s3select-api/src/lib.rs index 5c1436b7f..1b5cd1539 100644 --- a/crates/s3select-api/src/lib.rs +++ b/crates/s3select-api/src/lib.rs @@ -14,7 +14,12 @@ #![recursion_limit = "256"] -use datafusion::{common::DataFusionError, sql::sqlparser::parser::ParserError}; +use datafusion::{ + arrow::error::ArrowError, + common::{DataFusionError, SchemaError}, + parquet::errors::ParquetError, + sql::sqlparser::parser::ParserError, +}; use std::{error::Error as StdError, fmt::Display}; use thiserror::Error; @@ -67,23 +72,88 @@ pub enum QueryError { StoreError { e: String }, } -#[derive(Debug, Error)] -#[non_exhaustive] -pub enum S3SelectPolicyError { +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum SelectError { + #[error("The file is not in a supported compression format. Only GZIP and BZIP2 are supported.")] + InvalidCompressionFormat, + + #[error("The data source type is not valid. Only CSV, JSON, and Parquet are supported.")] + InvalidDataSource, + + #[error( + "Object decompression failed. Check that the object is properly compressed using the format specified in the request." + )] + TruncatedInput, + + #[error("An error occurred while parsing the CSV file. Check the file and try again.")] + CsvParsingError, + + #[error("An error occurred while parsing the JSON file. Check the file and try again.")] + JsonParsingError, + + #[error("An error occurred while parsing the Parquet file. Check the file and try again.")] + ParquetParsingError, + + #[error("{message}")] + ParseSelectFailure { message: String }, + + #[error("The SQL expression is invalid.")] + InvalidQuery, + + #[error("The SQL expression contains a data type that is not valid.")] + InvalidDataType, + + #[error("An incorrect argument type was specified in a function call in the SQL expression.")] + IncorrectSqlFunctionArgumentType, + + #[error("The data source path in the SQL expression is not supported.")] + DataSourcePathUnsupported, + #[error("Unsupported S3 Select SQL structure: {message}")] UnsupportedSqlStructure { message: String }, + #[error("We encountered an unsupported SQL operation.")] + UnsupportedSqlOperation, + + #[error("A column name or a path provided does not exist in the SQL expression.")] + EvaluatorBindingDoesNotExist, + + #[error("The field name matches to multiple fields in the file. Check the SQL expression and the file, and try again.")] + AmbiguousFieldName, + + #[error("The value of a parameter in ScanRange element is invalid. Check the service API documentation and try again.")] + InvalidScanRange, + #[error("S3 Select query concurrency limit reached")] QueryConcurrencyLimit, #[error("S3 Select query exceeded the {seconds}-second execution limit")] QueryTimeout { seconds: u64 }, + + #[error("S3 Select query resource limit exceeded")] + ResourceExhausted, + + #[error("The specified bucket does not exist.")] + BucketNotFound, + + #[error("The specified key does not exist.")] + ObjectNotFound, + + #[error("The query was canceled")] + Canceled, + + #[error("An internal error occurred.")] + InternalError, } +pub type S3SelectPolicyError = SelectError; + +const MAX_ERROR_SOURCE_DEPTH: usize = 16; + impl QueryError { fn source_error(&self) -> Option<&T> { let mut err: &(dyn StdError + 'static) = self; - for _ in 0..16 { + for _ in 0..MAX_ERROR_SOURCE_DEPTH { if let Some(source) = err.downcast_ref::() { return Some(source); } @@ -99,10 +169,113 @@ impl QueryError { pub fn s3_select_policy_error(&self) -> Option<&S3SelectPolicyError> { self.source_error() } + + pub fn select_error(&self) -> SelectError { + let mut err: &(dyn StdError + 'static) = match self { + Self::Datafusion { source } => source.as_ref(), + _ => self, + }; + for _ in 0..MAX_ERROR_SOURCE_DEPTH { + if let Some(select_error) = classify_select_error_source(err) { + return select_error; + } + let Some(source) = err.source() else { + break; + }; + err = source; + } + + match self { + QueryError::NotImplemented { .. } => SelectError::UnsupportedSqlOperation, + QueryError::MultiStatement { .. } => SelectError::UnsupportedSqlStructure { + message: "multiple SQL statements are not supported".to_string(), + }, + QueryError::BuildQueryDispatcher { .. } | QueryError::FunctionExists { .. } | QueryError::StoreError { .. } => { + SelectError::InternalError + } + QueryError::Cancel => SelectError::Canceled, + QueryError::FunctionNotExists { .. } => SelectError::InvalidQuery, + QueryError::Datafusion { .. } | QueryError::Parser { .. } => SelectError::InternalError, + } + } } -impl From for QueryError { - fn from(value: S3SelectPolicyError) -> Self { +fn classify_select_error_source(err: &(dyn StdError + 'static)) -> Option { + if let Some(error) = err.downcast_ref::() { + return Some(error.clone()); + } + if let Some(error) = err.downcast_ref::() { + return Some(error.select_error()); + } + if let Some(error) = err.downcast_ref::() { + return match error { + datafusion::object_store::Error::NotFound { source, .. } => Some( + source + .downcast_ref::() + .map_or(SelectError::ObjectNotFound, object_store::SelectObjectStoreError::select_error), + ), + _ => None, + }; + } + if let Some(error) = err.downcast_ref::() { + return Some(SelectError::ParseSelectFailure { + message: error.to_string(), + }); + } + if let Some(error) = err.downcast_ref::() { + return match error { + ArrowError::CsvError(_) => Some(SelectError::CsvParsingError), + ArrowError::JsonError(_) => Some(SelectError::JsonParsingError), + ArrowError::ParquetError(_) => Some(SelectError::ParquetParsingError), + ArrowError::CastError(_) | ArrowError::ParseError(_) => Some(SelectError::InvalidDataType), + ArrowError::MemoryError(_) => Some(SelectError::ResourceExhausted), + ArrowError::ExternalError(_) | ArrowError::IoError(_, _) => None, + _ => Some(SelectError::InternalError), + }; + } + if let Some(error) = err.downcast_ref::() { + return match error { + ParquetError::External(_) => None, + _ => Some(SelectError::ParquetParsingError), + }; + } + if let Some(error) = err.downcast_ref::() { + return Some(match error { + SchemaError::FieldNotFound { .. } => SelectError::EvaluatorBindingDoesNotExist, + SchemaError::AmbiguousReference { .. } + | SchemaError::DuplicateQualifiedField { .. } + | SchemaError::DuplicateUnqualifiedField { .. } => SelectError::AmbiguousFieldName, + }); + } + if let Some(error) = err.downcast_ref::() { + return match error { + DataFusionError::NotImplemented(_) => Some(SelectError::UnsupportedSqlOperation), + DataFusionError::Plan(_) => Some(SelectError::InvalidQuery), + DataFusionError::ResourcesExhausted(_) => Some(SelectError::ResourceExhausted), + DataFusionError::Internal(_) + | DataFusionError::Execution(_) + | DataFusionError::Configuration(_) + | DataFusionError::Substrait(_) + | DataFusionError::Ffi(_) => Some(SelectError::InternalError), + DataFusionError::ArrowError(_, _) + | DataFusionError::ParquetError(_) + | DataFusionError::ObjectStore(_) + | DataFusionError::IoError(_) + | DataFusionError::SQL(_, _) + | DataFusionError::SchemaError(_, _) + | DataFusionError::ExecutionJoin(_) + | DataFusionError::External(_) + | DataFusionError::Context(_, _) + | DataFusionError::Diagnostic(_, _) + | DataFusionError::Collection(_) + | DataFusionError::Shared(_) => None, + }; + } + None +} + +impl From for QueryError { + fn from(value: SelectError) -> Self { Self::Datafusion { source: Box::new(DataFusionError::External(Box::new(value))), } @@ -161,7 +334,7 @@ mod tests { }; assert_eq!(err.to_string(), "Multi-statement not allow, found num:2, sql:SELECT 1; SELECT 2;"); - let err = S3SelectPolicyError::UnsupportedSqlStructure { + let err = SelectError::UnsupportedSqlStructure { message: "JOIN is not supported".to_string(), }; assert_eq!(err.to_string(), "Unsupported S3 Select SQL structure: JOIN is not supported"); @@ -170,11 +343,11 @@ mod tests { assert_eq!(err.to_string(), "The query has been canceled"); assert_eq!( - S3SelectPolicyError::QueryConcurrencyLimit.to_string(), + SelectError::QueryConcurrencyLimit.to_string(), "S3 Select query concurrency limit reached" ); assert_eq!( - S3SelectPolicyError::QueryTimeout { seconds: 300 }.to_string(), + SelectError::QueryTimeout { seconds: 300 }.to_string(), "S3 Select query exceeded the 300-second execution limit" ); @@ -223,12 +396,132 @@ mod tests { #[test] fn policy_error_is_recoverable_from_query_error() { - let err: QueryError = S3SelectPolicyError::QueryTimeout { seconds: 300 }.into(); + let err: QueryError = SelectError::QueryTimeout { seconds: 300 }.into(); - assert!(matches!( - err.s3_select_policy_error(), - Some(S3SelectPolicyError::QueryTimeout { seconds: 300 }) - )); + assert!(matches!(err.s3_select_policy_error(), Some(SelectError::QueryTimeout { seconds: 300 }))); + } + + #[test] + fn query_error_classifies_data_errors_without_display_matching() { + let cases = [ + ( + DataFusionError::ArrowError(Box::new(ArrowError::CsvError("private csv detail".to_string())), None), + SelectError::CsvParsingError, + ), + ( + DataFusionError::ArrowError(Box::new(ArrowError::JsonError("private json detail".to_string())), None), + SelectError::JsonParsingError, + ), + ( + DataFusionError::ParquetError(Box::new(ParquetError::General("private parquet detail".to_string()))), + SelectError::ParquetParsingError, + ), + ( + DataFusionError::External(Box::new(SelectError::TruncatedInput)), + SelectError::TruncatedInput, + ), + ( + DataFusionError::ArrowError( + Box::new(ArrowError::InvalidArgumentError("private implementation detail".to_string())), + None, + ), + SelectError::InternalError, + ), + ( + DataFusionError::ArrowError(Box::new(ArrowError::CastError("invalid cast".to_string())), None), + SelectError::InvalidDataType, + ), + ( + DataFusionError::ArrowError(Box::new(ArrowError::MemoryError("query memory limit".to_string())), None), + SelectError::ResourceExhausted, + ), + ( + DataFusionError::Execution("private execution detail".to_string()), + SelectError::InternalError, + ), + (DataFusionError::Plan("invalid expression".to_string()), SelectError::InvalidQuery), + ( + DataFusionError::NotImplemented("unsupported expression".to_string()), + SelectError::UnsupportedSqlOperation, + ), + ( + DataFusionError::SchemaError( + Box::new(SchemaError::FieldNotFound { + field: Box::new(datafusion::common::Column::from_name("missing")), + valid_fields: Vec::new(), + }), + Box::new(None), + ), + SelectError::EvaluatorBindingDoesNotExist, + ), + ( + DataFusionError::SchemaError( + Box::new(SchemaError::AmbiguousReference { + field: Box::new(datafusion::common::Column::from_name("duplicate")), + }), + Box::new(None), + ), + SelectError::AmbiguousFieldName, + ), + ]; + + for (source, expected) in cases { + let error = QueryError::from(source); + assert_eq!(error.select_error(), expected, "wrong classification for {error:?}"); + } + } + + #[test] + fn query_error_preserves_typed_object_store_classification() { + let bucket_error = QueryError::from(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::NotFound { + path: "private-bucket/private-object".to_string(), + source: Box::new(object_store::SelectObjectStoreError::BucketNotFound { + source: SelectStorageError::BucketNotFound("private-bucket".to_string()), + }), + }))); + let object_error = QueryError::from(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::NotFound { + path: "private-bucket/private-object".to_string(), + source: Box::new(object_store::SelectObjectStoreError::ObjectNotFound { + source: SelectStorageError::ObjectNotFound("private-bucket".to_string(), "private-object".to_string()), + }), + }))); + let scan_range_error = + QueryError::from(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::Generic { + store: "test", + source: Box::new(object_store::SelectObjectStoreError::InvalidScanRange), + }))); + let storage_error = QueryError::from(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::Generic { + store: "test", + source: Box::new(object_store::SelectObjectStoreError::Storage { + source: SelectStorageError::LessData, + }), + }))); + + assert_eq!(bucket_error.select_error(), SelectError::BucketNotFound); + assert_eq!(object_error.select_error(), SelectError::ObjectNotFound); + assert_eq!(scan_range_error.select_error(), SelectError::InvalidScanRange); + assert_eq!(storage_error.select_error(), SelectError::InternalError); + } + + #[test] + fn select_error_source_traversal_stops_at_the_depth_bound() { + #[derive(Debug)] + struct CyclicError; + + impl std::fmt::Display for CyclicError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("cyclic error") + } + } + + impl StdError for CyclicError { + fn source(&self) -> Option<&(dyn StdError + 'static)> { + Some(self) + } + } + + let error = QueryError::from(DataFusionError::External(Box::new(CyclicError))); + assert_eq!(error.select_error(), SelectError::InternalError); } #[test] diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 5a9080390..476cbab5f 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -13,7 +13,7 @@ // limitations under the License. use crate::{ - PrepareSelectObjectSnapshotError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectOptions, + PrepareSelectObjectSnapshotError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectError, SelectGetObjectReader, SelectObjectOptions, SelectObjectSnapshot, SelectObjectSnapshotReadError, SelectStorageError, SelectStore, SnapshotConsistencyError, query::{ parser::RustFsDialect, @@ -115,6 +115,38 @@ pub(crate) enum EcObjectStoreBuildError { Snapshot(#[source] SnapshotConsistencyError), } +#[derive(Debug, thiserror::Error)] +pub(crate) enum SelectObjectStoreError { + #[error("SelectObjectContent bucket does not exist")] + BucketNotFound { + #[source] + source: SelectStorageError, + }, + #[error("SelectObjectContent object does not exist")] + ObjectNotFound { + #[source] + source: SelectStorageError, + }, + #[error("SelectObjectContent storage failure")] + Storage { + #[source] + source: SelectStorageError, + }, + #[error("SelectObjectContent ScanRange is invalid")] + InvalidScanRange, +} + +impl SelectObjectStoreError { + pub(crate) fn select_error(&self) -> SelectError { + match self { + Self::BucketNotFound { .. } => SelectError::BucketNotFound, + Self::ObjectNotFound { .. } => SelectError::ObjectNotFound, + Self::InvalidScanRange => SelectError::InvalidScanRange, + Self::Storage { .. } => SelectError::InternalError, + } + } +} + #[derive(Clone, Copy, Debug)] pub struct SelectScanRange { start: u64, @@ -502,8 +534,7 @@ fn map_prepare_snapshot_error(bucket: &str, object: &str, err: PrepareSelectObje } fn map_build_error_to_s3(error: EcObjectStoreBuildError) -> S3Error { - let message = error.to_string(); - let mut s3_error = S3Error::with_message(S3ErrorCode::InternalError, message); + let mut s3_error = S3Error::with_message(S3ErrorCode::InternalError, SelectError::InternalError.to_string()); s3_error.set_source(Box::new(error)); s3_error } @@ -519,15 +550,21 @@ fn snapshot_read_error(bucket: &str, object: &str, err: SelectObjectSnapshotRead } fn map_storage_error(bucket: &str, object: &str, err: SelectStorageError) -> o_Error { - if select_is_err_bucket_not_found(&err) || select_is_err_object_not_found(&err) || select_is_err_version_not_found(&err) { + if select_is_err_bucket_not_found(&err) { return o_Error::NotFound { path: format!("{bucket}/{object}"), - source: Box::new(err), + source: Box::new(SelectObjectStoreError::BucketNotFound { source: err }), + }; + } + if select_is_err_object_not_found(&err) || select_is_err_version_not_found(&err) { + return o_Error::NotFound { + path: format!("{bucket}/{object}"), + source: Box::new(SelectObjectStoreError::ObjectNotFound { source: err }), }; } o_Error::Generic { store: "EcObjectStore", - source: Box::new(err), + source: Box::new(SelectObjectStoreError::Storage { source: err }), } } @@ -602,7 +639,7 @@ fn parse_scan_range_from_bounds( fn invalid_scan_range_store_error() -> o_Error { o_Error::Generic { store: "EcObjectStore", - source: format!("ScanRange: {INVALID_SCAN_RANGE_MESSAGE}").into(), + source: Box::new(SelectObjectStoreError::InvalidScanRange), } } @@ -1150,7 +1187,11 @@ where })? .map_err(|e| o_Error::Generic { store: "EcObjectStore", - source: Box::new(e), + source: if e.kind() == std::io::ErrorKind::InvalidData { + Box::new(SelectError::JsonParsingError) + } else { + Box::new(e) + }, })?; // ── 3. Yield phase (one Bytes per NDJSON line) ─────────────────── @@ -1341,12 +1382,13 @@ mod test { SELECT_DEFAULT_READ_BUFFER_SIZE, SelectObjectOptions, SelectObjectSnapshot, SelectScanRange, SnapshotConsistencyError, bytes_stream, convert_csv_delimiter_stream, convert_field_delimiter_stream, convert_record_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, - json_document_ndjson_stream, json_document_ndjson_stream_with_parser, scan_range_from_bounds, scan_range_stream, - select_read_headers, snapshot_last_modified, validate_json_document_size, + json_document_ndjson_stream, json_document_ndjson_stream_with_parser, map_storage_error, scan_range_from_bounds, + scan_range_stream, select_read_headers, snapshot_last_modified, validate_json_document_size, }; use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker}; use crate::storage_api::SelectPutObjReader; use crate::storage_api::object_store::ObjectIO as _; + use crate::{QueryError, SelectError, SelectStorageError}; use bytes::Bytes; use datafusion::{ common::DataFusionError, @@ -2688,6 +2730,69 @@ mod test { assert!(output.next().await.is_none()); } + #[tokio::test] + async fn malformed_json_document_stream_has_typed_select_error() { + let input = b"{bad".to_vec(); + let memory_pool: Arc = + Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let mut output = json_document_ndjson_stream( + Box::new(std::io::Cursor::new(input.clone())), + input.len() as u64, + None, + memory_pool, + None, + ); + + let source = output + .next() + .await + .expect("malformed JSON should produce one stream error") + .expect_err("malformed JSON DOCUMENT must fail"); + let error = QueryError::from(DataFusionError::ObjectStore(Box::new(source))); + + assert_eq!(error.select_error(), SelectError::JsonParsingError); + assert!(output.next().await.is_none()); + } + + #[test] + fn storage_error_mapper_preserves_protocol_classification() { + let classify = |source| QueryError::from(DataFusionError::ObjectStore(Box::new(source))).select_error(); + + assert_eq!( + classify(map_storage_error( + "private-bucket", + "private-object", + SelectStorageError::BucketNotFound("private-bucket".to_string()), + )), + SelectError::BucketNotFound + ); + assert_eq!( + classify(map_storage_error( + "private-bucket", + "private-object", + SelectStorageError::ObjectNotFound("private-bucket".to_string(), "private-object".to_string()), + )), + SelectError::ObjectNotFound + ); + assert_eq!( + classify(map_storage_error("private-bucket", "private-object", SelectStorageError::LessData)), + SelectError::InternalError + ); + assert_eq!( + classify(scan_range_from_bounds(Some(10), None, 10).expect_err("out-of-bounds range must fail")), + SelectError::InvalidScanRange + ); + let parquet_source = map_storage_error( + "private-bucket", + "private-object", + SelectStorageError::ObjectNotFound("private-bucket".to_string(), "private-object".to_string()), + ); + let parquet_error = QueryError::from(DataFusionError::ParquetError(Box::new( + datafusion::parquet::errors::ParquetError::External(Box::new(parquet_source)), + ))); + assert_eq!(parquet_error.select_error(), SelectError::ObjectNotFound); + } + #[test] fn test_json_document_size_error_is_resource_exhausted() { assert!(validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES).is_ok()); diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 39a70e383..40cdf4696 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -433,7 +433,7 @@ impl SessionCtxFactory { let path = Path::from(context.input.key.clone()); store.put(&path, data_bytes.into()).await.map_err(|e| { error!("put data into memory failed: {}", e.to_string()); - QueryError::StoreError { e: e.to_string() } + QueryError::from(DataFusionError::from(e)) })?; df_session_state.with_object_store(&store_url, store).build() @@ -477,16 +477,11 @@ fn test_parquet_bytes() -> QueryResult> { let mut bytes = Vec::new(); { - let mut writer = - ArrowWriter::try_new(&mut bytes, schema, None).map_err(|e| QueryError::StoreError { e: e.to_string() })?; - writer - .write(&first_batch) - .map_err(|e| QueryError::StoreError { e: e.to_string() })?; - writer.flush().map_err(|e| QueryError::StoreError { e: e.to_string() })?; - writer - .write(&second_batch) - .map_err(|e| QueryError::StoreError { e: e.to_string() })?; - writer.close().map_err(|e| QueryError::StoreError { e: e.to_string() })?; + let mut writer = ArrowWriter::try_new(&mut bytes, schema, None).map_err(DataFusionError::from)?; + writer.write(&first_batch).map_err(DataFusionError::from)?; + writer.flush().map_err(DataFusionError::from)?; + writer.write(&second_batch).map_err(DataFusionError::from)?; + writer.close().map_err(DataFusionError::from)?; } Ok(bytes) } @@ -509,7 +504,8 @@ fn test_parquet_batch( Arc::new(Int32Array::from(salaries.to_vec())), ], ) - .map_err(|e| QueryError::StoreError { e: e.to_string() }) + .map_err(DataFusionError::from) + .map_err(QueryError::from) } #[cfg(test)] diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index c9398db2b..993d77c01 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -38,7 +38,7 @@ use datafusion::{ use futures::Stream; use parking_lot::Mutex; use rustfs_s3select_api::{ - QueryError, QueryResult, S3SelectPolicyError, + QueryError, QueryResult, SelectError, query::{ Query, ast::ExtStatement, @@ -128,7 +128,7 @@ impl QueryDispatcher for SimpleQueryDispatcher { .query_admission .clone() .try_acquire_owned() - .map_err(|_| QueryError::from(S3SelectPolicyError::QueryConcurrencyLimit))?; + .map_err(|_| QueryError::from(SelectError::QueryConcurrencyLimit))?; Ok(QueryAdmission::new(Arc::new(permit))) } @@ -245,7 +245,7 @@ impl SimpleQueryDispatcher { .query_admission .clone() .try_acquire_owned() - .map_err(|_| QueryError::from(S3SelectPolicyError::QueryConcurrencyLimit))?; + .map_err(|_| QueryError::from(SelectError::QueryConcurrencyLimit))?; Arc::new(permit) } }; @@ -293,7 +293,7 @@ impl SimpleQueryDispatcher { ) -> QueryResult { let deadline = query_tracker.deadline(); let timeout_error = || { - S3SelectPolicyError::QueryTimeout { + SelectError::QueryTimeout { seconds: query_tracker.timeout_seconds(), } .into() @@ -343,11 +343,11 @@ impl SimpleQueryDispatcher { return QueryError::Cancel; } match query_tracker.status() { - QueryExecutionStatus::TimedOut => S3SelectPolicyError::QueryTimeout { + QueryExecutionStatus::TimedOut => SelectError::QueryTimeout { seconds: query_tracker.timeout_seconds(), } .into(), - QueryExecutionStatus::Active if Instant::now() >= query_tracker.deadline() => S3SelectPolicyError::QueryTimeout { + QueryExecutionStatus::Active if Instant::now() >= query_tracker.deadline() => SelectError::QueryTimeout { seconds: query_tracker.timeout_seconds(), } .into(), @@ -430,15 +430,11 @@ impl SimpleQueryDispatcher { } else if *info == *USE { file_format = file_format.with_has_header(true); } else { - return Err(QueryError::NotImplemented { - err: "unsupported FileHeaderInfo".to_string(), - }); + return Err(SelectError::InvalidDataSource.into()); } } _ => { - return Err(QueryError::NotImplemented { - err: "unsupported FileHeaderInfo".to_string(), - }); + return Err(SelectError::InvalidDataSource.into()); } } if let Some(quote) = csv.quote_character.as_ref() { @@ -462,9 +458,7 @@ impl SimpleQueryDispatcher { .unwrap_or_else(|| ".json".to_string()); (ListingOptions::new(Arc::new(file_format)).with_file_extension(file_ext), false, false) } else { - return Err(QueryError::NotImplemented { - err: "not support this file type".to_string(), - }); + return Err(SelectError::InvalidDataSource.into()); }; let resolve_schema = listing_options.infer_schema(session.inner(), &table_path).await?; @@ -642,7 +636,7 @@ impl Stream for TrackedRecordBatchStream { } fn query_timeout_error(timeout_seconds: u64) -> datafusion::common::DataFusionError { - datafusion::common::DataFusionError::External(Box::new(S3SelectPolicyError::QueryTimeout { + datafusion::common::DataFusionError::External(Box::new(SelectError::QueryTimeout { seconds: timeout_seconds, })) } @@ -791,7 +785,7 @@ mod tests { }; use futures::{StreamExt, TryStreamExt, stream}; use rustfs_s3select_api::{ - QueryError, QueryResult, S3SelectPolicyError, + QueryError, QueryResult, SelectError, query::{ Context as QueryContext, Query, dispatcher::QueryDispatcher, @@ -1338,6 +1332,47 @@ mod tests { assert_eq!(dispatcher.memory_limit_bytes, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES); } + #[tokio::test] + async fn invalid_csv_header_info_is_typed_invalid_data_source() { + let mut input = test_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("test input should use CSV") + .file_header_info = Some(FileHeaderInfo::from_static("INVALID")); + let input = Arc::new(input); + let optimizer = Arc::new(CascadeOptimizerBuilder::default().build()); + let scheduler = Arc::new(LocalScheduler {}); + let dispatcher = SimpleQueryDispatcherBuilder::default() + .with_input(Arc::clone(&input)) + .with_default_table_provider(Arc::new(BaseTableProvider::default())) + .with_session_factory(Arc::new(SessionCtxFactory::new(true))) + .with_parser(Arc::new(DefaultParser::default())) + .with_query_execution_factory(Arc::new(SqlQueryExecutionFactory::new(optimizer, scheduler))) + .with_func_manager(Arc::new(SimpleFunctionMetadataManager::default())) + .build() + .expect("query dispatcher should build"); + let query = Query::new( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + ); + let query_state_machine = dispatcher + .build_query_state_machine(query) + .await + .expect("query should acquire admission"); + + let error = match dispatcher.build_logical_plan(query_state_machine).await { + Err(error) => error, + Ok(_) => panic!("invalid FileHeaderInfo must fail while building the provider"), + }; + + assert_eq!(error.select_error(), SelectError::InvalidDataSource); + } + #[tokio::test] async fn csv_query_uses_custom_record_delimiter_across_file_partitions() { const ROW_COUNT: usize = 200_000; @@ -1420,7 +1455,7 @@ mod tests { assert!(matches!( result, - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::QueryConcurrencyLimit)) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::QueryConcurrencyLimit)) )); } @@ -1474,7 +1509,7 @@ mod tests { assert!(matches!( result, - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::QueryConcurrencyLimit)) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::QueryConcurrencyLimit)) )); } @@ -1571,7 +1606,7 @@ mod tests { assert!(matches!( result, - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::QueryTimeout { seconds: 0 })) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::QueryTimeout { seconds: 0 })) )); assert_eq!(admission.available_permits(), 1); } @@ -1601,7 +1636,7 @@ mod tests { assert!(matches!( dispatcher.execute_logical_plan(logical_plan, query_state_machine).await, - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::QueryTimeout { seconds: 300 })) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::QueryTimeout { seconds: 300 })) )); assert_eq!(admission.available_permits(), 1); } @@ -1742,7 +1777,7 @@ mod tests { assert_eq!(admission.available_permits(), 1); assert!(matches!( dispatcher.build_logical_plan(query_state_machine).await, - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::QueryTimeout { seconds: 1 })) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::QueryTimeout { seconds: 1 })) )); } @@ -1843,7 +1878,7 @@ mod tests { release_drop_tx.send(()).expect("release result drop"); assert!(matches!( task.await.expect("deadline task should finish"), - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::QueryTimeout { seconds: 1 })) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::QueryTimeout { seconds: 1 })) )); assert_eq!(admission.available_permits(), 1); } @@ -1985,8 +2020,8 @@ mod tests { panic!("expected external query error"); }; assert!(matches!( - source.downcast_ref::(), - Some(S3SelectPolicyError::QueryTimeout { seconds: 300 }) + source.downcast_ref::(), + Some(SelectError::QueryTimeout { seconds: 300 }) )); assert!(inner_dropped.load(Ordering::SeqCst)); assert_eq!(admission.available_permits(), 1); @@ -2124,8 +2159,8 @@ mod tests { panic!("expected external query error"); }; assert!(matches!( - source.downcast_ref::(), - Some(S3SelectPolicyError::QueryTimeout { seconds: 300 }) + source.downcast_ref::(), + Some(SelectError::QueryTimeout { seconds: 300 }) )); assert_eq!(admission.available_permits(), 1); assert!(output.next().await.is_none()); @@ -2174,8 +2209,8 @@ mod tests { panic!("expected external query error"); }; assert!(matches!( - source.downcast_ref::(), - Some(S3SelectPolicyError::QueryTimeout { seconds: 1 }) + source.downcast_ref::(), + Some(SelectError::QueryTimeout { seconds: 1 }) )); assert_eq!(admission.available_permits(), 1); assert!(output.next().await.is_none()); diff --git a/crates/s3select-query/src/dispatcher/parquet_table.rs b/crates/s3select-query/src/dispatcher/parquet_table.rs index 9ef1eac65..97949371c 100644 --- a/crates/s3select-query/src/dispatcher/parquet_table.rs +++ b/crates/s3select-query/src/dispatcher/parquet_table.rs @@ -38,7 +38,7 @@ use datafusion::{ }; use futures::{FutureExt, TryFutureExt, future::BoxFuture}; use rustfs_s3select_api::{ - QueryError, QueryResult, + QueryResult, object_store::{SelectScanRange, scan_range_from_bounds}, }; use s3s::dto::SelectObjectContentInput; @@ -106,7 +106,10 @@ impl ParquetSelectTable { let object_store_url = table_path.object_store(); let object_location = Path::from(input.key.clone()); let store = state.runtime_env().object_store(&object_store_url)?; - let object_meta = store.head(&object_location).await.map_err(query_store_error)?; + let object_meta = store + .head(&object_location) + .await + .map_err(datafusion::common::DataFusionError::from)?; let reader = ObjectStoreParquetReader { store: Arc::clone(&store), @@ -115,7 +118,7 @@ impl ParquetSelectTable { }; let builder = ParquetRecordBatchStreamBuilder::new(reader) .await - .map_err(query_store_error)?; + .map_err(datafusion::common::DataFusionError::from)?; let schema = Arc::clone(builder.schema()); let metadata = Arc::clone(builder.metadata()); let access_plan = parquet_access_plan(input, object_meta.size, metadata.as_ref())?; @@ -180,7 +183,8 @@ fn parquet_access_plan( let Some(scan_range) = input.request.scan_range.as_ref() else { return Ok(None); }; - let scan_range = scan_range_from_bounds(scan_range.start, scan_range.end, object_size).map_err(query_store_error)?; + let scan_range = scan_range_from_bounds(scan_range.start, scan_range.end, object_size) + .map_err(datafusion::common::DataFusionError::from)?; Ok(scan_range.map(|range| Arc::new(access_plan_for_scan_range(range, metadata)))) } @@ -214,10 +218,6 @@ fn parquet_store_error(err: ObjectStoreError) -> ParquetError { ParquetError::External(Box::new(err)) } -fn query_store_error(err: impl fmt::Display) -> QueryError { - QueryError::StoreError { e: err.to_string() } -} - #[cfg(test)] mod tests { use super::*; @@ -227,7 +227,13 @@ mod tests { datatypes::{DataType, Field, Schema, SchemaRef}, record_batch::RecordBatch, }, + object_store::memory::InMemory, parquet::arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder}, + prelude::SessionContext, + }; + use rustfs_s3select_api::SelectError; + use s3s::dto::{ + CSVOutput, ExpressionType, InputSerialization, OutputSerialization, ParquetInput, ScanRange, SelectObjectContentRequest, }; use std::{ fs::File, @@ -275,6 +281,85 @@ mod tests { assert!(!plan.should_scan(1)); } + #[test] + fn parquet_access_plan_has_typed_invalid_scan_range_error() { + let metadata = two_row_group_metadata(); + let mut input = parquet_input("test.parquet"); + input.request.scan_range = Some(ScanRange { + start: Some(10), + end: None, + }); + + let error = parquet_access_plan(&input, 10, metadata.as_ref()).expect_err("out-of-bounds range must fail"); + + assert_eq!(error.select_error(), SelectError::InvalidScanRange); + } + + #[tokio::test] + async fn try_new_preserves_missing_object_error() { + let store = Arc::new(InMemory::new()); + let context = parquet_session(store); + let state = context.state(); + + let error = match ParquetSelectTable::try_new(&state, &parquet_input("missing.parquet")).await { + Ok(_) => panic!("missing parquet object must fail"), + Err(error) => error, + }; + + assert_eq!(error.select_error(), SelectError::ObjectNotFound); + } + + #[tokio::test] + async fn try_new_preserves_parquet_metadata_error() { + let store = Arc::new(InMemory::new()); + let object = Path::from("corrupt.parquet"); + store + .put(&object, Bytes::from_static(b"not a parquet file").into()) + .await + .expect("put corrupt parquet object"); + let context = parquet_session(store); + let state = context.state(); + + let error = match ParquetSelectTable::try_new(&state, &parquet_input(object.as_ref())).await { + Ok(_) => panic!("corrupt parquet metadata must fail"), + Err(error) => error, + }; + + assert_eq!(error.select_error(), SelectError::ParquetParsingError); + } + + fn parquet_session(store: Arc) -> SessionContext { + let context = SessionContext::new(); + let store_url = ObjectStoreUrl::parse("s3://test-bucket").expect("valid test object store URL"); + context.register_object_store(store_url.as_ref(), store); + context + } + + fn parquet_input(key: &str) -> SelectObjectContentInput { + SelectObjectContentInput { + bucket: "test-bucket".to_string(), + expected_bucket_owner: None, + key: key.to_string(), + sse_customer_algorithm: None, + sse_customer_key: None, + sse_customer_key_md5: None, + request: SelectObjectContentRequest { + expression: "SELECT * FROM S3Object".to_string(), + expression_type: ExpressionType::from_static(ExpressionType::SQL), + input_serialization: InputSerialization { + parquet: Some(ParquetInput::default()), + ..Default::default() + }, + output_serialization: OutputSerialization { + csv: Some(CSVOutput::default()), + ..Default::default() + }, + request_progress: None, + scan_range: None, + }, + } + } + fn two_row_group_metadata() -> Arc { let now = SystemTime::now() .duration_since(UNIX_EPOCH) diff --git a/crates/s3select-query/src/sql/planner.rs b/crates/s3select-query/src/sql/planner.rs index c262166f0..a5bbb2881 100644 --- a/crates/s3select-query/src/sql/planner.rs +++ b/crates/s3select-query/src/sql/planner.rs @@ -23,7 +23,7 @@ use datafusion::sql::{ }, }; use rustfs_s3select_api::{ - QueryError, QueryResult, S3SelectPolicyError, + QueryError, QueryResult, SelectError, query::{ ast::ExtStatement, logical_planner::{LogicalPlanner, Plan, QueryPlan}, @@ -68,7 +68,7 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> { match stmt { Statement::Query(_) => { validate_s3_select_statement(&stmt)?; - let df_plan = self.df_planner.sql_statement_to_plan(stmt)?; + let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?; let plan = Plan::Query(QueryPlan { df_plan, is_tag_scan: false, @@ -76,11 +76,25 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> { Ok(plan) } - _ => Err(QueryError::NotImplemented { err: stmt.to_string() }), + _ => Err(unsupported_structure("only SELECT queries are supported")), } } } +fn classify_planner_error(error: datafusion::common::DataFusionError) -> QueryError { + if matches!( + &error, + datafusion::common::DataFusionError::Plan(message) + if message.starts_with("Failed to coerce arguments to satisfy a call to") + || (message.starts_with("Internal error: Function '") + && message.contains("' failed to match any signature, errors:")) + ) { + return SelectError::IncorrectSqlFunctionArgumentType.into(); + } + + error.into() +} + fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> { let Statement::Query(query) = statement else { return Err(unsupported_structure("only SELECT queries are supported")); @@ -191,7 +205,7 @@ fn validate_select(select: &Select) -> QueryResult<()> { let ([ObjectNamePart::Identifier(table_name)] | [ObjectNamePart::Identifier(table_name), ObjectNamePart::Identifier(_)]) = name.0.as_slice() else { - return Err(unsupported_structure("the source must be S3Object")); + return Err(SelectError::DataSourcePathUnsupported.into()); }; let is_s3_object = if table_name.quote_style.is_some() { table_name.value == "S3Object" @@ -199,14 +213,14 @@ fn validate_select(select: &Select) -> QueryResult<()> { table_name.value.eq_ignore_ascii_case("S3Object") }; if !is_s3_object { - return Err(unsupported_structure("the source must be S3Object")); + return Err(SelectError::DataSourcePathUnsupported.into()); } Ok(()) } fn unsupported_structure(message: &str) -> QueryError { - S3SelectPolicyError::UnsupportedSqlStructure { + SelectError::UnsupportedSqlStructure { message: message.to_string(), } .into() @@ -234,7 +248,7 @@ mod tests { use super::validate_s3_select_statement; use crate::sql::parser::ExtParser; use datafusion::sql::sqlparser::ast::Statement; - use rustfs_s3select_api::{S3SelectPolicyError, query::ast::ExtStatement}; + use rustfs_s3select_api::{SelectError, query::ast::ExtStatement}; fn parse_statement(sql: &str) -> Statement { let mut statements = ExtParser::parse_sql(sql).expect("SQL should parse"); @@ -271,7 +285,7 @@ mod tests { validate_s3_select_statement(&statement), Err(ref err) if matches!( err.s3_select_policy_error(), - Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "JOIN is not supported" + Some(SelectError::UnsupportedSqlStructure { message }) if message == "JOIN is not supported" ) )); } @@ -284,7 +298,7 @@ mod tests { validate_s3_select_statement(&statement), Err(ref err) if matches!( err.s3_select_policy_error(), - Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "subqueries are not supported" + Some(SelectError::UnsupportedSqlStructure { message }) if message == "subqueries are not supported" ) )); } @@ -297,7 +311,7 @@ mod tests { validate_s3_select_statement(&statement), Err(ref err) if matches!( err.s3_select_policy_error(), - Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "subqueries are not supported" + Some(SelectError::UnsupportedSqlStructure { message }) if message == "subqueries are not supported" ) )); } @@ -310,7 +324,7 @@ mod tests { validate_s3_select_statement(&statement), Err(ref err) if matches!( err.s3_select_policy_error(), - Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "the source must be S3Object" + Some(SelectError::DataSourcePathUnsupported) ) )); } @@ -326,7 +340,7 @@ mod tests { assert!( matches!( validate_s3_select_statement(&statement), - Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::UnsupportedSqlStructure { .. })) + Err(ref err) if matches!(err.s3_select_policy_error(), Some(SelectError::UnsupportedSqlStructure { .. })) ), "query should be rejected: {sql}" ); diff --git a/crates/s3select-query/src/test/error_handling_test.rs b/crates/s3select-query/src/test/error_handling_test.rs index 72a2144fe..1462719ac 100644 --- a/crates/s3select-query/src/test/error_handling_test.rs +++ b/crates/s3select-query/src/test/error_handling_test.rs @@ -16,7 +16,7 @@ mod error_handling_tests { use crate::get_global_db; use rustfs_s3select_api::{ - QueryError, + QueryError, SelectError, query::{Context, Query}, }; use s3s::dto::{ @@ -98,7 +98,6 @@ mod error_handling_tests { "INSERT INTO S3Object VALUES (1, 'test')", "UPDATE S3Object SET name = 'test'", "DELETE FROM S3Object", - "CREATE TABLE test (id INT)", "DROP TABLE S3Object", ]; @@ -113,6 +112,68 @@ mod error_handling_tests { } } + #[tokio::test] + async fn test_non_select_statement_is_typed_unsupported_structure() { + let sql = "CREATE TABLE test (id INT)"; + let input = create_test_input_with_sql(sql); + let db = get_global_db(input.clone(), true).await.unwrap(); + let query = Query::new(Context { input: Arc::new(input) }, sql.to_string()); + + let error = match db.execute(&query).await { + Err(error) => error, + Ok(_) => panic!("non-SELECT statement must fail"), + }; + + assert!(matches!(error.select_error(), SelectError::UnsupportedSqlStructure { .. })); + } + + #[tokio::test] + async fn test_function_argument_coercion_failure_is_typed() { + for sql in ["SELECT ROUND(3.14, 1.1) FROM S3Object", "SELECT SQRT(1, 2) FROM S3Object"] { + let input = create_test_input_with_sql(sql); + let db = get_global_db(input.clone(), true) + .await + .expect("test database should initialize"); + let query = Query::new(Context { input: Arc::new(input) }, sql.to_string()); + + let error = match db.execute(&query).await { + Err(error) => error, + Ok(_) => panic!("invalid function arguments must fail during planning: {sql}"), + }; + + assert_eq!( + error.select_error(), + SelectError::IncorrectSqlFunctionArgumentType, + "unexpected planner error for {sql}: {error:?}" + ); + } + } + + #[tokio::test] + async fn test_other_planner_failures_remain_invalid_query() { + for sql in [ + "SELECT DEFINITELY_UNKNOWN_FUNCTION(1) FROM S3Object", + "SELECT 1 + 'text' FROM S3Object", + ] { + let input = create_test_input_with_sql(sql); + let db = get_global_db(input.clone(), true) + .await + .expect("test database should initialize"); + let query = Query::new(Context { input: Arc::new(input) }, sql.to_string()); + + let error = match db.execute(&query).await { + Err(error) => error, + Ok(_) => panic!("invalid query must fail during planning: {sql}"), + }; + + assert_eq!( + error.select_error(), + SelectError::InvalidQuery, + "unexpected planner error for {sql}: {error:?}" + ); + } + } + #[tokio::test] async fn test_invalid_column_references() { let invalid_column_sqls = vec![ diff --git a/rustfs/src/app/select_object.rs b/rustfs/src/app/select_object.rs index 0b4904893..743b70f37 100644 --- a/rustfs/src/app/select_object.rs +++ b/rustfs/src/app/select_object.rs @@ -15,12 +15,13 @@ use datafusion::arrow::{ json::{WriterBuilder as JsonWriterBuilder, writer::LineDelimited}, record_batch::RecordBatch, }; +#[cfg(test)] use datafusion::common::DataFusionError; use datafusion::physical_plan::SendableRecordBatchStream; use futures::StreamExt; use http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header::RANGE}; use rustfs_s3select_api::{ - QueryError, S3SelectPolicyError, + QueryError, SelectError, object_store::{INVALID_SCAN_RANGE_MESSAGE, validate_scan_range_bounds}, query::{Context, Query}, }; @@ -49,8 +50,13 @@ use tracing::info; const MAX_SELECT_EXPRESSION_BYTES: usize = 256 * 1024; const RECORDS_CHUNK_TARGET: usize = 128 * 1024; +const DATA_SOURCE_PATH_UNSUPPORTED_CODE: &str = "DataSourcePathUnsupported"; +const INVALID_QUERY_CODE: &str = "InvalidQuery"; const PARSE_SELECT_FAILURE_CODE: &str = "ParseSelectFailure"; +const BUSY_MESSAGE: &str = "The service is unavailable. Try again later."; const EMPTY_SELECT_EXPRESSION_MESSAGE: &str = "empty SQL expression"; +const SLOW_DOWN_MESSAGE: &str = "Reduce your request rate."; +const UNSUPPORTED_SQL_STRUCTURE_MESSAGE: &str = "We encountered an unsupported SQL structure. Check the SQL Reference."; const SELECT_MINIO_SSEC_SEALED_KEY: &str = "X-Minio-Internal-Server-Side-Encryption-Sealed-Key"; const SELECT_MINIO_S3_SEALED_KEY: &str = "X-Minio-Internal-Server-Side-Encryption-S3-Sealed-Key"; const SELECT_MINIO_KMS_SEALED_KEY: &str = "X-Minio-Internal-Server-Side-Encryption-Kms-Sealed-Key"; @@ -82,12 +88,7 @@ trait SelectSnapshotFence { impl SelectSnapshotFence for Arc { fn ensure_snapshot_valid(&self) -> S3Result<()> { - self.ensure_valid().map_err(|error| { - let message = error.to_string(); - let mut s3_error = S3Error::with_message(S3ErrorCode::InternalError, message); - s3_error.set_source(Box::new(error)); - s3_error - }) + self.ensure_valid().map_err(internal_select_error) } } @@ -128,7 +129,7 @@ pub async fn execute_select_object_content( let terminal_permit = tx .clone() .try_reserve_owned() - .map_err(|_| s3_error!(InternalError, "can't reserve Select terminal event capacity"))?; + .map_err(|_| map_select_error_to_s3(&SelectError::InternalError))?; let response = select_object_response(rx, &snapshot.object_info().user_defined, &req.headers)?; spawn_traced(async move { send_select_events_until_deadline( @@ -338,7 +339,7 @@ async fn send_select_events_until_deadline( let outcome = match timeout_at(deadline, send_select_events(output, &tx, validation, &snapshot_lease)).await { Ok(outcome) => outcome, Err(_) => SelectProducerOutcome::Terminal(Err(map_query_error_to_s3( - S3SelectPolicyError::QueryTimeout { + SelectError::QueryTimeout { seconds: timeout_seconds, } .into(), @@ -416,12 +417,14 @@ async fn send_select_events( let stats = SelectObjectContentEvent::Stats(StatsEvent { details: Some(progress.to_stats()), }); - if tx.send(Ok(stats)).await.is_err() { - return SelectProducerOutcome::ReceiverClosed; - } + let stats_permit = match tx.reserve().await { + Ok(permit) => permit, + Err(_) => return SelectProducerOutcome::ReceiverClosed, + }; if let Err(error) = snapshot_fence.ensure_snapshot_valid() { return SelectProducerOutcome::Terminal(Err(error)); } + stats_permit.send(Ok(stats)); SelectProducerOutcome::Terminal(Ok(SelectObjectContentEvent::End(EndEvent::default()))) } @@ -441,7 +444,9 @@ fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectCo let output_format = normalize_output_serialization(&mut input.request.output_serialization)?; if input.request.expression.trim().is_empty() { - return Err(parse_select_failure(EMPTY_SELECT_EXPRESSION_MESSAGE)); + return Err(map_select_error_to_s3(&SelectError::ParseSelectFailure { + message: EMPTY_SELECT_EXPRESSION_MESSAGE.to_string(), + })); } let progress_enabled = input .request @@ -466,13 +471,17 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> return Err(S3Error::new(S3ErrorCode::ObjectSerializationConflict)); } - if let Some(compression) = input.compression_type.as_ref() - && compression.as_str() != CompressionType::NONE - { - return Err(s3_error!( - NotImplemented, - "SelectObjectContent currently supports only uncompressed input" - )); + if let Some(compression) = input.compression_type.as_ref() { + match compression.as_str() { + CompressionType::NONE => {} + CompressionType::GZIP | CompressionType::BZIP2 => { + return Err(s3_error!( + NotImplemented, + "SelectObjectContent currently supports only uncompressed input" + )); + } + _ => return Err(map_select_error_to_s3(&SelectError::InvalidCompressionFormat)), + } } input.compression_type = Some(CompressionType::from_static(CompressionType::NONE)); @@ -483,8 +492,18 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> "CSV AllowQuotedRecordDelimiter is not supported by SelectObjectContent" )); } - csv.file_header_info + let file_header_info = csv + .file_header_info .get_or_insert_with(|| FileHeaderInfo::from_static(FileHeaderInfo::NONE)); + if !matches!( + file_header_info.as_str(), + FileHeaderInfo::NONE | FileHeaderInfo::USE | FileHeaderInfo::IGNORE + ) { + return Err(S3Error::with_message( + S3ErrorCode::InvalidFileHeaderInfo, + "The FileHeaderInfo value is not valid. Only NONE, USE, and IGNORE are supported.", + )); + } validate_single_byte(csv.comments.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; @@ -575,12 +594,6 @@ fn invalid_scan_range_error() -> S3Error { S3Error::with_message(S3ErrorCode::InvalidRequestParameter, INVALID_SCAN_RANGE_MESSAGE.to_string()) } -fn parse_select_failure(message: impl Into) -> S3Error { - let mut err = S3Error::with_message(S3ErrorCode::Custom(PARSE_SELECT_FAILURE_CODE.into()), message.into()); - err.set_status_code(StatusCode::BAD_REQUEST); - err -} - fn validate_single_byte(value: Option<&str>, code: S3ErrorCode) -> S3Result<()> { if let Some(value) = value && value.len() != 1 @@ -646,12 +659,14 @@ async fn prepare_select_object_snapshot( fn map_prepare_snapshot_error(err: StoragePrepareSelectObjectSnapshotError) -> S3Error { match err { - StoragePrepareSelectObjectSnapshotError::Storage(err) => ApiError::from(err).into(), - err => { - let mut s3_error = S3Error::with_message(S3ErrorCode::InternalError, err.to_string()); - s3_error.set_source(Box::new(err)); + StoragePrepareSelectObjectSnapshotError::Storage(err) => { + let mut s3_error: S3Error = ApiError::from(err).into(); + if s3_error.code() == &S3ErrorCode::InternalError { + s3_error.set_message(SelectError::InternalError.to_string()); + } s3_error } + err => internal_select_error(err), } } @@ -719,9 +734,7 @@ fn encode_csv_batch(batch: &RecordBatch, config: &CSVOutput) -> S3Result } let mut writer = builder.build(&mut buffer); - writer - .write(batch) - .map_err(|err| s3_error!(InternalError, "can't encode Select output to CSV: {}", err))?; + writer.write(batch).map_err(internal_select_error)?; drop(writer); Ok(buffer) } @@ -739,12 +752,8 @@ fn encode_json_batch(batch: &RecordBatch, config: &JSONOutput) -> S3Result(&mut buffer); - writer - .write(batch) - .map_err(|err| s3_error!(InternalError, "can't encode Select output to JSON: {}", err))?; - writer - .finish() - .map_err(|err| s3_error!(InternalError, "can't finish Select JSON output: {}", err))?; + writer.write(batch).map_err(internal_select_error)?; + writer.finish().map_err(internal_select_error)?; drop(writer); if let Some(delimiter) = config.record_delimiter.as_deref() @@ -813,122 +822,60 @@ fn clamp_i64(value: u64) -> i64 { } fn map_query_error_to_s3(err: QueryError) -> S3Error { - if err.is_snapshot_consistency_error() { - let message = err.to_string(); - let mut s3_error = S3Error::with_message(S3ErrorCode::InternalError, message); - s3_error.set_source(Box::new(err)); - return s3_error; - } - if let Some(policy_error) = err.s3_select_policy_error() { - let message = policy_error.to_string(); - return match policy_error { - S3SelectPolicyError::UnsupportedSqlStructure { .. } => { - S3Error::with_message(S3ErrorCode::UnsupportedSqlStructure, message) - } - S3SelectPolicyError::QueryConcurrencyLimit => S3Error::with_message(S3ErrorCode::SlowDown, message), - S3SelectPolicyError::QueryTimeout { .. } => S3Error::with_message(S3ErrorCode::Busy, message), - _ => S3Error::with_message(S3ErrorCode::InternalError, message), - }; - } - let message = err.to_string(); + let select_error = err.select_error(); + map_select_error_to_s3(&select_error) +} + +fn map_select_error_to_s3(err: &SelectError) -> S3Error { match err { - QueryError::Parser { .. } => parse_select_failure(message), - QueryError::MultiStatement { .. } => S3Error::with_message(S3ErrorCode::UnsupportedSqlStructure, message), - QueryError::NotImplemented { .. } => S3Error::with_message(S3ErrorCode::NotImplemented, message), - QueryError::Datafusion { source } if is_resource_exhausted(source.as_ref()) => { - S3Error::with_message(S3ErrorCode::Busy, message) + SelectError::InvalidCompressionFormat => S3Error::with_message(S3ErrorCode::InvalidCompressionFormat, err.to_string()), + SelectError::InvalidDataSource => S3Error::with_message(S3ErrorCode::InvalidDataSource, err.to_string()), + SelectError::TruncatedInput => S3Error::with_message(S3ErrorCode::TruncatedInput, err.to_string()), + SelectError::CsvParsingError => S3Error::with_message(S3ErrorCode::CSVParsingError, err.to_string()), + SelectError::JsonParsingError => S3Error::with_message(S3ErrorCode::JSONParsingError, err.to_string()), + SelectError::ParquetParsingError => S3Error::with_message(S3ErrorCode::ParquetParsingError, err.to_string()), + SelectError::ParseSelectFailure { message } => custom_bad_request(PARSE_SELECT_FAILURE_CODE, message.clone()), + SelectError::InvalidQuery => custom_bad_request(INVALID_QUERY_CODE, err.to_string()), + SelectError::InvalidDataType => S3Error::with_message(S3ErrorCode::InvalidDataType, err.to_string()), + SelectError::IncorrectSqlFunctionArgumentType => { + S3Error::with_message(S3ErrorCode::IncorrectSqlFunctionArgumentType, err.to_string()) } - QueryError::Datafusion { source } if is_unexpected_eof(source.as_ref()) => { - S3Error::with_message(S3ErrorCode::InternalError, message) + SelectError::DataSourcePathUnsupported => custom_bad_request(DATA_SOURCE_PATH_UNSUPPORTED_CODE, err.to_string()), + SelectError::UnsupportedSqlStructure { .. } => { + S3Error::with_message(S3ErrorCode::UnsupportedSqlStructure, UNSUPPORTED_SQL_STRUCTURE_MESSAGE) } - QueryError::Datafusion { source } if is_invalid_object_size(source.as_ref()) => { - S3Error::with_message(S3ErrorCode::InternalError, message) + SelectError::UnsupportedSqlOperation => S3Error::with_message(S3ErrorCode::UnsupportedSqlOperation, err.to_string()), + SelectError::EvaluatorBindingDoesNotExist => { + S3Error::with_message(S3ErrorCode::EvaluatorBindingDoesNotExist, err.to_string()) } - QueryError::Datafusion { .. } if looks_like_invalid_scan_range(&message) => { + SelectError::AmbiguousFieldName => S3Error::with_message(S3ErrorCode::AmbiguousFieldName, err.to_string()), + SelectError::InvalidScanRange => { S3Error::with_message(S3ErrorCode::InvalidRequestParameter, INVALID_SCAN_RANGE_MESSAGE.to_string()) } - QueryError::Datafusion { .. } if looks_like_missing_binding(&message) => { - S3Error::with_message(S3ErrorCode::EvaluatorBindingDoesNotExist, message) + SelectError::QueryConcurrencyLimit => S3Error::with_message(S3ErrorCode::SlowDown, SLOW_DOWN_MESSAGE), + SelectError::QueryTimeout { .. } | SelectError::ResourceExhausted => { + S3Error::with_message(S3ErrorCode::Busy, BUSY_MESSAGE) } - QueryError::Datafusion { .. } => S3Error::with_message(S3ErrorCode::UnsupportedSqlOperation, message), - QueryError::StoreError { .. } if looks_like_invalid_scan_range(&message) => { - S3Error::with_message(S3ErrorCode::InvalidRequestParameter, INVALID_SCAN_RANGE_MESSAGE.to_string()) + SelectError::BucketNotFound => S3Error::with_message(S3ErrorCode::NoSuchBucket, err.to_string()), + SelectError::ObjectNotFound => S3Error::with_message(S3ErrorCode::NoSuchKey, err.to_string()), + SelectError::Canceled | SelectError::InternalError => { + S3Error::with_message(S3ErrorCode::InternalError, SelectError::InternalError.to_string()) } - QueryError::StoreError { .. } if looks_like_bucket_not_found(&message) => { - S3Error::with_message(S3ErrorCode::NoSuchBucket, message) - } - QueryError::StoreError { .. } if looks_like_object_not_found(&message) => { - S3Error::with_message(S3ErrorCode::NoSuchKey, message) - } - QueryError::StoreError { .. } => S3Error::with_message(S3ErrorCode::InternalError, message), - QueryError::BuildQueryDispatcher { .. } - | QueryError::Cancel - | QueryError::FunctionNotExists { .. } - | QueryError::FunctionExists { .. } => S3Error::with_message(S3ErrorCode::InternalError, message), } } +fn internal_select_error(_error: impl std::error::Error + Send + Sync + 'static) -> S3Error { + map_select_error_to_s3(&SelectError::InternalError) +} + +fn custom_bad_request(code: &'static str, message: String) -> S3Error { + let mut err = S3Error::with_message(S3ErrorCode::Custom(code.into()), message); + err.set_status_code(StatusCode::BAD_REQUEST); + err +} + fn select_query_timeout_error(seconds: u64) -> S3Error { - map_query_error_to_s3(S3SelectPolicyError::QueryTimeout { seconds }.into()) -} - -fn looks_like_bucket_not_found(message: &str) -> bool { - message.contains("NoSuchBucket") || message.contains("bucket not found") || message.contains("BucketNotFound") -} - -const MAX_ERROR_SOURCE_DEPTH: usize = 16; - -fn error_chain_any( - mut err: &(dyn std::error::Error + 'static), - predicate: impl Fn(&(dyn std::error::Error + 'static)) -> bool, -) -> bool { - for _ in 0..MAX_ERROR_SOURCE_DEPTH { - if predicate(err) { - return true; - } - let Some(source) = err.source() else { - return false; - }; - err = source; - } - false -} - -fn is_resource_exhausted(err: &(dyn std::error::Error + 'static)) -> bool { - error_chain_any(err, |err| { - err.downcast_ref::() - .is_some_and(|err| matches!(err, DataFusionError::ResourcesExhausted(_))) - }) -} - -fn is_unexpected_eof(err: &(dyn std::error::Error + 'static)) -> bool { - error_chain_any(err, |err| { - err.downcast_ref::() - .is_some_and(|err| err.kind() == std::io::ErrorKind::UnexpectedEof) - }) -} - -fn is_invalid_object_size(err: &(dyn std::error::Error + 'static)) -> bool { - error_chain_any(err, |err| err.downcast_ref::().is_some()) -} - -fn looks_like_object_not_found(message: &str) -> bool { - message.contains("NoSuchKey") - || message.contains("NoSuchVersion") - || message.contains("ObjectNotFound") - || message.contains("object not found") - || message.contains("NotFound") -} - -fn looks_like_missing_binding(message: &str) -> bool { - message.contains("No field named") - || message.contains("field not found") - || message.contains("Schema error") - || message.contains("No such column") -} - -fn looks_like_invalid_scan_range(message: &str) -> bool { - message.contains("ScanRange:") || message.contains(INVALID_SCAN_RANGE_MESSAGE) + map_query_error_to_s3(SelectError::QueryTimeout { seconds }.into()) } fn is_json_document(json: &JSONInput) -> bool { @@ -942,8 +889,9 @@ mod tests { use super::*; use datafusion::{ arrow::{ - array::{Array, ListArray}, - datatypes::{Field, Int32Type, Schema}, + array::{Array, ListArray, StringArray}, + datatypes::{DataType, Field, Int32Type, Schema}, + error::ArrowError, }, physical_plan::stream::RecordBatchStreamAdapter, sql::sqlparser::parser::ParserError, @@ -951,6 +899,53 @@ mod tests { use rustfs_test_utils::TestECStoreEnv; use s3s::dto::{CSVInput, ParquetInput, ScanRange}; + fn event_stream_headers(mut bytes: &[u8]) -> Vec> { + let mut messages = Vec::new(); + while !bytes.is_empty() { + assert!(bytes.len() >= 16, "event-stream message is truncated"); + let total_len = u32::from_be_bytes(bytes[0..4].try_into().expect("event-stream total length")) as usize; + let headers_len = u32::from_be_bytes(bytes[4..8].try_into().expect("event-stream headers length")) as usize; + assert!(total_len >= 16 && total_len <= bytes.len(), "invalid event-stream message length"); + assert!(12 + headers_len <= total_len - 4, "invalid event-stream headers length"); + + let mut headers = &bytes[12..12 + headers_len]; + let mut decoded = Vec::new(); + while !headers.is_empty() { + let name_len = headers[0] as usize; + assert!(headers.len() >= name_len + 4, "event-stream header is truncated"); + let name = std::str::from_utf8(&headers[1..1 + name_len]) + .expect("event-stream header name should be UTF-8") + .to_string(); + assert_eq!(headers[1 + name_len], 7, "expected an event-stream string header"); + let value_len = u16::from_be_bytes( + headers[2 + name_len..4 + name_len] + .try_into() + .expect("event-stream header value length"), + ) as usize; + assert!(headers.len() >= name_len + 4 + value_len, "event-stream header value is truncated"); + let value = std::str::from_utf8(&headers[4 + name_len..4 + name_len + value_len]) + .expect("event-stream header value should be UTF-8") + .to_string(); + decoded.push((name, value)); + headers = &headers[4 + name_len + value_len..]; + } + messages.push(decoded); + bytes = &bytes[total_len..]; + } + messages + } + + async fn http_xml_error(error: S3Error) -> (StatusCode, String) { + let response = error.to_http_response().expect("S3 error should serialize to HTTP"); + let status = response.status(); + let body = http_body_util::BodyExt::collect(response.into_body()) + .await + .expect("S3 error body should be readable") + .to_bytes(); + let body = std::str::from_utf8(&body).expect("S3 error XML should be UTF-8").to_string(); + (status, body) + } + struct LeaseDropSignal(Option>); impl Drop for LeaseDropSignal { @@ -1010,22 +1005,8 @@ mod tests { .expect_err("production fence adapter must reject a lost storage snapshot"); assert_eq!(error.code(), &S3ErrorCode::InternalError); - assert!(error.to_string().contains("namespace lock was lost")); - } - - #[derive(Debug)] - struct CyclicError; - - impl std::fmt::Display for CyclicError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str("cyclic error") - } - } - - impl std::error::Error for CyclicError { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - Some(self) - } + assert_eq!(error.message(), Some("An internal error occurred.")); + assert!(error.source().is_none()); } fn base_input() -> SelectObjectContentInput { @@ -1272,15 +1253,15 @@ mod tests { #[test] fn map_query_policy_errors_to_s3_errors() { let unsupported = map_query_error_to_s3( - S3SelectPolicyError::UnsupportedSqlStructure { + SelectError::UnsupportedSqlStructure { message: "JOIN is not supported".to_string(), } .into(), ); - let saturated = map_query_error_to_s3(S3SelectPolicyError::QueryConcurrencyLimit.into()); - let timed_out = map_query_error_to_s3(S3SelectPolicyError::QueryTimeout { seconds: 300 }.into()); + let saturated = map_query_error_to_s3(SelectError::QueryConcurrencyLimit.into()); + let timed_out = map_query_error_to_s3(SelectError::QueryTimeout { seconds: 300 }.into()); let stream_timed_out = map_query_error_to_s3(QueryError::Datafusion { - source: Box::new(DataFusionError::External(Box::new(S3SelectPolicyError::QueryTimeout { seconds: 300 }))), + source: Box::new(DataFusionError::External(Box::new(SelectError::QueryTimeout { seconds: 300 }))), }); let exhausted = map_query_error_to_s3(QueryError::Datafusion { source: Box::new(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::Generic { @@ -1289,6 +1270,9 @@ mod tests { }))), }); let truncated = map_query_error_to_s3(QueryError::Datafusion { + source: Box::new(DataFusionError::External(Box::new(SelectError::TruncatedInput))), + }); + let raw_storage_short_read = map_query_error_to_s3(QueryError::Datafusion { source: Box::new(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::Generic { store: "EcObjectStore", source: Box::new(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated object stream")), @@ -1302,19 +1286,114 @@ mod tests { }); assert_eq!(unsupported.code(), &S3ErrorCode::UnsupportedSqlStructure); - assert_eq!(unsupported.message(), Some("Unsupported S3 Select SQL structure: JOIN is not supported")); + assert_eq!(unsupported.message(), Some(UNSUPPORTED_SQL_STRUCTURE_MESSAGE)); assert_eq!(saturated.code(), &S3ErrorCode::SlowDown); - assert_eq!(saturated.message(), Some("S3 Select query concurrency limit reached")); assert_eq!(timed_out.code(), &S3ErrorCode::Busy); - assert_eq!(timed_out.message(), Some("S3 Select query exceeded the 300-second execution limit")); assert_eq!(stream_timed_out.code(), &S3ErrorCode::Busy); - assert_eq!( - stream_timed_out.message(), - Some("S3 Select query exceeded the 300-second execution limit") - ); assert_eq!(exhausted.code(), &S3ErrorCode::Busy); - assert_eq!(truncated.code(), &S3ErrorCode::InternalError); + assert_eq!(truncated.code(), &S3ErrorCode::TruncatedInput); + assert_eq!(raw_storage_short_read.code(), &S3ErrorCode::InternalError); assert_eq!(invalid_object_size.code(), &S3ErrorCode::InternalError); + assert_eq!(invalid_object_size.message(), Some("An internal error occurred.")); + } + + #[test] + fn every_select_error_has_an_explicit_protocol_mapping() { + let cases = vec![ + ( + SelectError::InvalidCompressionFormat, + S3ErrorCode::InvalidCompressionFormat, + StatusCode::BAD_REQUEST, + ), + (SelectError::InvalidDataSource, S3ErrorCode::InvalidDataSource, StatusCode::BAD_REQUEST), + (SelectError::TruncatedInput, S3ErrorCode::TruncatedInput, StatusCode::BAD_REQUEST), + (SelectError::CsvParsingError, S3ErrorCode::CSVParsingError, StatusCode::BAD_REQUEST), + (SelectError::JsonParsingError, S3ErrorCode::JSONParsingError, StatusCode::BAD_REQUEST), + ( + SelectError::ParquetParsingError, + S3ErrorCode::ParquetParsingError, + StatusCode::BAD_REQUEST, + ), + ( + SelectError::ParseSelectFailure { + message: "invalid SELECT expression".to_string(), + }, + S3ErrorCode::Custom(PARSE_SELECT_FAILURE_CODE.into()), + StatusCode::BAD_REQUEST, + ), + ( + SelectError::InvalidQuery, + S3ErrorCode::Custom(INVALID_QUERY_CODE.into()), + StatusCode::BAD_REQUEST, + ), + (SelectError::InvalidDataType, S3ErrorCode::InvalidDataType, StatusCode::BAD_REQUEST), + ( + SelectError::IncorrectSqlFunctionArgumentType, + S3ErrorCode::IncorrectSqlFunctionArgumentType, + StatusCode::BAD_REQUEST, + ), + ( + SelectError::DataSourcePathUnsupported, + S3ErrorCode::Custom(DATA_SOURCE_PATH_UNSUPPORTED_CODE.into()), + StatusCode::BAD_REQUEST, + ), + ( + SelectError::UnsupportedSqlStructure { + message: "JOIN is not supported".to_string(), + }, + S3ErrorCode::UnsupportedSqlStructure, + StatusCode::BAD_REQUEST, + ), + ( + SelectError::UnsupportedSqlOperation, + S3ErrorCode::UnsupportedSqlOperation, + StatusCode::BAD_REQUEST, + ), + ( + SelectError::EvaluatorBindingDoesNotExist, + S3ErrorCode::EvaluatorBindingDoesNotExist, + StatusCode::BAD_REQUEST, + ), + (SelectError::AmbiguousFieldName, S3ErrorCode::AmbiguousFieldName, StatusCode::BAD_REQUEST), + ( + SelectError::InvalidScanRange, + S3ErrorCode::InvalidRequestParameter, + StatusCode::BAD_REQUEST, + ), + (SelectError::QueryConcurrencyLimit, S3ErrorCode::SlowDown, StatusCode::SERVICE_UNAVAILABLE), + ( + SelectError::QueryTimeout { seconds: 300 }, + S3ErrorCode::Busy, + StatusCode::SERVICE_UNAVAILABLE, + ), + (SelectError::ResourceExhausted, S3ErrorCode::Busy, StatusCode::SERVICE_UNAVAILABLE), + (SelectError::BucketNotFound, S3ErrorCode::NoSuchBucket, StatusCode::NOT_FOUND), + (SelectError::ObjectNotFound, S3ErrorCode::NoSuchKey, StatusCode::NOT_FOUND), + (SelectError::Canceled, S3ErrorCode::InternalError, StatusCode::INTERNAL_SERVER_ERROR), + (SelectError::InternalError, S3ErrorCode::InternalError, StatusCode::INTERNAL_SERVER_ERROR), + ]; + + for (select_error, expected_code, expected_status) in cases { + let error = map_select_error_to_s3(&select_error); + assert_eq!(error.code(), &expected_code, "wrong mapping for {select_error:?}"); + assert_eq!(error.status_code(), Some(expected_status), "wrong status for {select_error:?}"); + assert!( + error.message().is_some_and(|message| !message.is_empty()), + "missing protocol message for {select_error:?}" + ); + } + } + + #[test] + fn internal_query_details_are_not_exposed_to_clients() { + let private_detail = "node-1:/private/object/path physical_plan=secret"; + let error = map_query_error_to_s3(QueryError::from(DataFusionError::Internal(private_detail.to_string()))); + + assert_eq!(error.code(), &S3ErrorCode::InternalError); + assert_eq!(error.message(), Some("An internal error occurred.")); + assert!(!error.message().is_some_and(|message| message.contains(private_detail))); + assert!(!format!("{error:?}").contains(private_detail)); + assert!(error.source().is_none()); } #[test] @@ -1329,23 +1408,12 @@ mod tests { } #[test] - fn prepare_snapshot_invalid_logical_size_fails_with_internal_error_and_source() { + fn prepare_snapshot_invalid_logical_size_fails_with_redacted_internal_error() { let err = map_prepare_snapshot_error(StoragePrepareSelectObjectSnapshotError::InvalidLogicalSize { size: -1 }); assert_eq!(err.code(), &S3ErrorCode::InternalError); - assert!( - err.source() - .is_some_and(|source| source.downcast_ref::().is_some()) - ); - } - - #[test] - fn error_source_matching_stops_at_the_depth_bound() { - let err = CyclicError; - - assert!(!is_resource_exhausted(&err)); - assert!(!is_unexpected_eof(&err)); - assert!(!is_invalid_object_size(&err)); + assert!(err.source().is_none()); + assert_eq!(err.message(), Some("An internal error occurred.")); } #[tokio::test(start_paused = true)] @@ -1422,6 +1490,39 @@ mod tests { assert!(lease_released.await.is_ok(), "End should release the snapshot lease"); } + #[tokio::test(start_paused = true)] + async fn successful_stream_serializes_records_stats_and_end_without_error() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, false)])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(vec!["row"]))]) + .expect("test record batch should be valid"); + let output = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { Ok::<_, DataFusionError>(batch) }), + )); + let (producer, rx, lease_released) = spawn_test_producer(output, 4); + producer.await.expect("producer should finish successfully"); + + let mut byte_stream = SelectObjectContentEventStream::new(ReceiverStream::new(rx)).into_byte_stream(); + let mut encoded = Vec::new(); + while let Some(chunk) = byte_stream.next().await { + encoded.extend_from_slice(&chunk.expect("event-stream message should serialize")); + } + let messages = event_stream_headers(&encoded); + let event_types = messages + .iter() + .filter_map(|headers| { + headers + .iter() + .find_map(|(name, value)| (name == ":event-type").then_some(value.as_str())) + }) + .collect::>(); + assert_eq!(event_types, ["Cont", "Records", "Stats", "End"]); + assert!(!messages.iter().flatten().any(|(name, value)| { + (name == ":message-type" && value == "error") || name == ":error-code" || name == ":error-message" + })); + assert!(lease_released.await.is_ok(), "End should release the snapshot lease"); + } + #[tokio::test(start_paused = true)] async fn eof_at_deadline_uses_reserved_slot_for_stats_then_end() { let output = Box::pin(RecordBatchStreamAdapter::new( @@ -1455,7 +1556,7 @@ mod tests { Arc::new(Schema::empty()), futures::stream::once(async { tokio::time::sleep(std::time::Duration::from_secs(1)).await; - Err(DataFusionError::External(Box::new(S3SelectPolicyError::QueryConcurrencyLimit))) + Err(DataFusionError::External(Box::new(SelectError::QueryConcurrencyLimit))) }), )); let (producer, mut rx, lease_released) = spawn_test_producer(output, 2); @@ -1475,6 +1576,172 @@ mod tests { assert!(lease_released.await.is_ok(), "stream error should release the snapshot lease"); } + #[tokio::test(start_paused = true)] + async fn select_errors_use_http_codes_before_stream_and_error_frames_after_stream() { + fn csv_error() -> DataFusionError { + DataFusionError::ArrowError(Box::new(ArrowError::CsvError("private CSV parser state".to_string())), None) + } + fn json_error() -> DataFusionError { + DataFusionError::ArrowError(Box::new(ArrowError::JsonError("private JSON parser state".to_string())), None) + } + fn parquet_error() -> DataFusionError { + DataFusionError::ParquetError(Box::new(datafusion::parquet::errors::ParquetError::General( + "private Parquet parser state".to_string(), + ))) + } + fn truncated_error() -> DataFusionError { + DataFusionError::External(Box::new(SelectError::TruncatedInput)) + } + fn timeout_error() -> DataFusionError { + DataFusionError::External(Box::new(SelectError::QueryTimeout { seconds: 300 })) + } + + let cases = [ + ( + csv_error as fn() -> DataFusionError, + S3ErrorCode::CSVParsingError, + StatusCode::BAD_REQUEST, + b"CSVParsingError" as &[u8], + ), + (json_error, S3ErrorCode::JSONParsingError, StatusCode::BAD_REQUEST, b"JSONParsingError"), + ( + parquet_error, + S3ErrorCode::ParquetParsingError, + StatusCode::BAD_REQUEST, + b"ParquetParsingError", + ), + (truncated_error, S3ErrorCode::TruncatedInput, StatusCode::BAD_REQUEST, b"TruncatedInput"), + (timeout_error, S3ErrorCode::Busy, StatusCode::SERVICE_UNAVAILABLE, b"Busy"), + ]; + + for (source, expected_code, expected_status, encoded_code) in cases { + let pre_stream = map_query_error_to_s3(QueryError::from(source())); + assert_eq!(pre_stream.code(), &expected_code); + assert_eq!(pre_stream.status_code(), Some(expected_status)); + let expected_code_text = expected_code.as_str().to_string(); + let (status, body) = http_xml_error(pre_stream).await; + assert_eq!(status, expected_status); + assert!(body.contains(&format!("{expected_code_text}"))); + assert!(body.contains("")); + + let output = Box::pin(RecordBatchStreamAdapter::new( + Arc::new(Schema::empty()), + futures::stream::once(async move { Err(source()) }), + )); + let (producer, rx, lease_released) = spawn_test_producer(output, 2); + producer.await.expect("producer should emit the terminal Select error"); + + let mut byte_stream = SelectObjectContentEventStream::new(ReceiverStream::new(rx)).into_byte_stream(); + let mut encoded = Vec::new(); + while let Some(chunk) = byte_stream.next().await { + encoded.extend_from_slice(&chunk.expect("event-stream message should serialize")); + } + let messages = event_stream_headers(&encoded); + let terminal_headers = messages.last().expect("event stream should contain a terminal error"); + let encoded_code = std::str::from_utf8(encoded_code).expect("test error code should be UTF-8"); + assert!( + terminal_headers + .iter() + .any(|(name, value)| name == ":message-type" && value == "error") + ); + assert!( + terminal_headers + .iter() + .any(|(name, value)| name == ":error-code" && value == encoded_code) + ); + assert!( + terminal_headers + .iter() + .any(|(name, value)| name == ":error-message" && !value.is_empty() && !value.contains("private")) + ); + assert!( + !messages + .iter() + .flatten() + .any(|(name, value)| { name == ":event-type" && matches!(value.as_str(), "Stats" | "End") }) + ); + assert!(lease_released.await.is_ok(), "error frame should release the snapshot lease"); + } + } + + #[tokio::test] + async fn sql_and_compression_errors_serialize_as_http_xml() { + let sql_error = map_query_error_to_s3(QueryError::Parser { + source: ParserError::ParserError("unexpected token".to_string()), + }); + let (sql_status, sql_body) = http_xml_error(sql_error).await; + assert_eq!(sql_status, StatusCode::BAD_REQUEST); + assert!(sql_body.contains("ParseSelectFailure")); + assert!(sql_body.contains("")); + + let mut input = base_input(); + input.request.input_serialization.compression_type = Some(CompressionType::from_static("SNAPPY")); + let compression_error = + validate_select_request(&HeaderMap::new(), &mut input).expect_err("unknown compression must fail before streaming"); + let (compression_status, compression_body) = http_xml_error(compression_error).await; + assert_eq!(compression_status, StatusCode::BAD_REQUEST); + assert!(compression_body.contains("InvalidCompressionFormat")); + assert!(compression_body.contains("")); + } + + #[tokio::test(start_paused = true)] + async fn stream_error_after_records_omits_stats_and_end() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, false)])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(vec!["row"]))]) + .expect("test record batch should be valid"); + let output = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter([ + Ok(batch), + Err(DataFusionError::ArrowError( + Box::new(ArrowError::CsvError("private CSV parser state".to_string())), + None, + )), + ]), + )); + let (producer, rx, lease_released) = spawn_test_producer(output, 3); + + producer + .await + .expect("producer should emit records followed by the terminal error"); + + let mut byte_stream = SelectObjectContentEventStream::new(ReceiverStream::new(rx)).into_byte_stream(); + let mut encoded = Vec::new(); + while let Some(chunk) = byte_stream.next().await { + encoded.extend_from_slice(&chunk.expect("event-stream message should serialize")); + } + let messages = event_stream_headers(&encoded); + assert_eq!( + messages + .iter() + .filter_map(|headers| { + headers + .iter() + .find_map(|(name, value)| (name == ":event-type").then_some(value.as_str())) + }) + .collect::>(), + ["Cont", "Records"] + ); + let terminal_headers = messages.last().expect("stream should contain a terminal error"); + assert!( + terminal_headers + .iter() + .any(|(name, value)| name == ":message-type" && value == "error") + ); + assert!( + terminal_headers + .iter() + .any(|(name, value)| name == ":error-code" && value == "CSVParsingError") + ); + assert!( + !messages + .iter() + .flatten() + .any(|(name, value)| { name == ":event-type" && matches!(value.as_str(), "Stats" | "End") }) + ); + assert!(lease_released.await.is_ok(), "error should release the snapshot lease"); + } + #[tokio::test(start_paused = true)] async fn encoder_error_uses_reserved_terminal_slot() { let values = ListArray::from_iter_primitive::([Some([Some(1)])]); @@ -1497,6 +1764,8 @@ mod tests { .expect("encoder failure should send one terminal error") .expect_err("terminal event should be an error"); assert_eq!(encoder_error.code(), &S3ErrorCode::InternalError); + assert_eq!(encoder_error.message(), Some("An internal error occurred.")); + assert!(encoder_error.source().is_none()); assert!(rx.recv().await.is_none()); assert!(lease_released.await.is_ok(), "encoder error should release the snapshot lease"); } @@ -1628,8 +1897,7 @@ mod tests { }; assert_eq!(error.code(), &S3ErrorCode::InternalError); assert_eq!(snapshot_fence.0.load(std::sync::atomic::Ordering::Relaxed), 2); - assert!(matches!(rx.recv().await, Some(Ok(SelectObjectContentEvent::Stats(_))))); - assert!(rx.try_recv().is_err(), "snapshot loss after Stats must not enqueue End"); + assert!(rx.try_recv().is_err(), "snapshot loss must not enqueue Stats or End"); } #[test] @@ -1658,6 +1926,27 @@ mod tests { ); } + #[test] + fn validate_rejects_unknown_csv_header_mode_before_streaming() { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV") + .file_header_info = Some(FileHeaderInfo::from_static("INVALID")); + + let error = validate_select_request(&HeaderMap::new(), &mut input).expect_err("unknown header mode must fail"); + + assert_eq!(error.code(), &S3ErrorCode::InvalidFileHeaderInfo); + assert_eq!(error.status_code(), Some(StatusCode::BAD_REQUEST)); + assert_eq!( + error.message(), + Some("The FileHeaderInfo value is not valid. Only NONE, USE, and IGNORE are supported.") + ); + } + #[test] fn validate_accepts_two_byte_csv_input_record_delimiter() { let mut input = base_input(); @@ -1980,26 +2269,8 @@ mod tests { } #[test] - fn map_store_error_not_found_to_no_such_key() { - let err = map_query_error_to_s3(QueryError::StoreError { - e: "ObjectStore NotFound: bucket/object.csv".to_string(), - }); - assert_eq!(err.code(), &S3ErrorCode::NoSuchKey); - } - - #[test] - fn map_store_error_bucket_not_found_to_no_such_bucket() { - let err = map_query_error_to_s3(QueryError::StoreError { - e: "bucket not found".to_string(), - }); - assert_eq!(err.code(), &S3ErrorCode::NoSuchBucket); - } - - #[test] - fn map_scan_range_store_error_to_invalid_request_parameter() { - let err = map_query_error_to_s3(QueryError::StoreError { - e: "ScanRange: Start after EOF".to_string(), - }); + fn map_typed_scan_range_error_to_invalid_request_parameter() { + let err = map_query_error_to_s3(SelectError::InvalidScanRange.into()); assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter); assert_eq!(err.message(), Some(INVALID_SCAN_RANGE_MESSAGE)); }