mirror of
https://github.com/rustfs/rustfs.git
synced 2026-07-26 08:18:18 +00:00
fix(s3select): enforce query and resource limits (#5028)
* fix(s3select): enforce query and resource limits * fix(s3select): close query resource limit gaps * fix(s3select): preserve timeout and stream invariants * fix(s3select): enforce staged query limits * fix(s3select): preserve policy error compatibility * fix(s3select): bound error source traversal
This commit is contained in:
GatewayJ
Generated
+1
@@ -9868,6 +9868,7 @@ dependencies = [
|
||||
"rustfs-test-utils",
|
||||
"s3s",
|
||||
"serde_json",
|
||||
"serial_test",
|
||||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
|
||||
@@ -49,6 +49,7 @@ url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rustfs-test-utils.workspace = true
|
||||
serial_test.workspace = true
|
||||
tempfile.workspace = true
|
||||
|
||||
[lib]
|
||||
|
||||
@@ -64,6 +64,48 @@ pub enum QueryError {
|
||||
StoreError { e: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
#[non_exhaustive]
|
||||
pub enum S3SelectPolicyError {
|
||||
#[error("Unsupported S3 Select SQL structure: {message}")]
|
||||
UnsupportedSqlStructure { message: String },
|
||||
|
||||
#[error("S3 Select query concurrency limit reached")]
|
||||
QueryConcurrencyLimit,
|
||||
|
||||
#[error("S3 Select query exceeded the {seconds}-second execution limit")]
|
||||
QueryTimeout { seconds: u64 },
|
||||
}
|
||||
|
||||
impl S3SelectPolicyError {
|
||||
fn from_error<'a>(mut err: &'a (dyn std::error::Error + 'static)) -> Option<&'a Self> {
|
||||
for _ in 0..16 {
|
||||
if let Some(policy_error) = err.downcast_ref::<Self>() {
|
||||
return Some(policy_error);
|
||||
}
|
||||
err = err.source()?;
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl QueryError {
|
||||
pub fn s3_select_policy_error(&self) -> Option<&S3SelectPolicyError> {
|
||||
match self {
|
||||
Self::Datafusion { source } => S3SelectPolicyError::from_error(source.as_ref()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<S3SelectPolicyError> for QueryError {
|
||||
fn from(value: S3SelectPolicyError) -> Self {
|
||||
Self::Datafusion {
|
||||
source: Box::new(DataFusionError::External(Box::new(value))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<DataFusionError> for QueryError {
|
||||
fn from(value: DataFusionError) -> Self {
|
||||
match value {
|
||||
@@ -116,9 +158,23 @@ mod tests {
|
||||
};
|
||||
assert_eq!(err.to_string(), "Multi-statement not allow, found num:2, sql:SELECT 1; SELECT 2;");
|
||||
|
||||
let err = S3SelectPolicyError::UnsupportedSqlStructure {
|
||||
message: "JOIN is not supported".to_string(),
|
||||
};
|
||||
assert_eq!(err.to_string(), "Unsupported S3 Select SQL structure: JOIN is not supported");
|
||||
|
||||
let err = QueryError::Cancel;
|
||||
assert_eq!(err.to_string(), "The query has been canceled");
|
||||
|
||||
assert_eq!(
|
||||
S3SelectPolicyError::QueryConcurrencyLimit.to_string(),
|
||||
"S3 Select query concurrency limit reached"
|
||||
);
|
||||
assert_eq!(
|
||||
S3SelectPolicyError::QueryTimeout { seconds: 300 }.to_string(),
|
||||
"S3 Select query exceeded the 300-second execution limit"
|
||||
);
|
||||
|
||||
let err = QueryError::FunctionNotExists {
|
||||
name: "my_func".to_string(),
|
||||
};
|
||||
@@ -143,6 +199,35 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_error_variants_remain_source_compatible() {
|
||||
fn exhaustive_match(err: QueryError) {
|
||||
match err {
|
||||
QueryError::Datafusion { .. }
|
||||
| QueryError::NotImplemented { .. }
|
||||
| QueryError::MultiStatement { .. }
|
||||
| QueryError::BuildQueryDispatcher { .. }
|
||||
| QueryError::Cancel
|
||||
| QueryError::Parser { .. }
|
||||
| QueryError::FunctionNotExists { .. }
|
||||
| QueryError::FunctionExists { .. }
|
||||
| QueryError::StoreError { .. } => {}
|
||||
}
|
||||
}
|
||||
|
||||
exhaustive_match(QueryError::Cancel);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_error_is_recoverable_from_query_error() {
|
||||
let err: QueryError = S3SelectPolicyError::QueryTimeout { seconds: 300 }.into();
|
||||
|
||||
assert!(matches!(
|
||||
err.s3_select_policy_error(),
|
||||
Some(S3SelectPolicyError::QueryTimeout { seconds: 300 })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_query_error_from_parser_error() {
|
||||
let parser_error = ParserError::ParserError("syntax error".to_string());
|
||||
|
||||
@@ -14,20 +14,27 @@
|
||||
|
||||
use crate::{
|
||||
SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectInfo, SelectObjectOptions, SelectStorageError,
|
||||
SelectStore, resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
|
||||
SelectStore,
|
||||
query::session::{QueryExecutionGuard, QueryExecutionTracker},
|
||||
resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
|
||||
select_is_err_version_not_found,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use chrono::Utc;
|
||||
use datafusion::object_store::{
|
||||
Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload,
|
||||
ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path,
|
||||
use datafusion::{
|
||||
common::{DataFusionError, runtime::SpawnedTask},
|
||||
execution::memory_pool::{MemoryConsumer, MemoryPool, UnboundedMemoryPool},
|
||||
object_store::{
|
||||
Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult,
|
||||
MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path,
|
||||
},
|
||||
};
|
||||
use futures::pin_mut;
|
||||
use futures::{Stream, StreamExt, future::ready, stream};
|
||||
use futures_core::stream::BoxStream;
|
||||
use http::{HeaderMap, HeaderValue, header::HeaderName};
|
||||
use parking_lot::Mutex;
|
||||
use rustfs_common::DEFAULT_DELIMITER;
|
||||
use s3s::S3Result;
|
||||
use s3s::dto::SelectObjectContentInput;
|
||||
@@ -49,6 +56,13 @@ fn select_default_read_buffer_size_u64() -> u64 {
|
||||
u64::try_from(SELECT_DEFAULT_READ_BUFFER_SIZE).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
||||
fn validated_object_size(size: i64) -> Result<u64> {
|
||||
u64::try_from(size).map_err(|err| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(err),
|
||||
})
|
||||
}
|
||||
|
||||
/// Maximum allowed object size for JSON DOCUMENT mode.
|
||||
///
|
||||
/// JSON DOCUMENT format requires loading the entire file into memory for DOM
|
||||
@@ -61,8 +75,13 @@ fn select_default_read_buffer_size_u64() -> u64 {
|
||||
/// size limit.
|
||||
///
|
||||
/// Default: 128 MiB. This matches the AWS S3 Select limit for JSON DOCUMENT
|
||||
/// inputs.
|
||||
/// inputs. The query memory pool also applies: RustFS reserves 64 times the
|
||||
/// input size for parsing and output. With the default 64 MiB query memory
|
||||
/// limit, JSON DOCUMENT inputs larger than 1 MiB are rejected; raise
|
||||
/// `RUSTFS_S3SELECT_MEMORY_LIMIT_BYTES` to process larger inputs, up to this
|
||||
/// hard cap.
|
||||
pub const MAX_JSON_DOCUMENT_BYTES: u64 = 128 * 1024 * 1024;
|
||||
const JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER: usize = 64;
|
||||
pub const INVALID_SCAN_RANGE_MESSAGE: &str =
|
||||
"The value of a parameter in ScanRange element is invalid. Check the service API documentation and try again.";
|
||||
|
||||
@@ -79,6 +98,8 @@ pub struct EcObjectStore {
|
||||
/// expression. When set, `flatten_json_document_to_ndjson` navigates to
|
||||
/// this key in the root JSON object before flattening.
|
||||
json_sub_path: Option<String>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
|
||||
store: Arc<SelectStore>,
|
||||
}
|
||||
@@ -108,7 +129,29 @@ pub struct InvalidScanRange;
|
||||
|
||||
impl EcObjectStore {
|
||||
pub fn new(input: Arc<SelectObjectContentInput>) -> S3Result<Self> {
|
||||
let Some(store) = resolve_select_object_store_handle() else {
|
||||
Self::build(input, Arc::new(UnboundedMemoryPool::default()), None, None)
|
||||
}
|
||||
|
||||
pub(crate) fn new_with_memory_pool(input: Arc<SelectObjectContentInput>, memory_pool: Arc<dyn MemoryPool>) -> S3Result<Self> {
|
||||
Self::build(input, memory_pool, None, None)
|
||||
}
|
||||
|
||||
pub(crate) fn new_with_query_tracker(
|
||||
input: Arc<SelectObjectContentInput>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
query_tracker: QueryExecutionTracker,
|
||||
store: Option<Arc<SelectStore>>,
|
||||
) -> S3Result<Self> {
|
||||
Self::build(input, memory_pool, Some(query_tracker), store)
|
||||
}
|
||||
|
||||
fn build(
|
||||
input: Arc<SelectObjectContentInput>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
store: Option<Arc<SelectStore>>,
|
||||
) -> S3Result<Self> {
|
||||
let Some(store) = store.or_else(resolve_select_object_store_handle) else {
|
||||
return Err(s3_error!(InternalError, "ec store not inited"));
|
||||
};
|
||||
|
||||
@@ -151,6 +194,8 @@ impl EcObjectStore {
|
||||
delimiter,
|
||||
is_json_document,
|
||||
json_sub_path,
|
||||
memory_pool,
|
||||
query_tracker,
|
||||
store,
|
||||
})
|
||||
}
|
||||
@@ -213,13 +258,30 @@ impl EcObjectStore {
|
||||
if range.is_empty() {
|
||||
return Ok(Bytes::new());
|
||||
}
|
||||
let reader = self.object_reader(Some(http_range_spec_from_range(range)), opts).await?;
|
||||
let reader = self
|
||||
.object_reader(Some(http_range_spec_from_range(range.clone())), opts)
|
||||
.await?;
|
||||
let object_size = validated_object_size(reader.object_info.size)?;
|
||||
let resolved_range = GetRange::Bounded(range)
|
||||
.as_range(object_size)
|
||||
.map_err(|err| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(err),
|
||||
})?;
|
||||
let expected_size = usize::try_from(resolved_range.end - resolved_range.start).map_err(|err| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(err),
|
||||
})?;
|
||||
let mut reader = reader.stream;
|
||||
let mut bytes = Vec::new();
|
||||
reader.read_to_end(&mut bytes).await.map_err(|err| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(err),
|
||||
})?;
|
||||
if bytes.len() < expected_size {
|
||||
return Err(incomplete_object_stream_error(expected_size - bytes.len()));
|
||||
}
|
||||
bytes.truncate(expected_size);
|
||||
Ok(Bytes::from(bytes))
|
||||
}
|
||||
|
||||
@@ -421,7 +483,7 @@ impl ObjectStore for EcObjectStore {
|
||||
let opts = self.object_options(&options);
|
||||
let needs_scan_context = options.range.is_none() && !options.head && self.input.request.scan_range.is_some();
|
||||
let source_size = if needs_scan_context {
|
||||
Some(self.object_info(&opts).await?.size as u64)
|
||||
Some(validated_object_size(self.object_info(&opts).await?.size)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -442,7 +504,10 @@ impl ObjectStore for EcObjectStore {
|
||||
self.object_reader(range, &opts).await?
|
||||
};
|
||||
|
||||
let original_size = source_size.unwrap_or(reader.object_info.size as u64);
|
||||
let original_size = match source_size {
|
||||
Some(source_size) => source_size,
|
||||
None => validated_object_size(reader.object_info.size)?,
|
||||
};
|
||||
let etag = reader.object_info.etag;
|
||||
let version = reader.object_info.version_id.map(|version| version.to_string());
|
||||
let attributes = Attributes::default();
|
||||
@@ -473,20 +538,14 @@ impl ObjectStore for EcObjectStore {
|
||||
// which must load the whole file into memory; rejecting oversized
|
||||
// files upfront is safer than risking OOM. Users should convert
|
||||
// their data to JSON LINES (NDJSON) format for large files.
|
||||
if original_size > MAX_JSON_DOCUMENT_BYTES {
|
||||
return Err(o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: format!(
|
||||
"JSON DOCUMENT object is {original_size} bytes, which exceeds the \
|
||||
maximum allowed size of {MAX_JSON_DOCUMENT_BYTES} bytes \
|
||||
({} MiB). Convert the input to JSON LINES (NDJSON) to process \
|
||||
large files.",
|
||||
MAX_JSON_DOCUMENT_BYTES / (1024 * 1024)
|
||||
)
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
let stream = json_document_ndjson_stream(reader.stream, original_size, self.json_sub_path.clone());
|
||||
validate_json_document_size(original_size)?;
|
||||
let stream = json_document_ndjson_stream(
|
||||
reader.stream,
|
||||
original_size,
|
||||
self.json_sub_path.clone(),
|
||||
Arc::clone(&self.memory_pool),
|
||||
self.query_tracker.clone(),
|
||||
);
|
||||
GetResultPayload::Stream(stream)
|
||||
} else if let Some((_, scan_range)) = scan_context {
|
||||
let delimiter = self.record_delimiter();
|
||||
@@ -503,6 +562,7 @@ impl ObjectStore for EcObjectStore {
|
||||
scan_range,
|
||||
include_header && header.is_none(),
|
||||
read_start,
|
||||
original_size,
|
||||
)
|
||||
.boxed();
|
||||
let stream = if let Some(header) = header {
|
||||
@@ -664,6 +724,7 @@ struct ScanRangeState<S> {
|
||||
record: Vec<u8>,
|
||||
pending: VecDeque<Bytes>,
|
||||
done: bool,
|
||||
expected_end: u64,
|
||||
}
|
||||
|
||||
fn scan_range_stream<S>(
|
||||
@@ -672,6 +733,7 @@ fn scan_range_stream<S>(
|
||||
range: SelectScanRange,
|
||||
include_header: bool,
|
||||
base_offset: u64,
|
||||
expected_end: u64,
|
||||
) -> BoxStream<'static, Result<Bytes>>
|
||||
where
|
||||
S: Stream<Item = std::io::Result<Bytes>> + Send + Unpin + 'static,
|
||||
@@ -686,6 +748,7 @@ where
|
||||
record: Vec::new(),
|
||||
pending: VecDeque::new(),
|
||||
done: false,
|
||||
expected_end,
|
||||
};
|
||||
|
||||
stream::unfold(state, |mut state| async move {
|
||||
@@ -709,6 +772,10 @@ where
|
||||
));
|
||||
}
|
||||
None => {
|
||||
if state.offset < state.expected_end {
|
||||
state.done = true;
|
||||
return Some((Err(incomplete_object_stream_error(state.expected_end - state.offset)), state));
|
||||
}
|
||||
state.finish_pending_record();
|
||||
state.done = true;
|
||||
}
|
||||
@@ -855,11 +922,57 @@ fn json_document_ndjson_stream(
|
||||
stream: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
|
||||
original_size: u64,
|
||||
json_sub_path: Option<String>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
) -> futures_core::stream::BoxStream<'static, Result<Bytes>> {
|
||||
json_document_ndjson_stream_with_parser(
|
||||
stream,
|
||||
original_size,
|
||||
json_sub_path,
|
||||
memory_pool,
|
||||
query_tracker,
|
||||
|all_bytes, json_sub_path| parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()),
|
||||
)
|
||||
}
|
||||
|
||||
fn json_document_ndjson_stream_with_parser<P>(
|
||||
stream: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
|
||||
original_size: u64,
|
||||
json_sub_path: Option<String>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
parser: P,
|
||||
) -> futures_core::stream::BoxStream<'static, Result<Bytes>>
|
||||
where
|
||||
P: FnOnce(Vec<u8>, Option<String>) -> std::io::Result<Vec<Bytes>> + Send + 'static,
|
||||
{
|
||||
AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
|
||||
// Compact JSON can expand substantially into a serde_json DOM and
|
||||
// per-record output buffers, so reserve a conservative upper bound
|
||||
// before the source buffer is allocated.
|
||||
let buffer_capacity = usize::try_from(original_size).map_err(|_| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(DataFusionError::ResourcesExhausted(format!(
|
||||
"JSON DOCUMENT input size {original_size} does not fit in memory"
|
||||
))),
|
||||
})?;
|
||||
let reservation_bytes = buffer_capacity
|
||||
.checked_mul(JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)
|
||||
.ok_or_else(|| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(DataFusionError::ResourcesExhausted(format!(
|
||||
"JSON DOCUMENT memory reservation overflow for {original_size} input bytes"
|
||||
))),
|
||||
})?;
|
||||
let reservation = MemoryConsumer::new("S3 Select JSON document").register(&memory_pool);
|
||||
reservation.try_resize(reservation_bytes).map_err(|err| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(err),
|
||||
})?;
|
||||
|
||||
pin_mut!(stream);
|
||||
// ── 1. Read phase (lazy: only runs when the stream is polled) ────
|
||||
let mut all_bytes = Vec::with_capacity(original_size as usize);
|
||||
let mut all_bytes = Vec::with_capacity(buffer_capacity);
|
||||
stream
|
||||
.take(original_size)
|
||||
.read_to_end(&mut all_bytes)
|
||||
@@ -868,18 +981,26 @@ fn json_document_ndjson_stream(
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(e),
|
||||
})?;
|
||||
if all_bytes.len() != buffer_capacity {
|
||||
return Err(incomplete_object_stream_error(buffer_capacity - all_bytes.len()));
|
||||
}
|
||||
|
||||
// ── 2. Parse phase (blocking thread pool, non-blocking runtime) ──
|
||||
let lines = tokio::task::spawn_blocking(move || parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()))
|
||||
.await
|
||||
.map_err(|e| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: e.to_string().into(),
|
||||
})?
|
||||
.map_err(|e| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(e),
|
||||
})?;
|
||||
let pending_query_guard = PendingQueryExecutionGuard::new(query_tracker);
|
||||
let task_query_guard = pending_query_guard.task_state();
|
||||
let (lines, _reservation, _query_guard) = SpawnedTask::spawn_blocking(move || {
|
||||
let query_guard = PendingQueryExecutionGuard::start(&task_query_guard)?;
|
||||
parser(all_bytes, json_sub_path).map(|lines| (lines, reservation, query_guard))
|
||||
})
|
||||
.await
|
||||
.map_err(|e| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: e.to_string().into(),
|
||||
})?
|
||||
.map_err(|e| o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(e),
|
||||
})?;
|
||||
|
||||
// ── 3. Yield phase (one Bytes per NDJSON line) ───────────────────
|
||||
for line in lines {
|
||||
@@ -890,6 +1011,66 @@ fn json_document_ndjson_stream(
|
||||
.boxed()
|
||||
}
|
||||
|
||||
struct PendingQueryExecutionGuard {
|
||||
state: Arc<Mutex<QueryExecutionGuardState>>,
|
||||
}
|
||||
|
||||
enum QueryExecutionGuardState {
|
||||
Pending(Option<QueryExecutionTracker>),
|
||||
Started,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl PendingQueryExecutionGuard {
|
||||
fn new(query_tracker: Option<QueryExecutionTracker>) -> Self {
|
||||
Self {
|
||||
state: Arc::new(Mutex::new(QueryExecutionGuardState::Pending(query_tracker))),
|
||||
}
|
||||
}
|
||||
|
||||
fn task_state(&self) -> Arc<Mutex<QueryExecutionGuardState>> {
|
||||
Arc::clone(&self.state)
|
||||
}
|
||||
|
||||
fn start(state: &Mutex<QueryExecutionGuardState>) -> std::io::Result<Option<QueryExecutionGuard>> {
|
||||
let mut state = state.lock();
|
||||
match std::mem::replace(&mut *state, QueryExecutionGuardState::Started) {
|
||||
QueryExecutionGuardState::Pending(None) => Ok(None),
|
||||
QueryExecutionGuardState::Pending(Some(query_tracker)) => query_tracker.query_guard().map(Some).ok_or_else(|| {
|
||||
std::io::Error::new(std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started")
|
||||
}),
|
||||
QueryExecutionGuardState::Cancelled => {
|
||||
*state = QueryExecutionGuardState::Cancelled;
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::Interrupted,
|
||||
"JSON DOCUMENT parse was cancelled before it started",
|
||||
))
|
||||
}
|
||||
QueryExecutionGuardState::Started => {
|
||||
*state = QueryExecutionGuardState::Started;
|
||||
Err(std::io::Error::other("JSON DOCUMENT parse started more than once"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PendingQueryExecutionGuard {
|
||||
fn drop(&mut self) {
|
||||
let query_guard = {
|
||||
let mut state = self.state.lock();
|
||||
match std::mem::replace(&mut *state, QueryExecutionGuardState::Cancelled) {
|
||||
QueryExecutionGuardState::Pending(query_guard) => query_guard,
|
||||
QueryExecutionGuardState::Started => {
|
||||
*state = QueryExecutionGuardState::Started;
|
||||
None
|
||||
}
|
||||
QueryExecutionGuardState::Cancelled => None,
|
||||
}
|
||||
};
|
||||
drop(query_guard);
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a JSON DOCUMENT (a single JSON value, possibly multi-line) into a
|
||||
/// list of NDJSON lines – one [`Bytes`] per record.
|
||||
///
|
||||
@@ -907,19 +1088,11 @@ fn parse_json_document_to_lines(bytes: &[u8], json_sub_path: Option<&str>) -> st
|
||||
|
||||
// Navigate into the sub-path when the root is an object and a path was
|
||||
// extracted from the SQL FROM clause (e.g. `FROM s3object.employees`).
|
||||
let value = if let Some(path) = json_sub_path {
|
||||
if let serde_json::Value::Object(ref obj) = root {
|
||||
match obj.get(path) {
|
||||
Some(sub) => sub.clone(),
|
||||
// Path not found – fall back to emitting the whole root object.
|
||||
None => root,
|
||||
}
|
||||
} else {
|
||||
// Root is already an array or scalar; ignore the path hint.
|
||||
root
|
||||
let value = match (root, json_sub_path) {
|
||||
(serde_json::Value::Object(mut object), Some(path)) => {
|
||||
object.remove(path).unwrap_or_else(|| serde_json::Value::Object(object))
|
||||
}
|
||||
} else {
|
||||
root
|
||||
(root, _) => root,
|
||||
};
|
||||
|
||||
let mut lines: Vec<Bytes> = Vec::new();
|
||||
@@ -979,30 +1152,55 @@ where
|
||||
y.yield_ok(bytes).await;
|
||||
}
|
||||
if remaining > 0 {
|
||||
return Err(o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
format!("object stream ended with {remaining} bytes remaining"),
|
||||
)),
|
||||
});
|
||||
return Err(incomplete_object_stream_error(remaining));
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_json_document_size(original_size: u64) -> Result<()> {
|
||||
if original_size <= MAX_JSON_DOCUMENT_BYTES {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
Err(o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(DataFusionError::ResourcesExhausted(format!(
|
||||
"JSON DOCUMENT object is {original_size} bytes, which exceeds the maximum allowed size of \
|
||||
{MAX_JSON_DOCUMENT_BYTES} bytes ({} MiB). Convert the input to JSON LINES (NDJSON) to process large files.",
|
||||
MAX_JSON_DOCUMENT_BYTES / (1024 * 1024)
|
||||
))),
|
||||
})
|
||||
}
|
||||
|
||||
fn incomplete_object_stream_error(remaining: impl std::fmt::Display) -> o_Error {
|
||||
o_Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
format!("object stream ended with {remaining} bytes remaining"),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::{
|
||||
SelectScanRange, bytes_stream, convert_field_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter,
|
||||
flatten_json_document_to_ndjson, http_range_spec_from_get_range, replace_symbol, scan_range_from_bounds,
|
||||
scan_range_read_start, scan_range_stream, select_read_headers,
|
||||
EcObjectStore, JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER, SelectScanRange, bytes_stream,
|
||||
convert_field_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, replace_symbol,
|
||||
scan_range_from_bounds, scan_range_read_start, scan_range_stream, select_read_headers, validate_json_document_size,
|
||||
validated_object_size,
|
||||
};
|
||||
use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker};
|
||||
use crate::storage_api::SelectPutObjReader;
|
||||
use crate::storage_api::object_store::ObjectIO as _;
|
||||
use bytes::Bytes;
|
||||
use datafusion::object_store::{self, GetRange};
|
||||
use datafusion::object_store::{GetOptions, GetResultPayload, ObjectStore as _, path::Path};
|
||||
use datafusion::{
|
||||
common::DataFusionError,
|
||||
execution::memory_pool::{GreedyMemoryPool, MemoryPool},
|
||||
object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path},
|
||||
};
|
||||
use futures::{StreamExt, TryStreamExt, stream};
|
||||
use s3s::dto::{
|
||||
CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput,
|
||||
@@ -1017,16 +1215,28 @@ mod test {
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn ec_object_store_constructor_remains_source_compatible() {
|
||||
let _constructor: fn(Arc<SelectObjectContentInput>) -> s3s::S3Result<EcObjectStore> = EcObjectStore::new;
|
||||
}
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
#[test]
|
||||
fn test_replace() {
|
||||
let result = replace_symbol(b"&&", b"dandan&&is&&best");
|
||||
assert_eq!(result, b"dandan,is,best");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validated_object_size_rejects_negative_metadata() {
|
||||
assert_eq!(validated_object_size(0).expect("zero object size should be valid"), 0);
|
||||
assert!(validated_object_size(-1).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_scan_range_stream_keeps_header_and_selected_record() {
|
||||
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0, 18);
|
||||
let mut output = Vec::new();
|
||||
while let Some(bytes) = stream.next().await {
|
||||
output.extend_from_slice(&bytes.unwrap());
|
||||
@@ -1037,7 +1247,7 @@ mod test {
|
||||
#[tokio::test]
|
||||
async fn test_scan_range_stream_skips_record_when_start_is_in_middle() {
|
||||
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0, 12);
|
||||
let mut output = Vec::new();
|
||||
while let Some(bytes) = stream.next().await {
|
||||
output.extend_from_slice(&bytes.unwrap());
|
||||
@@ -1048,7 +1258,7 @@ mod test {
|
||||
#[tokio::test]
|
||||
async fn test_scan_range_stream_keeps_record_when_end_is_in_middle() {
|
||||
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0, 12);
|
||||
let mut output = Vec::new();
|
||||
while let Some(bytes) = stream.next().await {
|
||||
output.extend_from_slice(&bytes.unwrap());
|
||||
@@ -1059,7 +1269,7 @@ mod test {
|
||||
#[tokio::test]
|
||||
async fn test_scan_range_stream_uses_base_offset_for_range_reader() {
|
||||
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"\n2,b\n3,c\n"))]);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3);
|
||||
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3, 12);
|
||||
let mut output = Vec::new();
|
||||
while let Some(bytes) = stream.next().await {
|
||||
output.extend_from_slice(&bytes.unwrap());
|
||||
@@ -1074,7 +1284,7 @@ mod test {
|
||||
Ok::<_, std::io::Error>(Bytes::from_static(b"\n1,a\r\n2,b\r")),
|
||||
Ok::<_, std::io::Error>(Bytes::from_static(b"\n3,c\r\n")),
|
||||
]);
|
||||
let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0);
|
||||
let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0, 22);
|
||||
let mut output = Vec::new();
|
||||
while let Some(bytes) = stream.next().await {
|
||||
output.extend_from_slice(&bytes.unwrap());
|
||||
@@ -1082,6 +1292,26 @@ mod test {
|
||||
assert_eq!(output, b"h1,h2\r\n2,b\r\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_scan_range_stream_rejects_early_eof() {
|
||||
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n"))]);
|
||||
let mut output = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 7), false, 0, 8);
|
||||
|
||||
assert_eq!(output.next().await.expect("first stream item").expect("first record"), b"1,a\n"[..]);
|
||||
let err = output
|
||||
.next()
|
||||
.await
|
||||
.expect("early EOF error")
|
||||
.expect_err("short ScanRange stream must fail");
|
||||
let object_store::Error::Generic { source, .. } = err else {
|
||||
panic!("expected generic object store error");
|
||||
};
|
||||
let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
|
||||
assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
|
||||
assert!(source.to_string().contains("4 bytes remaining"));
|
||||
assert!(output.next().await.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_range_read_start_keeps_full_delimiter_boundary() {
|
||||
let range = SelectScanRange::new(10, 20);
|
||||
@@ -1157,7 +1387,7 @@ mod test {
|
||||
#[tokio::test]
|
||||
async fn test_scan_range_output_can_convert_field_delimiter() {
|
||||
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"a&&1\nb&&2\n"))]);
|
||||
let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0);
|
||||
let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0, 10);
|
||||
let mut stream = convert_field_delimiter_stream(stream, Some("&&".to_string()));
|
||||
let mut output = Vec::new();
|
||||
while let Some(bytes) = stream.next().await {
|
||||
@@ -1195,6 +1425,7 @@ mod test {
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial_test::serial]
|
||||
async fn test_get_opts_validates_raw_length_before_delimiter_conversion() {
|
||||
let temp_root = tempfile::tempdir().expect("create s3select test temp root");
|
||||
let env = rustfs_test_utils::TestECStoreEnv::builder()
|
||||
@@ -1242,6 +1473,8 @@ mod test {
|
||||
delimiter: "&&".to_string(),
|
||||
is_json_document: false,
|
||||
json_sub_path: None,
|
||||
memory_pool: Arc::new(GreedyMemoryPool::new(1024)),
|
||||
query_tracker: None,
|
||||
store: env.ecstore,
|
||||
};
|
||||
|
||||
@@ -1255,6 +1488,13 @@ mod test {
|
||||
let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect converted object stream");
|
||||
|
||||
assert_eq!(chunks.concat(), b"a,1\n");
|
||||
|
||||
let requested_range = 3..10;
|
||||
let ranges = store
|
||||
.get_ranges(&Path::from(object), std::slice::from_ref(&requested_range))
|
||||
.await
|
||||
.expect("bounded range past EOF should return the object remainder");
|
||||
assert_eq!(ranges, vec![Bytes::from_static(b"1\n")]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1302,6 +1542,278 @@ mod test {
|
||||
assert!(output.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_document_stream_respects_query_memory_pool() {
|
||||
let input = b"{}".to_vec();
|
||||
let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
|
||||
let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(required - 1));
|
||||
let mut output = json_document_ndjson_stream(
|
||||
Box::new(std::io::Cursor::new(input.clone())),
|
||||
input.len() as u64,
|
||||
None,
|
||||
memory_pool,
|
||||
None,
|
||||
);
|
||||
|
||||
let err = output
|
||||
.next()
|
||||
.await
|
||||
.expect("memory error")
|
||||
.expect_err("reservation should exceed the pool");
|
||||
let object_store::Error::Generic { source, .. } = err else {
|
||||
panic!("expected generic object store error");
|
||||
};
|
||||
assert!(matches!(
|
||||
source.downcast_ref::<DataFusionError>(),
|
||||
Some(DataFusionError::ResourcesExhausted(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_document_stream_releases_memory_reservation() {
|
||||
let input = b"[1,2]".to_vec();
|
||||
let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
|
||||
let memory_pool = Arc::new(GreedyMemoryPool::new(required));
|
||||
let output: Vec<Bytes> = json_document_ndjson_stream(
|
||||
Box::new(std::io::Cursor::new(input.clone())),
|
||||
input.len() as u64,
|
||||
None,
|
||||
memory_pool.clone(),
|
||||
None,
|
||||
)
|
||||
.try_collect()
|
||||
.await
|
||||
.expect("JSON conversion should fit the pool");
|
||||
|
||||
assert_eq!(output, vec![Bytes::from_static(b"1\n"), Bytes::from_static(b"2\n")]);
|
||||
assert_eq!(memory_pool.reserved(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_document_stream_rejects_early_eof() {
|
||||
let input = b"{}".to_vec();
|
||||
let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
|
||||
let mut output = json_document_ndjson_stream(Box::new(std::io::Cursor::new(input)), 4, None, memory_pool, None);
|
||||
|
||||
let err = output
|
||||
.next()
|
||||
.await
|
||||
.expect("early EOF error")
|
||||
.expect_err("short JSON document must fail");
|
||||
let object_store::Error::Generic { source, .. } = err else {
|
||||
panic!("expected generic object store error");
|
||||
};
|
||||
let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
|
||||
assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
|
||||
assert!(source.to_string().contains("2 bytes remaining"));
|
||||
assert!(output.next().await.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_document_size_error_is_resource_exhausted() {
|
||||
assert!(validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES).is_ok());
|
||||
|
||||
let err = validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES + 1).expect_err("oversized JSON document must fail");
|
||||
let object_store::Error::Generic { source, .. } = err else {
|
||||
panic!("expected generic object store error");
|
||||
};
|
||||
assert!(matches!(
|
||||
source.downcast_ref::<DataFusionError>(),
|
||||
Some(DataFusionError::ResourcesExhausted(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_document_queued_parse_releases_query_guard_when_cancelled() {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(2)
|
||||
.max_blocking_threads(1)
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("build test runtime");
|
||||
|
||||
runtime.block_on(async {
|
||||
let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
|
||||
let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
|
||||
let blocker = tokio::task::spawn_blocking(move || {
|
||||
let _ = blocking_started_tx.send(());
|
||||
release_blocking_rx.recv().expect("release blocking worker");
|
||||
});
|
||||
blocking_started_rx.await.expect("blocking worker should start");
|
||||
|
||||
let admission = Arc::new(Semaphore::new(1));
|
||||
let permit = Arc::clone(&admission)
|
||||
.acquire_owned()
|
||||
.await
|
||||
.expect("query permit should be available");
|
||||
let query_guard: QueryExecutionGuard = Arc::new(permit);
|
||||
let query_tracker = QueryExecutionTracker::new(
|
||||
&QueryExecutionOwner::new(),
|
||||
query_guard,
|
||||
tokio::time::Instant::now() + std::time::Duration::from_secs(30),
|
||||
30,
|
||||
);
|
||||
let input = b"{}".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,
|
||||
Some(query_tracker),
|
||||
);
|
||||
|
||||
{
|
||||
let next = output.next();
|
||||
futures::pin_mut!(next);
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
}
|
||||
drop(output);
|
||||
|
||||
let recovered_permit =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
|
||||
.await
|
||||
.expect("queued JSON parse should be cancelled")
|
||||
.expect("query admission should remain open");
|
||||
release_blocking_tx.send(()).expect("release blocking worker");
|
||||
blocker.await.expect("blocking worker should finish");
|
||||
drop(recovered_permit);
|
||||
assert_eq!(admission.available_permits(), 1);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_document_expired_queued_parse_does_not_start() {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(2)
|
||||
.max_blocking_threads(1)
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("build test runtime");
|
||||
|
||||
runtime.block_on(async {
|
||||
let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
|
||||
let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
|
||||
let blocker = tokio::task::spawn_blocking(move || {
|
||||
let _ = blocking_started_tx.send(());
|
||||
release_blocking_rx.recv().expect("release blocking worker");
|
||||
});
|
||||
blocking_started_rx.await.expect("blocking worker should start");
|
||||
|
||||
let admission = Arc::new(Semaphore::new(1));
|
||||
let permit = Arc::clone(&admission)
|
||||
.acquire_owned()
|
||||
.await
|
||||
.expect("query permit should be available");
|
||||
let owner = QueryExecutionOwner::new();
|
||||
let query_tracker = QueryExecutionTracker::new(
|
||||
&owner,
|
||||
Arc::new(permit),
|
||||
tokio::time::Instant::now() + std::time::Duration::from_secs(30),
|
||||
30,
|
||||
);
|
||||
let input = b"{}".to_vec();
|
||||
let memory_pool: Arc<dyn MemoryPool> =
|
||||
Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
|
||||
let parser_started = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let parser_started_in_task = Arc::clone(&parser_started);
|
||||
let mut output = json_document_ndjson_stream_with_parser(
|
||||
Box::new(std::io::Cursor::new(input.clone())),
|
||||
input.len() as u64,
|
||||
None,
|
||||
memory_pool,
|
||||
Some(query_tracker.clone()),
|
||||
move |_, _| {
|
||||
parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
Ok(vec![Bytes::from_static(b"{}\n")])
|
||||
},
|
||||
);
|
||||
|
||||
{
|
||||
let next = output.next();
|
||||
futures::pin_mut!(next);
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
}
|
||||
query_tracker.expire(&owner);
|
||||
assert_eq!(admission.available_permits(), 1);
|
||||
release_blocking_tx.send(()).expect("release blocking worker");
|
||||
blocker.await.expect("blocking worker should finish");
|
||||
|
||||
let err = tokio::time::timeout(std::time::Duration::from_secs(5), output.next())
|
||||
.await
|
||||
.expect("queued parser should resume")
|
||||
.expect("queued parser should return an error")
|
||||
.expect_err("expired queued parser must not run");
|
||||
let object_store::Error::Generic { source, .. } = err else {
|
||||
panic!("expected generic object store error");
|
||||
};
|
||||
let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
|
||||
assert_eq!(source.kind(), std::io::ErrorKind::Interrupted);
|
||||
assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_document_started_parse_retains_query_guard_when_cancelled() {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(2)
|
||||
.max_blocking_threads(1)
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("build test runtime");
|
||||
|
||||
runtime.block_on(async {
|
||||
let admission = Arc::new(Semaphore::new(1));
|
||||
let permit = Arc::clone(&admission)
|
||||
.acquire_owned()
|
||||
.await
|
||||
.expect("query permit should be available");
|
||||
let query_guard: QueryExecutionGuard = Arc::new(permit);
|
||||
let query_tracker = QueryExecutionTracker::new(
|
||||
&QueryExecutionOwner::new(),
|
||||
query_guard,
|
||||
tokio::time::Instant::now() + std::time::Duration::from_secs(30),
|
||||
30,
|
||||
);
|
||||
let input = b"{}".to_vec();
|
||||
let memory_pool: Arc<dyn MemoryPool> =
|
||||
Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
|
||||
let (parse_started_tx, parse_started_rx) = tokio::sync::oneshot::channel();
|
||||
let (release_parse_tx, release_parse_rx) = std::sync::mpsc::channel();
|
||||
let mut output = json_document_ndjson_stream_with_parser(
|
||||
Box::new(std::io::Cursor::new(input.clone())),
|
||||
input.len() as u64,
|
||||
None,
|
||||
memory_pool,
|
||||
Some(query_tracker),
|
||||
move |_, _| {
|
||||
let _ = parse_started_tx.send(());
|
||||
release_parse_rx.recv().expect("release JSON parser");
|
||||
Ok(vec![Bytes::from_static(b"{}\n")])
|
||||
},
|
||||
);
|
||||
|
||||
{
|
||||
let next = output.next();
|
||||
futures::pin_mut!(next);
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
}
|
||||
parse_started_rx.await.expect("JSON parser should start");
|
||||
drop(output);
|
||||
|
||||
assert!(Arc::clone(&admission).try_acquire_owned().is_err());
|
||||
release_parse_tx.send(()).expect("release JSON parser");
|
||||
let recovered_permit =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
|
||||
.await
|
||||
.expect("started JSON parse should release the query guard")
|
||||
.expect("query admission should remain open");
|
||||
drop(recovered_permit);
|
||||
assert_eq!(admission.available_permits(), 1);
|
||||
});
|
||||
}
|
||||
|
||||
/// A JSON array is split into one NDJSON line per element.
|
||||
#[test]
|
||||
fn test_flatten_array_produces_one_line_per_element() {
|
||||
|
||||
@@ -30,7 +30,7 @@ use crate::{QueryError, QueryResult};
|
||||
|
||||
use super::Query;
|
||||
use super::logical_planner::Plan;
|
||||
use super::session::SessionCtx;
|
||||
use super::session::{QueryExecutionTracker, SessionCtx};
|
||||
|
||||
pub struct PhaseTimer {
|
||||
phase_name: &'static str,
|
||||
@@ -172,6 +172,7 @@ pub struct QueryStateMachine {
|
||||
pub session: SessionCtx,
|
||||
pub query: Query,
|
||||
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
state: RwLock<QueryState>,
|
||||
start: Instant,
|
||||
}
|
||||
@@ -195,11 +196,31 @@ impl QueryStateMachine {
|
||||
Self {
|
||||
session,
|
||||
query,
|
||||
query_tracker: None,
|
||||
state: RwLock::new(QueryState::ACCEPTING),
|
||||
start: Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn begin_tracked(query: Query, session: SessionCtx, query_tracker: QueryExecutionTracker) -> QueryResult<Self> {
|
||||
if !session.is_bound_to(&query_tracker) {
|
||||
return Err(QueryError::Cancel);
|
||||
}
|
||||
let mut state_machine = Self::begin(query, session);
|
||||
state_machine.query_tracker = Some(query_tracker);
|
||||
Ok(state_machine)
|
||||
}
|
||||
|
||||
pub fn query_tracker(&self) -> Option<&QueryExecutionTracker> {
|
||||
self.query_tracker.as_ref()
|
||||
}
|
||||
|
||||
pub fn tracker_matches_session(&self) -> bool {
|
||||
self.query_tracker
|
||||
.as_ref()
|
||||
.is_some_and(|query_tracker| self.session.is_bound_to(query_tracker))
|
||||
}
|
||||
|
||||
pub fn begin_analyze(&self) {
|
||||
self.record_phase_timestamp("analyze", "start");
|
||||
self.translate_to(QueryState::RUNNING(RUNNING::ANALYZING));
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
use crate::query::Context;
|
||||
use crate::{QueryError, QueryResult, object_store::EcObjectStore};
|
||||
use crate::{QueryError, QueryResult, SelectStore, object_store::EcObjectStore};
|
||||
use datafusion::{
|
||||
arrow::{
|
||||
array::{Int32Array, StringArray},
|
||||
@@ -25,19 +25,237 @@ use datafusion::{
|
||||
parquet::arrow::ArrowWriter,
|
||||
prelude::SessionContext,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use parking_lot::Mutex;
|
||||
use std::sync::{
|
||||
Arc, Weak,
|
||||
atomic::{AtomicU8, Ordering},
|
||||
};
|
||||
use tokio::{
|
||||
sync::OwnedSemaphorePermit,
|
||||
task::AbortHandle,
|
||||
time::{Instant, sleep_until},
|
||||
};
|
||||
use tracing::error;
|
||||
|
||||
pub type QueryExecutionGuard = Arc<OwnedSemaphorePermit>;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct QueryExecutionOwner {
|
||||
identity: Arc<()>,
|
||||
}
|
||||
|
||||
impl QueryExecutionOwner {
|
||||
pub fn new() -> Self {
|
||||
Self { identity: Arc::new(()) }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QueryExecutionTracker {
|
||||
inner: Arc<QueryExecutionTrackerInner>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum QueryExecutionStatus {
|
||||
Active,
|
||||
Finished,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
const EXECUTION_SETTING_UP: u8 = 0;
|
||||
const EXECUTION_ADMITTED: u8 = 1;
|
||||
const EXECUTION_PLANNING: u8 = 2;
|
||||
const EXECUTION_PLANNED: u8 = 3;
|
||||
const EXECUTION_STARTING: u8 = 4;
|
||||
const EXECUTION_RUNNING: u8 = 5;
|
||||
const EXECUTION_FINISHED: u8 = 6;
|
||||
const EXECUTION_TIMED_OUT: u8 = 7;
|
||||
|
||||
struct QueryExecutionTrackerInner {
|
||||
owner_identity: Arc<()>,
|
||||
query_guard: Mutex<Option<QueryExecutionGuard>>,
|
||||
deadline: Instant,
|
||||
timeout_seconds: u64,
|
||||
state: AtomicU8,
|
||||
deadline_task: Mutex<Option<AbortHandle>>,
|
||||
}
|
||||
|
||||
impl QueryExecutionTracker {
|
||||
pub fn new(owner: &QueryExecutionOwner, query_guard: QueryExecutionGuard, deadline: Instant, timeout_seconds: u64) -> Self {
|
||||
let inner = Arc::new(QueryExecutionTrackerInner {
|
||||
owner_identity: Arc::clone(&owner.identity),
|
||||
query_guard: Mutex::new(Some(query_guard)),
|
||||
deadline,
|
||||
timeout_seconds,
|
||||
state: AtomicU8::new(EXECUTION_SETTING_UP),
|
||||
deadline_task: Mutex::new(None),
|
||||
});
|
||||
let deadline_inner = Arc::downgrade(&inner);
|
||||
let deadline_task = tokio::spawn(async move {
|
||||
sleep_until(deadline).await;
|
||||
if let Some(inner) = Weak::upgrade(&deadline_inner) {
|
||||
inner.expire_at_deadline();
|
||||
}
|
||||
});
|
||||
*inner.deadline_task.lock() = Some(deadline_task.abort_handle());
|
||||
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
pub fn deadline(&self) -> Instant {
|
||||
self.inner.deadline
|
||||
}
|
||||
|
||||
pub fn timeout_seconds(&self) -> u64 {
|
||||
self.inner.timeout_seconds
|
||||
}
|
||||
|
||||
pub fn status(&self) -> QueryExecutionStatus {
|
||||
match self.inner.state.load(Ordering::Acquire) {
|
||||
EXECUTION_TIMED_OUT => QueryExecutionStatus::TimedOut,
|
||||
state if state < EXECUTION_FINISHED => QueryExecutionStatus::Active,
|
||||
_ => QueryExecutionStatus::Finished,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_owned_by(&self, owner: &QueryExecutionOwner) -> bool {
|
||||
Arc::ptr_eq(&self.inner.owner_identity, &owner.identity)
|
||||
}
|
||||
|
||||
pub(crate) fn is_same_execution(&self, other: &Self) -> bool {
|
||||
Arc::ptr_eq(&self.inner, &other.inner)
|
||||
}
|
||||
|
||||
pub fn mark_admitted(&self, owner: &QueryExecutionOwner) -> bool {
|
||||
self.transition(owner, EXECUTION_SETTING_UP, EXECUTION_ADMITTED)
|
||||
}
|
||||
|
||||
pub fn claim_planning(&self, owner: &QueryExecutionOwner) -> bool {
|
||||
self.transition(owner, EXECUTION_ADMITTED, EXECUTION_PLANNING)
|
||||
}
|
||||
|
||||
pub fn mark_planned(&self, owner: &QueryExecutionOwner) -> bool {
|
||||
self.transition(owner, EXECUTION_PLANNING, EXECUTION_PLANNED)
|
||||
}
|
||||
|
||||
pub fn claim_execution(&self, owner: &QueryExecutionOwner) -> bool {
|
||||
self.transition(owner, EXECUTION_PLANNED, EXECUTION_STARTING)
|
||||
}
|
||||
|
||||
pub fn mark_running(&self, owner: &QueryExecutionOwner) -> bool {
|
||||
self.transition(owner, EXECUTION_STARTING, EXECUTION_RUNNING)
|
||||
}
|
||||
|
||||
pub fn handoff_deadline(&self, owner: &QueryExecutionOwner) {
|
||||
if !self.is_owned_by(owner) || self.inner.state.load(Ordering::Acquire) != EXECUTION_RUNNING {
|
||||
return;
|
||||
}
|
||||
if let Some(deadline_task) = self.inner.deadline_task.lock().take() {
|
||||
deadline_task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn finish(&self, owner: &QueryExecutionOwner) {
|
||||
if self.is_owned_by(owner) {
|
||||
self.inner.finish();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn expire(&self, owner: &QueryExecutionOwner) {
|
||||
if self.is_owned_by(owner) {
|
||||
self.inner.expire();
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn query_guard(&self) -> Option<QueryExecutionGuard> {
|
||||
let query_guard = self.inner.query_guard.lock();
|
||||
if Instant::now() >= self.inner.deadline || self.status() != QueryExecutionStatus::Active {
|
||||
return None;
|
||||
}
|
||||
query_guard.clone()
|
||||
}
|
||||
|
||||
fn transition(&self, owner: &QueryExecutionOwner, from: u8, to: u8) -> bool {
|
||||
self.is_owned_by(owner)
|
||||
&& self
|
||||
.inner
|
||||
.state
|
||||
.compare_exchange(from, to, Ordering::AcqRel, Ordering::Acquire)
|
||||
.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
impl QueryExecutionTrackerInner {
|
||||
fn finish(&self) {
|
||||
self.state.fetch_max(EXECUTION_FINISHED, Ordering::AcqRel);
|
||||
self.release();
|
||||
}
|
||||
|
||||
fn expire(&self) {
|
||||
self.mark_timed_out();
|
||||
self.release();
|
||||
}
|
||||
|
||||
fn expire_at_deadline(&self) {
|
||||
if self
|
||||
.mark_timed_out()
|
||||
.is_some_and(|state| matches!(state, EXECUTION_ADMITTED | EXECUTION_PLANNED))
|
||||
{
|
||||
self.release();
|
||||
}
|
||||
}
|
||||
|
||||
fn mark_timed_out(&self) -> Option<u8> {
|
||||
self.state
|
||||
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
|
||||
(state < EXECUTION_FINISHED).then_some(EXECUTION_TIMED_OUT)
|
||||
})
|
||||
.ok()
|
||||
}
|
||||
|
||||
fn release(&self) {
|
||||
self.query_guard.lock().take();
|
||||
if let Some(deadline_task) = self.deadline_task.lock().take() {
|
||||
deadline_task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for QueryExecutionTrackerInner {
|
||||
fn drop(&mut self) {
|
||||
if let Some(deadline_task) = self.deadline_task.get_mut().take() {
|
||||
deadline_task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for QueryExecutionTracker {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("QueryExecutionTracker")
|
||||
.field("deadline", &self.deadline())
|
||||
.field("timeout_seconds", &self.timeout_seconds())
|
||||
.field("status", &self.status())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SessionCtx {
|
||||
_desc: Arc<SessionCtxDesc>,
|
||||
inner: SessionState,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
}
|
||||
|
||||
impl SessionCtx {
|
||||
pub fn inner(&self) -> &SessionState {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
pub(crate) fn is_bound_to(&self, query_tracker: &QueryExecutionTracker) -> bool {
|
||||
self.query_tracker
|
||||
.as_ref()
|
||||
.is_some_and(|bound_tracker| bound_tracker.is_same_execution(query_tracker))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -45,12 +263,19 @@ pub struct SessionCtxDesc {
|
||||
// maybe we need some info
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct SessionCtxFactory {
|
||||
pub is_test: bool,
|
||||
pub target_partitions: usize,
|
||||
}
|
||||
|
||||
pub const DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
impl Default for SessionCtxFactory {
|
||||
fn default() -> Self {
|
||||
Self::new(false)
|
||||
}
|
||||
}
|
||||
|
||||
impl SessionCtxFactory {
|
||||
pub fn new(is_test: bool) -> Self {
|
||||
Self {
|
||||
@@ -65,19 +290,66 @@ impl SessionCtxFactory {
|
||||
}
|
||||
|
||||
pub async fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
|
||||
let df_session_ctx = self.build_df_session_context(context).await?;
|
||||
self.create_session_ctx_inner(context, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn create_session_ctx_with_tracker_and_memory_limit(
|
||||
&self,
|
||||
context: &Context,
|
||||
query_tracker: QueryExecutionTracker,
|
||||
memory_limit_bytes: usize,
|
||||
) -> QueryResult<SessionCtx> {
|
||||
self.create_session_ctx_inner(context, Some(query_tracker), None, memory_limit_bytes)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
async fn create_session_ctx_with_tracker_and_store(
|
||||
&self,
|
||||
context: &Context,
|
||||
query_tracker: QueryExecutionTracker,
|
||||
store: Arc<SelectStore>,
|
||||
) -> QueryResult<SessionCtx> {
|
||||
self.create_session_ctx_inner(context, Some(query_tracker), Some(store), DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_session_ctx_inner(
|
||||
&self,
|
||||
context: &Context,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
store: Option<Arc<SelectStore>>,
|
||||
memory_limit_bytes: usize,
|
||||
) -> QueryResult<SessionCtx> {
|
||||
let df_session_ctx = self
|
||||
.build_df_session_context(context, query_tracker.clone(), store, memory_limit_bytes)
|
||||
.await?;
|
||||
|
||||
Ok(SessionCtx {
|
||||
_desc: Arc::new(SessionCtxDesc {}),
|
||||
inner: df_session_ctx.state(),
|
||||
query_tracker,
|
||||
})
|
||||
}
|
||||
|
||||
async fn build_df_session_context(&self, context: &Context) -> QueryResult<SessionContext> {
|
||||
async fn build_df_session_context(
|
||||
&self,
|
||||
context: &Context,
|
||||
query_tracker: Option<QueryExecutionTracker>,
|
||||
store: Option<Arc<SelectStore>>,
|
||||
memory_limit_bytes: usize,
|
||||
) -> QueryResult<SessionContext> {
|
||||
let path = format!("s3://{}", context.input.bucket);
|
||||
let store_url = url::Url::parse(&path).unwrap();
|
||||
let rt = RuntimeEnvBuilder::new().build()?;
|
||||
let memory_limit_bytes = if memory_limit_bytes == 0 {
|
||||
DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES
|
||||
} else {
|
||||
memory_limit_bytes
|
||||
};
|
||||
let rt = RuntimeEnvBuilder::new().with_memory_limit(memory_limit_bytes, 1.0).build()?;
|
||||
let config = SessionConfig::new().with_target_partitions(self.target_partitions);
|
||||
let memory_pool = Arc::clone(&rt.memory_pool);
|
||||
let df_session_state = SessionStateBuilder::new()
|
||||
.with_config(config)
|
||||
.with_runtime_env(Arc::new(rt))
|
||||
@@ -127,8 +399,13 @@ impl SessionCtxFactory {
|
||||
|
||||
df_session_state.with_object_store(&store_url, store).build()
|
||||
} else {
|
||||
let store: EcObjectStore =
|
||||
EcObjectStore::new(context.input.clone()).map_err(|_| QueryError::NotImplemented { err: String::new() })?;
|
||||
let store: EcObjectStore = match query_tracker {
|
||||
Some(query_tracker) => {
|
||||
EcObjectStore::new_with_query_tracker(context.input.clone(), memory_pool, query_tracker, store)
|
||||
}
|
||||
None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool),
|
||||
}
|
||||
.map_err(|_| QueryError::NotImplemented { err: String::new() })?;
|
||||
df_session_state.with_object_store(&store_url, Arc::new(store)).build()
|
||||
};
|
||||
|
||||
@@ -197,6 +474,7 @@ fn test_parquet_batch(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use datafusion::execution::memory_pool::MemoryLimit;
|
||||
use s3s::dto::{
|
||||
CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput,
|
||||
SelectObjectContentRequest,
|
||||
@@ -229,6 +507,17 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_factory_fields_remain_source_compatible() {
|
||||
let factory = SessionCtxFactory {
|
||||
is_test: true,
|
||||
target_partitions: 0,
|
||||
};
|
||||
|
||||
assert!(factory.is_test);
|
||||
assert_eq!(factory.target_partitions, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_factory_applies_target_partitions() {
|
||||
let factory = SessionCtxFactory::new(true).with_target_partitions(3);
|
||||
@@ -250,4 +539,78 @@ mod tests {
|
||||
|
||||
assert_eq!(session.inner().config().target_partitions(), SessionConfig::new().target_partitions());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_factory_applies_memory_limit() {
|
||||
let factory = SessionCtxFactory::new(true);
|
||||
let session = factory
|
||||
.create_session_ctx_inner(&test_context(), None, None, 1024)
|
||||
.await
|
||||
.expect("session should be created with a bounded memory pool");
|
||||
|
||||
assert!(matches!(
|
||||
session.inner().runtime_env().memory_pool.memory_limit(),
|
||||
MemoryLimit::Finite(1024)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial_test::serial]
|
||||
async fn session_factory_propagates_query_guard_to_ec_store() {
|
||||
let temp_root = tempfile::tempdir().expect("create session test temp root");
|
||||
let env = rustfs_test_utils::TestECStoreEnv::builder()
|
||||
.base_dir(temp_root.path())
|
||||
.init_bucket_metadata(false)
|
||||
.build()
|
||||
.await;
|
||||
|
||||
let admission = Arc::new(tokio::sync::Semaphore::new(1));
|
||||
let permit = Arc::clone(&admission)
|
||||
.acquire_owned()
|
||||
.await
|
||||
.expect("query permit should be available");
|
||||
let query_guard = Arc::new(permit);
|
||||
let query_tracker = QueryExecutionTracker::new(
|
||||
&QueryExecutionOwner::new(),
|
||||
Arc::clone(&query_guard),
|
||||
Instant::now() + std::time::Duration::from_secs(300),
|
||||
300,
|
||||
);
|
||||
let session = SessionCtxFactory::new(false)
|
||||
.create_session_ctx_with_tracker_and_store(&test_context(), query_tracker, Arc::clone(&env.ecstore))
|
||||
.await
|
||||
.expect("production session should be created with the query guard");
|
||||
|
||||
assert!(Arc::strong_count(&query_guard) > 1);
|
||||
drop(session);
|
||||
assert_eq!(Arc::strong_count(&query_guard), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn elapsed_deadline_does_not_yield_query_guard_before_timer_poll() {
|
||||
let admission = Arc::new(tokio::sync::Semaphore::new(1));
|
||||
let permit = Arc::clone(&admission)
|
||||
.acquire_owned()
|
||||
.await
|
||||
.expect("query permit should be available");
|
||||
let query_tracker = QueryExecutionTracker::new(&QueryExecutionOwner::new(), Arc::new(permit), Instant::now(), 0);
|
||||
|
||||
assert!(query_tracker.query_guard().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_factory_default_uses_bounded_memory() {
|
||||
let factory = SessionCtxFactory::default();
|
||||
let session = SessionCtxFactory::new(true)
|
||||
.create_session_ctx(&test_context())
|
||||
.await
|
||||
.expect("default session should be created with a bounded memory pool");
|
||||
|
||||
assert!(!factory.is_test);
|
||||
assert_eq!(factory.target_partitions, 0);
|
||||
assert!(matches!(
|
||||
session.inner().runtime_env().memory_pool.memory_limit(),
|
||||
MemoryLimit::Finite(DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -34,7 +34,7 @@ derive_builder = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
parking_lot = { workspace = true }
|
||||
s3s = { workspace = true, features = ["minio"] }
|
||||
tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
|
||||
tokio = { workspace = true, features = ["fs", "rt-multi-thread", "sync", "time"] }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[lib]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -12,18 +12,26 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::{
|
||||
sync::{Arc, LazyLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use derive_builder::Builder;
|
||||
use rustfs_s3select_api::{
|
||||
QueryResult,
|
||||
query::{
|
||||
Query, dispatcher::QueryDispatcher, execution::QueryStateMachineRef, logical_planner::Plan, session::SessionCtxFactory,
|
||||
Query,
|
||||
dispatcher::QueryDispatcher,
|
||||
execution::QueryStateMachineRef,
|
||||
logical_planner::Plan,
|
||||
session::{DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES as DEFAULT_MEMORY_LIMIT_BYTES, SessionCtxFactory},
|
||||
},
|
||||
server::dbms::{DatabaseManagerSystem, QueryHandle},
|
||||
};
|
||||
use s3s::dto::SelectObjectContentInput;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::{
|
||||
dispatcher::manager::SimpleQueryDispatcherBuilder,
|
||||
@@ -34,6 +42,16 @@ use crate::{
|
||||
};
|
||||
|
||||
const ENV_RUSTFS_S3SELECT_TARGET_PARTITIONS: &str = "RUSTFS_S3SELECT_TARGET_PARTITIONS";
|
||||
const ENV_RUSTFS_S3SELECT_MEMORY_LIMIT_BYTES: &str = "RUSTFS_S3SELECT_MEMORY_LIMIT_BYTES";
|
||||
const ENV_RUSTFS_S3SELECT_QUERY_TIMEOUT_SECS: &str = "RUSTFS_S3SELECT_QUERY_TIMEOUT_SECS";
|
||||
const ENV_RUSTFS_S3SELECT_MAX_CONCURRENT_QUERIES: &str = "RUSTFS_S3SELECT_MAX_CONCURRENT_QUERIES";
|
||||
pub(crate) const DEFAULT_QUERY_TIMEOUT_SECS: u64 = 300;
|
||||
pub(crate) const DEFAULT_MAX_CONCURRENT_QUERIES: usize = 4;
|
||||
const MAX_QUERY_TIMEOUT_SECS: u64 = 24 * 60 * 60;
|
||||
const TEST_MAX_CONCURRENT_QUERIES: usize = 1024;
|
||||
|
||||
static QUERY_ADMISSION: LazyLock<Arc<Semaphore>> =
|
||||
LazyLock::new(|| Arc::new(Semaphore::new(S3SelectRuntimeConfig::from_env().max_concurrent_queries)));
|
||||
|
||||
#[derive(Builder)]
|
||||
pub struct RustFSms<D: QueryDispatcher> {
|
||||
@@ -79,9 +97,23 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
struct S3SelectRuntimeConfig {
|
||||
target_partitions: usize,
|
||||
memory_limit_bytes: usize,
|
||||
query_timeout: Duration,
|
||||
max_concurrent_queries: usize,
|
||||
}
|
||||
|
||||
impl Default for S3SelectRuntimeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
target_partitions: 0,
|
||||
memory_limit_bytes: DEFAULT_MEMORY_LIMIT_BYTES,
|
||||
query_timeout: Duration::from_secs(DEFAULT_QUERY_TIMEOUT_SECS),
|
||||
max_concurrent_queries: DEFAULT_MAX_CONCURRENT_QUERIES,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl S3SelectRuntimeConfig {
|
||||
@@ -90,14 +122,55 @@ impl S3SelectRuntimeConfig {
|
||||
target_partitions: target_partitions_from_env_value(
|
||||
std::env::var(ENV_RUSTFS_S3SELECT_TARGET_PARTITIONS).ok().as_deref(),
|
||||
),
|
||||
memory_limit_bytes: bounded_usize_from_env_value(
|
||||
std::env::var(ENV_RUSTFS_S3SELECT_MEMORY_LIMIT_BYTES).ok().as_deref(),
|
||||
DEFAULT_MEMORY_LIMIT_BYTES,
|
||||
usize::MAX,
|
||||
),
|
||||
query_timeout: s3_select_query_timeout(),
|
||||
max_concurrent_queries: bounded_usize_from_env_value(
|
||||
std::env::var(ENV_RUSTFS_S3SELECT_MAX_CONCURRENT_QUERIES).ok().as_deref(),
|
||||
DEFAULT_MAX_CONCURRENT_QUERIES,
|
||||
Semaphore::MAX_PERMITS,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn s3_select_query_timeout() -> Duration {
|
||||
Duration::from_secs(bounded_u64_from_env_value(
|
||||
std::env::var(ENV_RUSTFS_S3SELECT_QUERY_TIMEOUT_SECS).ok().as_deref(),
|
||||
DEFAULT_QUERY_TIMEOUT_SECS,
|
||||
MAX_QUERY_TIMEOUT_SECS,
|
||||
))
|
||||
}
|
||||
|
||||
fn target_partitions_from_env_value(value: Option<&str>) -> usize {
|
||||
value.and_then(|value| value.parse::<usize>().ok()).unwrap_or(0)
|
||||
}
|
||||
|
||||
fn bounded_usize_from_env_value(value: Option<&str>, default: usize, max: usize) -> usize {
|
||||
value
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.filter(|value| (1..=max).contains(value))
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn bounded_u64_from_env_value(value: Option<&str>, default: u64, max: u64) -> u64 {
|
||||
value
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.filter(|value| (1..=max).contains(value))
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn query_admission(is_test: bool) -> Arc<Semaphore> {
|
||||
if is_test {
|
||||
Arc::new(Semaphore::new(TEST_MAX_CONCURRENT_QUERIES))
|
||||
} else {
|
||||
Arc::clone(&QUERY_ADMISSION)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn make_rustfsms(input: Arc<SelectObjectContentInput>, is_test: bool) -> QueryResult<impl DatabaseManagerSystem> {
|
||||
// init Function Manager, we can define some UDF if need
|
||||
let func_manager = SimpleFunctionMetadataManager::default();
|
||||
@@ -116,8 +189,11 @@ pub async fn make_rustfsms(input: Arc<SelectObjectContentInput>, is_test: bool)
|
||||
.with_func_manager(Arc::new(func_manager))
|
||||
.with_default_table_provider(default_table_provider)
|
||||
.with_session_factory(session_factory)
|
||||
.with_memory_limit_bytes(runtime_config.memory_limit_bytes)
|
||||
.with_parser(parser)
|
||||
.with_query_execution_factory(query_execution_factory)
|
||||
.with_query_admission(query_admission(is_test))
|
||||
.with_query_timeout(runtime_config.query_timeout)
|
||||
.build()?;
|
||||
|
||||
let mut builder = RustFSmsBuilder::default();
|
||||
@@ -143,8 +219,11 @@ pub async fn make_rustfsms_with_components(
|
||||
.with_func_manager(func_manager)
|
||||
.with_default_table_provider(default_table_provider)
|
||||
.with_session_factory(session_factory)
|
||||
.with_memory_limit_bytes(runtime_config.memory_limit_bytes)
|
||||
.with_parser(parser)
|
||||
.with_query_execution_factory(query_execution_factory)
|
||||
.with_query_admission(query_admission(is_test))
|
||||
.with_query_timeout(runtime_config.query_timeout)
|
||||
.build()?;
|
||||
|
||||
let mut builder = RustFSmsBuilder::default();
|
||||
@@ -167,7 +246,10 @@ mod tests {
|
||||
|
||||
use crate::get_global_db;
|
||||
|
||||
use super::{S3SelectRuntimeConfig, target_partitions_from_env_value};
|
||||
use super::{
|
||||
DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_MEMORY_LIMIT_BYTES, DEFAULT_QUERY_TIMEOUT_SECS, MAX_QUERY_TIMEOUT_SECS,
|
||||
S3SelectRuntimeConfig, bounded_u64_from_env_value, bounded_usize_from_env_value, target_partitions_from_env_value,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn parses_target_partitions_from_env_value() {
|
||||
@@ -179,7 +261,24 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn default_runtime_config_uses_datafusion_default_partitions() {
|
||||
assert_eq!(S3SelectRuntimeConfig::default().target_partitions, 0);
|
||||
let config = S3SelectRuntimeConfig::default();
|
||||
|
||||
assert_eq!(config.target_partitions, 0);
|
||||
assert_eq!(config.memory_limit_bytes, DEFAULT_MEMORY_LIMIT_BYTES);
|
||||
assert_eq!(config.query_timeout.as_secs(), DEFAULT_QUERY_TIMEOUT_SECS);
|
||||
assert_eq!(config.max_concurrent_queries, DEFAULT_MAX_CONCURRENT_QUERIES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resource_limits_reject_invalid_and_out_of_range_values() {
|
||||
assert_eq!(bounded_usize_from_env_value(Some("1024"), 64, 2048), 1024);
|
||||
assert_eq!(bounded_usize_from_env_value(Some("0"), 64, 2048), 64);
|
||||
assert_eq!(bounded_usize_from_env_value(Some("4096"), 64, 2048), 64);
|
||||
assert_eq!(bounded_usize_from_env_value(Some("invalid"), 64, 2048), 64);
|
||||
assert_eq!(bounded_u64_from_env_value(Some("30"), 300, MAX_QUERY_TIMEOUT_SECS), 30);
|
||||
assert_eq!(bounded_u64_from_env_value(Some("0"), 300, MAX_QUERY_TIMEOUT_SECS), 300);
|
||||
assert_eq!(bounded_u64_from_env_value(Some("86401"), 300, MAX_QUERY_TIMEOUT_SECS), 300);
|
||||
assert_eq!(bounded_u64_from_env_value(None, 300, MAX_QUERY_TIMEOUT_SECS), 300);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -19,7 +19,7 @@ use datafusion::arrow::datatypes::DataType;
|
||||
use datafusion::common::{Result as DFResult, TableReference};
|
||||
use datafusion::datasource::TableProvider;
|
||||
use datafusion::logical_expr::var_provider::is_system_variables;
|
||||
use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF};
|
||||
use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF, planner::ExprPlanner};
|
||||
use datafusion::variable::VarType;
|
||||
use datafusion::{config::ConfigOptions, sql::planner::ContextProvider};
|
||||
use rustfs_s3select_api::query::{function::FuncMetaManagerRef, session::SessionCtx};
|
||||
@@ -81,6 +81,10 @@ impl ContextProviderExtension for MetadataProvider {
|
||||
}
|
||||
|
||||
impl ContextProvider for MetadataProvider {
|
||||
fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
|
||||
self.session.inner().expr_planners()
|
||||
}
|
||||
|
||||
fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
|
||||
self.func_manager
|
||||
.udf(name)
|
||||
|
||||
@@ -12,11 +12,18 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::ops::ControlFlow;
|
||||
|
||||
use async_recursion::async_recursion;
|
||||
use async_trait::async_trait;
|
||||
use datafusion::sql::{planner::SqlToRel, sqlparser::ast::Statement};
|
||||
use datafusion::sql::{
|
||||
planner::SqlToRel,
|
||||
sqlparser::ast::{
|
||||
GroupByExpr, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor, SetExpr, Statement, TableFactor, Visit, Visitor,
|
||||
},
|
||||
};
|
||||
use rustfs_s3select_api::{
|
||||
QueryError, QueryResult,
|
||||
QueryError, QueryResult, S3SelectPolicyError,
|
||||
query::{
|
||||
ast::ExtStatement,
|
||||
logical_planner::{LogicalPlanner, Plan, QueryPlan},
|
||||
@@ -60,6 +67,7 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
|
||||
async fn df_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult<Plan> {
|
||||
match stmt {
|
||||
Statement::Query(_) => {
|
||||
validate_s3_select_statement(&stmt)?;
|
||||
let df_plan = self.df_planner.sql_statement_to_plan(stmt)?;
|
||||
let plan = Plan::Query(QueryPlan {
|
||||
df_plan,
|
||||
@@ -72,3 +80,256 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> {
|
||||
let Statement::Query(query) = statement else {
|
||||
return Err(unsupported_structure("only SELECT queries are supported"));
|
||||
};
|
||||
|
||||
if query.with.is_some()
|
||||
|| query.order_by.as_ref().is_some_and(|order_by| {
|
||||
order_by.interpolate.is_some()
|
||||
|| match &order_by.kind {
|
||||
OrderByKind::Expressions(expressions) => expressions.iter().any(|expression| expression.with_fill.is_some()),
|
||||
OrderByKind::All(_) => true,
|
||||
}
|
||||
})
|
||||
|| query.fetch.is_some()
|
||||
|| !query.locks.is_empty()
|
||||
|| query.for_clause.is_some()
|
||||
|| query.settings.is_some()
|
||||
|| query.format_clause.is_some()
|
||||
|| !query.pipe_operators.is_empty()
|
||||
{
|
||||
return Err(unsupported_structure("the query contains an unsupported clause"));
|
||||
}
|
||||
if let Some(limit_clause) = query.limit_clause.as_ref()
|
||||
&& !matches!(
|
||||
limit_clause,
|
||||
datafusion::sql::sqlparser::ast::LimitClause::LimitOffset {
|
||||
limit: Some(_),
|
||||
offset: None,
|
||||
limit_by,
|
||||
} if limit_by.is_empty()
|
||||
)
|
||||
{
|
||||
return Err(unsupported_structure("only LIMIT without OFFSET is supported"));
|
||||
}
|
||||
if let Some(datafusion::sql::sqlparser::ast::LimitClause::LimitOffset { limit: Some(limit), .. }) =
|
||||
query.limit_clause.as_ref()
|
||||
&& limit.to_string().parse::<u64>().is_err()
|
||||
{
|
||||
return Err(unsupported_structure("LIMIT must be a non-negative integer"));
|
||||
}
|
||||
|
||||
let mut detector = SubqueryDetector { visited_root: false };
|
||||
if query.visit(&mut detector).is_break() {
|
||||
return Err(unsupported_structure("subqueries are not supported"));
|
||||
}
|
||||
|
||||
let SetExpr::Select(select) = query.body.as_ref() else {
|
||||
return Err(unsupported_structure("set operations and nested queries are not supported"));
|
||||
};
|
||||
validate_select(select)
|
||||
}
|
||||
|
||||
fn validate_select(select: &Select) -> QueryResult<()> {
|
||||
if !select.optimizer_hints.is_empty()
|
||||
|| select.distinct.is_some()
|
||||
|| select.select_modifiers.is_some()
|
||||
|| select.top.is_some()
|
||||
|| select.exclude.is_some()
|
||||
|| select.into.is_some()
|
||||
|| !select.lateral_views.is_empty()
|
||||
|| select.prewhere.is_some()
|
||||
|| !select.connect_by.is_empty()
|
||||
|| !select.cluster_by.is_empty()
|
||||
|| !select.distribute_by.is_empty()
|
||||
|| !select.sort_by.is_empty()
|
||||
|| select.having.is_some()
|
||||
|| !select.named_window.is_empty()
|
||||
|| select.qualify.is_some()
|
||||
|| select.value_table_mode.is_some()
|
||||
|| select.flavor != SelectFlavor::Standard
|
||||
|| !matches!(&select.group_by, GroupByExpr::Expressions(_, modifiers) if modifiers.is_empty())
|
||||
{
|
||||
return Err(unsupported_structure("the SELECT contains an unsupported clause"));
|
||||
}
|
||||
|
||||
let [table] = select.from.as_slice() else {
|
||||
return Err(unsupported_structure("exactly one S3Object source is required"));
|
||||
};
|
||||
if !table.joins.is_empty() {
|
||||
return Err(unsupported_structure("JOIN is not supported"));
|
||||
}
|
||||
let TableFactor::Table {
|
||||
name,
|
||||
alias,
|
||||
args,
|
||||
with_hints,
|
||||
version,
|
||||
with_ordinality,
|
||||
partitions,
|
||||
sample,
|
||||
index_hints,
|
||||
..
|
||||
} = &table.relation
|
||||
else {
|
||||
return Err(unsupported_structure("subqueries and table functions are not supported"));
|
||||
};
|
||||
if args.is_some()
|
||||
|| !with_hints.is_empty()
|
||||
|| version.is_some()
|
||||
|| *with_ordinality
|
||||
|| !partitions.is_empty()
|
||||
|| sample.is_some()
|
||||
|| !index_hints.is_empty()
|
||||
|| alias.as_ref().is_some_and(|alias| !alias.columns.is_empty())
|
||||
{
|
||||
return Err(unsupported_structure("the S3Object source contains unsupported modifiers"));
|
||||
}
|
||||
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"));
|
||||
};
|
||||
let is_s3_object = if table_name.quote_style.is_some() {
|
||||
table_name.value == "S3Object"
|
||||
} else {
|
||||
table_name.value.eq_ignore_ascii_case("S3Object")
|
||||
};
|
||||
if !is_s3_object {
|
||||
return Err(unsupported_structure("the source must be S3Object"));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unsupported_structure(message: &str) -> QueryError {
|
||||
S3SelectPolicyError::UnsupportedSqlStructure {
|
||||
message: message.to_string(),
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
struct SubqueryDetector {
|
||||
visited_root: bool,
|
||||
}
|
||||
|
||||
impl Visitor for SubqueryDetector {
|
||||
type Break = ();
|
||||
|
||||
fn pre_visit_query(&mut self, _query: &Query) -> ControlFlow<Self::Break> {
|
||||
if self.visited_root {
|
||||
ControlFlow::Break(())
|
||||
} else {
|
||||
self.visited_root = true;
|
||||
ControlFlow::Continue(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
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};
|
||||
|
||||
fn parse_statement(sql: &str) -> Statement {
|
||||
let mut statements = ExtParser::parse_sql(sql).expect("SQL should parse");
|
||||
let ExtStatement::SqlStatement(statement) = statements.pop_front().expect("one SQL statement");
|
||||
*statement
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_s3_select_query_shape() {
|
||||
let statement = parse_statement("SELECT s.id FROM S3Object AS s WHERE s.id = '1' LIMIT 10");
|
||||
|
||||
assert!(validate_s3_select_statement(&statement).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_json_sub_path_source() {
|
||||
let statement = parse_statement("SELECT e.name FROM S3Object.employees AS e");
|
||||
|
||||
assert!(validate_s3_select_statement(&statement).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_group_by_and_order_by() {
|
||||
let statement = parse_statement("SELECT department, COUNT(*) FROM S3Object GROUP BY department ORDER BY department");
|
||||
|
||||
assert!(validate_s3_select_statement(&statement).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_join() {
|
||||
let statement = parse_statement("SELECT * FROM S3Object a JOIN S3Object b ON a.id = b.id");
|
||||
|
||||
assert!(matches!(
|
||||
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"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_subquery() {
|
||||
let statement = parse_statement("SELECT * FROM S3Object WHERE id IN (SELECT id FROM S3Object)");
|
||||
|
||||
assert!(matches!(
|
||||
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"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_subquery_in_order_by() {
|
||||
let statement = parse_statement("SELECT id FROM S3Object ORDER BY (SELECT id FROM S3Object)");
|
||||
|
||||
assert!(matches!(
|
||||
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"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_s3_object_source() {
|
||||
let statement = parse_statement("SELECT * FROM other_table");
|
||||
|
||||
assert!(matches!(
|
||||
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"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unsupported_select_clauses() {
|
||||
for sql in [
|
||||
"SELECT DISTINCT id FROM S3Object",
|
||||
"SELECT * FROM S3Object OFFSET 1",
|
||||
"SELECT * FROM S3Object UNION SELECT * FROM S3Object",
|
||||
] {
|
||||
let statement = parse_statement(sql);
|
||||
assert!(
|
||||
matches!(
|
||||
validate_s3_select_statement(&statement),
|
||||
Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::UnsupportedSqlStructure { .. }))
|
||||
),
|
||||
"query should be rejected: {sql}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,7 +15,10 @@
|
||||
#[cfg(test)]
|
||||
mod integration_tests {
|
||||
use crate::{create_fresh_db, get_global_db, instance::make_rustfsms};
|
||||
use datafusion::arrow::array::{Array, StringArray};
|
||||
use datafusion::arrow::{
|
||||
array::{Array, Int64Array, StringArray},
|
||||
record_batch::RecordBatch,
|
||||
};
|
||||
use rustfs_s3select_api::{
|
||||
QueryError,
|
||||
query::{Context, Query},
|
||||
@@ -26,6 +29,52 @@ mod integration_tests {
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
fn assert_ages_descending(output: &[RecordBatch]) {
|
||||
let ages: Vec<i64> = output
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.expect("age column should be Int64")
|
||||
.values()
|
||||
.iter()
|
||||
.copied()
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert_eq!(ages, vec![40, 38, 35, 32, 30, 28, 26, 25, 24, 22]);
|
||||
}
|
||||
|
||||
fn assert_department_counts(output: &[RecordBatch]) {
|
||||
let mut counts: Vec<(&str, i64)> = output
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
let departments = batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.expect("department column should be Utf8");
|
||||
let counts = batch
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.expect("count column should be Int64");
|
||||
departments.iter().zip(counts.iter()).map(|(department, count)| {
|
||||
(
|
||||
department.expect("department should not be null"),
|
||||
count.expect("count should not be null"),
|
||||
)
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
counts.sort_unstable();
|
||||
|
||||
assert_eq!(counts, vec![("Finance", 3), ("HR", 2), ("IT", 3), ("Marketing", 2)]);
|
||||
}
|
||||
|
||||
fn create_test_input(sql: &str) -> SelectObjectContentInput {
|
||||
SelectObjectContentInput {
|
||||
bucket: "test-bucket".to_string(),
|
||||
@@ -292,21 +341,21 @@ mod integration_tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_with_aggregation() {
|
||||
let sql = "SELECT department, COUNT(*) as count FROM S3Object GROUP BY department";
|
||||
let sql = "SELECT department, COUNT(*) FROM S3Object GROUP BY department";
|
||||
let input = create_test_input(sql);
|
||||
let db = get_global_db(input.clone(), true).await.unwrap();
|
||||
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||
|
||||
let result = db.execute(&query).await;
|
||||
// Aggregation queries might fail due to lack of actual data, which is acceptable
|
||||
match result {
|
||||
Ok(_) => {
|
||||
// If successful, that's great
|
||||
}
|
||||
Err(_) => {
|
||||
// Expected to fail due to no actual data source
|
||||
}
|
||||
}
|
||||
let output = db
|
||||
.execute(&query)
|
||||
.await
|
||||
.expect("execute grouped CSV query")
|
||||
.result()
|
||||
.chunk_result()
|
||||
.await
|
||||
.expect("collect grouped CSV output");
|
||||
|
||||
assert_department_counts(&output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -381,13 +430,22 @@ mod integration_tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_with_order_by() {
|
||||
let sql = "SELECT name, age FROM S3Object ORDER BY age DESC";
|
||||
let sql = "SELECT name, CAST(age AS BIGINT) AS age FROM S3Object ORDER BY CAST(age AS BIGINT) DESC";
|
||||
let input = create_test_input(sql);
|
||||
let db = get_global_db(input.clone(), true).await.unwrap();
|
||||
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||
|
||||
let result = db.execute(&query).await;
|
||||
assert!(result.is_ok());
|
||||
let output = db
|
||||
.execute(&query)
|
||||
.await
|
||||
.expect("execute ordered CSV query")
|
||||
.result()
|
||||
.chunk_result()
|
||||
.await
|
||||
.expect("collect ordered CSV output");
|
||||
|
||||
assert_eq!(output.iter().map(|batch| batch.num_rows()).sum::<usize>(), 10);
|
||||
assert_ages_descending(&output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -582,21 +640,21 @@ mod integration_tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_with_aggregation_json() {
|
||||
let sql = "SELECT department, COUNT(*) as count FROM S3Object GROUP BY department";
|
||||
let sql = "SELECT department, COUNT(*) FROM S3Object GROUP BY department";
|
||||
let input = create_test_json_input(sql);
|
||||
let db = get_global_db(input.clone(), true).await.unwrap();
|
||||
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||
|
||||
let result = db.execute(&query).await;
|
||||
// Aggregation queries may fail due to lack of actual data, which is acceptable
|
||||
match result {
|
||||
Ok(_) => {
|
||||
// If successful, that's great
|
||||
}
|
||||
Err(_) => {
|
||||
// Expected to fail due to no actual data source
|
||||
}
|
||||
}
|
||||
let output = db
|
||||
.execute(&query)
|
||||
.await
|
||||
.expect("execute grouped JSON query")
|
||||
.result()
|
||||
.chunk_result()
|
||||
.await
|
||||
.expect("collect grouped JSON output");
|
||||
|
||||
assert_department_counts(&output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -672,8 +730,17 @@ mod integration_tests {
|
||||
let db = get_global_db(input.clone(), true).await.unwrap();
|
||||
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||
|
||||
let result = db.execute(&query).await;
|
||||
assert!(result.is_ok());
|
||||
let output = db
|
||||
.execute(&query)
|
||||
.await
|
||||
.expect("execute ordered JSON query")
|
||||
.result()
|
||||
.chunk_result()
|
||||
.await
|
||||
.expect("collect ordered JSON output");
|
||||
|
||||
assert_eq!(output.iter().map(|batch| batch.num_rows()).sum::<usize>(), 10);
|
||||
assert_ages_descending(&output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
+272
-61
@@ -10,13 +10,16 @@ use datafusion::arrow::{
|
||||
json::{WriterBuilder as JsonWriterBuilder, writer::LineDelimited},
|
||||
record_batch::RecordBatch,
|
||||
};
|
||||
use datafusion::common::DataFusionError;
|
||||
use datafusion::physical_plan::SendableRecordBatchStream;
|
||||
use futures::StreamExt;
|
||||
use http::{StatusCode, header::RANGE};
|
||||
use rustfs_s3select_api::{
|
||||
QueryError,
|
||||
QueryError, S3SelectPolicyError,
|
||||
object_store::{INVALID_SCAN_RANGE_MESSAGE, validate_scan_range_bounds},
|
||||
query::{Context, Query},
|
||||
};
|
||||
use rustfs_s3select_query::instance::s3_select_query_timeout;
|
||||
use s3s::dto::{
|
||||
CSVOutput, CompressionType, ContinuationEvent, EndEvent, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput,
|
||||
JSONOutput, JSONType, OutputSerialization, Progress, ProgressEvent, QuoteFields, RecordsEvent, SelectObjectContentEvent,
|
||||
@@ -26,6 +29,7 @@ use s3s::dto::{
|
||||
use s3s::{S3Error, S3ErrorCode, S3Request, S3Response, S3Result, s3_error};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::{Instant, timeout_at};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tracing::info;
|
||||
|
||||
@@ -60,6 +64,8 @@ pub async fn execute_select_object_content(
|
||||
validate_scan_range_for_object_size(&input.request, metadata.size)?;
|
||||
|
||||
let input = Arc::new(input);
|
||||
let query_timeout = s3_select_query_timeout();
|
||||
let query_deadline = Instant::now() + query_timeout;
|
||||
let db = current_s3select_db((*input).clone(), false)
|
||||
.await
|
||||
.map_err(map_query_error_to_s3)?;
|
||||
@@ -72,66 +78,13 @@ pub async fn execute_select_object_content(
|
||||
.into_record_batch_stream()
|
||||
.map_err(map_query_error_to_s3)?;
|
||||
|
||||
let (tx, rx) = mpsc::channel::<S3Result<SelectObjectContentEvent>>(8);
|
||||
let (tx, rx) = mpsc::channel::<S3Result<SelectObjectContentEvent>>(9);
|
||||
let terminal_permit = tx
|
||||
.clone()
|
||||
.try_reserve_owned()
|
||||
.map_err(|_| s3_error!(InternalError, "can't reserve Select terminal event capacity"))?;
|
||||
spawn_traced(async move {
|
||||
let mut encoder = SelectOutputEncoder::new(validation.output_format);
|
||||
let mut progress = SelectProgress::default();
|
||||
let mut output = output;
|
||||
|
||||
if tx
|
||||
.send(Ok(SelectObjectContentEvent::Cont(ContinuationEvent::default())))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
while let Some(result) = output.next().await {
|
||||
let batch = match result {
|
||||
Ok(batch) => batch,
|
||||
Err(err) => {
|
||||
let _ = tx.send(Err(map_query_error_to_s3(err.into()))).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match encoder.encode_batch(&batch) {
|
||||
Ok(payloads) => {
|
||||
for payload in payloads {
|
||||
progress.add_returned(payload.len());
|
||||
if tx
|
||||
.send(Ok(SelectObjectContentEvent::Records(RecordsEvent { payload: Some(payload) })))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
if validation.progress_enabled
|
||||
&& tx
|
||||
.send(Ok(SelectObjectContentEvent::Progress(ProgressEvent {
|
||||
details: Some(progress.to_progress()),
|
||||
})))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = tx.send(Err(err)).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let stats = SelectObjectContentEvent::Stats(StatsEvent {
|
||||
details: Some(progress.to_stats()),
|
||||
});
|
||||
if tx.send(Ok(stats)).await.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = tx.send(Ok(SelectObjectContentEvent::End(EndEvent::default()))).await;
|
||||
send_select_events_until_deadline(output, tx, terminal_permit, validation, query_deadline, query_timeout.as_secs()).await;
|
||||
});
|
||||
|
||||
Ok(S3Response::new(SelectObjectContentOutput {
|
||||
@@ -139,6 +92,91 @@ pub async fn execute_select_object_content(
|
||||
}))
|
||||
}
|
||||
|
||||
async fn send_select_events_until_deadline(
|
||||
output: SendableRecordBatchStream,
|
||||
tx: mpsc::Sender<S3Result<SelectObjectContentEvent>>,
|
||||
terminal_permit: mpsc::OwnedPermit<S3Result<SelectObjectContentEvent>>,
|
||||
validation: SelectValidation,
|
||||
deadline: Instant,
|
||||
timeout_seconds: u64,
|
||||
) {
|
||||
if timeout_at(deadline, send_select_events(output, &tx, validation))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
terminal_permit.send(Err(map_query_error_to_s3(
|
||||
S3SelectPolicyError::QueryTimeout {
|
||||
seconds: timeout_seconds,
|
||||
}
|
||||
.into(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_select_events(
|
||||
mut output: SendableRecordBatchStream,
|
||||
tx: &mpsc::Sender<S3Result<SelectObjectContentEvent>>,
|
||||
validation: SelectValidation,
|
||||
) {
|
||||
let mut encoder = SelectOutputEncoder::new(validation.output_format);
|
||||
let mut progress = SelectProgress::default();
|
||||
|
||||
if tx
|
||||
.send(Ok(SelectObjectContentEvent::Cont(ContinuationEvent::default())))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
while let Some(result) = output.next().await {
|
||||
let batch = match result {
|
||||
Ok(batch) => batch,
|
||||
Err(err) => {
|
||||
let _ = tx.send(Err(map_query_error_to_s3(err.into()))).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match encoder.encode_batch(&batch) {
|
||||
Ok(payloads) => {
|
||||
for payload in payloads {
|
||||
progress.add_returned(payload.len());
|
||||
if tx
|
||||
.send(Ok(SelectObjectContentEvent::Records(RecordsEvent { payload: Some(payload) })))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
if validation.progress_enabled
|
||||
&& tx
|
||||
.send(Ok(SelectObjectContentEvent::Progress(ProgressEvent {
|
||||
details: Some(progress.to_progress()),
|
||||
})))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = tx.send(Err(err)).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let stats = SelectObjectContentEvent::Stats(StatsEvent {
|
||||
details: Some(progress.to_stats()),
|
||||
});
|
||||
if tx.send(Ok(stats)).await.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = tx.send(Ok(SelectObjectContentEvent::End(EndEvent::default()))).await;
|
||||
}
|
||||
|
||||
fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectContentInput) -> S3Result<SelectValidation> {
|
||||
if headers.contains_key(RANGE) {
|
||||
return Err(S3Error::new(S3ErrorCode::UnsupportedRangeHeader));
|
||||
@@ -487,11 +525,31 @@ fn clamp_i64(value: u64) -> i64 {
|
||||
}
|
||||
|
||||
fn map_query_error_to_s3(err: QueryError) -> S3Error {
|
||||
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();
|
||||
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)
|
||||
}
|
||||
QueryError::Datafusion { source } if is_unexpected_eof(source.as_ref()) => {
|
||||
S3Error::with_message(S3ErrorCode::InternalError, message)
|
||||
}
|
||||
QueryError::Datafusion { source } if is_invalid_object_size(source.as_ref()) => {
|
||||
S3Error::with_message(S3ErrorCode::InternalError, message)
|
||||
}
|
||||
QueryError::Datafusion { .. } if looks_like_invalid_scan_range(&message) => {
|
||||
S3Error::with_message(S3ErrorCode::InvalidRequestParameter, INVALID_SCAN_RANGE_MESSAGE.to_string())
|
||||
}
|
||||
@@ -520,6 +578,42 @@ 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")
|
||||
@@ -548,10 +642,27 @@ fn is_json_document(json: &JSONInput) -> bool {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use datafusion::sql::sqlparser::parser::ParserError;
|
||||
use datafusion::{
|
||||
arrow::datatypes::Schema, physical_plan::stream::RecordBatchStreamAdapter, sql::sqlparser::parser::ParserError,
|
||||
};
|
||||
use http::HeaderMap;
|
||||
use s3s::dto::{CSVInput, ParquetInput, ScanRange};
|
||||
|
||||
#[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)
|
||||
}
|
||||
}
|
||||
|
||||
fn base_input() -> SelectObjectContentInput {
|
||||
SelectObjectContentInput {
|
||||
bucket: "bucket".to_string(),
|
||||
@@ -611,6 +722,106 @@ mod tests {
|
||||
assert_eq!(err.message(), Some("sql parser error: syntax error"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn map_query_policy_errors_to_s3_errors() {
|
||||
let unsupported = map_query_error_to_s3(
|
||||
S3SelectPolicyError::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 stream_timed_out = map_query_error_to_s3(QueryError::Datafusion {
|
||||
source: Box::new(DataFusionError::External(Box::new(S3SelectPolicyError::QueryTimeout { seconds: 300 }))),
|
||||
});
|
||||
let exhausted = map_query_error_to_s3(QueryError::Datafusion {
|
||||
source: Box::new(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(DataFusionError::ResourcesExhausted("memory limit".to_string())),
|
||||
}))),
|
||||
});
|
||||
let truncated = 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")),
|
||||
}))),
|
||||
});
|
||||
let invalid_object_size = map_query_error_to_s3(QueryError::Datafusion {
|
||||
source: Box::new(DataFusionError::ObjectStore(Box::new(datafusion::object_store::Error::Generic {
|
||||
store: "EcObjectStore",
|
||||
source: Box::new(u64::try_from(-1_i64).expect_err("negative size must fail conversion")),
|
||||
}))),
|
||||
});
|
||||
|
||||
assert_eq!(unsupported.code(), &S3ErrorCode::UnsupportedSqlStructure);
|
||||
assert_eq!(unsupported.message(), Some("Unsupported S3 Select SQL structure: JOIN is not supported"));
|
||||
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!(invalid_object_size.code(), &S3ErrorCode::InternalError);
|
||||
}
|
||||
|
||||
#[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));
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn producer_deadline_cancels_backpressured_send() {
|
||||
let output = Box::pin(RecordBatchStreamAdapter::new(
|
||||
Arc::new(Schema::empty()),
|
||||
futures::stream::pending::<Result<RecordBatch, DataFusionError>>(),
|
||||
));
|
||||
let (tx, mut rx) = mpsc::channel(2);
|
||||
let terminal_permit = tx
|
||||
.clone()
|
||||
.try_reserve_owned()
|
||||
.expect("test channel should reserve terminal capacity");
|
||||
tx.send(Ok(SelectObjectContentEvent::Cont(ContinuationEvent::default())))
|
||||
.await
|
||||
.expect("test channel should accept the prefilled event");
|
||||
let validation = SelectValidation {
|
||||
output_format: SelectOutputFormat::Csv(CSVOutput::default()),
|
||||
progress_enabled: false,
|
||||
};
|
||||
|
||||
let producer = tokio::spawn(send_select_events_until_deadline(
|
||||
output,
|
||||
tx,
|
||||
terminal_permit,
|
||||
validation,
|
||||
Instant::now() + std::time::Duration::from_secs(1),
|
||||
300,
|
||||
));
|
||||
|
||||
tokio::task::yield_now().await;
|
||||
tokio::time::advance(std::time::Duration::from_secs(1)).await;
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
producer.await.expect("producer should finish without draining the channel");
|
||||
assert!(matches!(rx.recv().await, Some(Ok(SelectObjectContentEvent::Cont(_)))));
|
||||
let timeout_error = rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("producer should send a terminal timeout error")
|
||||
.expect_err("terminal event should be an error");
|
||||
assert_eq!(timeout_error.code(), &S3ErrorCode::Busy);
|
||||
assert!(rx.recv().await.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_defaults_csv_header_and_compression() {
|
||||
let mut input = base_input();
|
||||
|
||||
Reference in New Issue
Block a user