mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-13 16:46:55 +00:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3f994c59eb | |||
| 7c1d9dec8f | |||
| 4a8759239d |
+308
-15
@@ -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<T: StdError + 'static>(&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::<T>() {
|
||||
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<S3SelectPolicyError> for QueryError {
|
||||
fn from(value: S3SelectPolicyError) -> Self {
|
||||
fn classify_select_error_source(err: &(dyn StdError + 'static)) -> Option<SelectError> {
|
||||
if let Some(error) = err.downcast_ref::<SelectError>() {
|
||||
return Some(error.clone());
|
||||
}
|
||||
if let Some(error) = err.downcast_ref::<object_store::SelectObjectStoreError>() {
|
||||
return Some(error.select_error());
|
||||
}
|
||||
if let Some(error) = err.downcast_ref::<datafusion::object_store::Error>() {
|
||||
return match error {
|
||||
datafusion::object_store::Error::NotFound { source, .. } => Some(
|
||||
source
|
||||
.downcast_ref::<object_store::SelectObjectStoreError>()
|
||||
.map_or(SelectError::ObjectNotFound, object_store::SelectObjectStoreError::select_error),
|
||||
),
|
||||
_ => None,
|
||||
};
|
||||
}
|
||||
if let Some(error) = err.downcast_ref::<ParserError>() {
|
||||
return Some(SelectError::ParseSelectFailure {
|
||||
message: error.to_string(),
|
||||
});
|
||||
}
|
||||
if let Some(error) = err.downcast_ref::<ArrowError>() {
|
||||
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::<ParquetError>() {
|
||||
return match error {
|
||||
ParquetError::External(_) => None,
|
||||
_ => Some(SelectError::ParquetParsingError),
|
||||
};
|
||||
}
|
||||
if let Some(error) = err.downcast_ref::<SchemaError>() {
|
||||
return Some(match error {
|
||||
SchemaError::FieldNotFound { .. } => SelectError::EvaluatorBindingDoesNotExist,
|
||||
SchemaError::AmbiguousReference { .. }
|
||||
| SchemaError::DuplicateQualifiedField { .. }
|
||||
| SchemaError::DuplicateUnqualifiedField { .. } => SelectError::AmbiguousFieldName,
|
||||
});
|
||||
}
|
||||
if let Some(error) = err.downcast_ref::<DataFusionError>() {
|
||||
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<SelectError> 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]
|
||||
|
||||
@@ -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<dyn MemoryPool> =
|
||||
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());
|
||||
|
||||
@@ -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<Vec<u8>> {
|
||||
|
||||
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)]
|
||||
|
||||
@@ -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<T> {
|
||||
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::<S3SelectPolicyError>(),
|
||||
Some(S3SelectPolicyError::QueryTimeout { seconds: 300 })
|
||||
source.downcast_ref::<SelectError>(),
|
||||
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::<S3SelectPolicyError>(),
|
||||
Some(S3SelectPolicyError::QueryTimeout { seconds: 300 })
|
||||
source.downcast_ref::<SelectError>(),
|
||||
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::<S3SelectPolicyError>(),
|
||||
Some(S3SelectPolicyError::QueryTimeout { seconds: 1 })
|
||||
source.downcast_ref::<SelectError>(),
|
||||
Some(SelectError::QueryTimeout { seconds: 1 })
|
||||
));
|
||||
assert_eq!(admission.available_permits(), 1);
|
||||
assert!(output.next().await.is_none());
|
||||
|
||||
@@ -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<dyn ObjectStore>) -> 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<ParquetMetaData> {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
|
||||
@@ -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![
|
||||
|
||||
+482
-211
@@ -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<StorageSelectObjectSnapshot> {
|
||||
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<L: SelectSnapshotFence>(
|
||||
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<String>) -> 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<Vec<u8>
|
||||
}
|
||||
|
||||
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<Vec<u
|
||||
let mut writer = JsonWriterBuilder::new()
|
||||
.with_explicit_nulls(true)
|
||||
.build::<_, LineDelimited>(&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::<DataFusionError>()
|
||||
.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::<std::io::Error>()
|
||||
.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::<std::num::TryFromIntError>().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<Vec<(String, String)>> {
|
||||
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<tokio::sync::oneshot::Sender<()>>);
|
||||
|
||||
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::<StoragePrepareSelectObjectSnapshotError>().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::<Vec<_>>();
|
||||
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!("<Code>{expected_code_text}</Code>")));
|
||||
assert!(body.contains("<Message>"));
|
||||
|
||||
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("<Code>ParseSelectFailure</Code>"));
|
||||
assert!(sql_body.contains("<Message>"));
|
||||
|
||||
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("<Code>InvalidCompressionFormat</Code>"));
|
||||
assert!(compression_body.contains("<Message>"));
|
||||
}
|
||||
|
||||
#[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::<Vec<_>>(),
|
||||
["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::<Int32Type, _, _>([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));
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user