diff --git a/crates/e2e_test/src/reliant/sql.rs b/crates/e2e_test/src/reliant/sql.rs index d9f821a15..ab2d385ab 100644 --- a/crates/e2e_test/src/reliant/sql.rs +++ b/crates/e2e_test/src/reliant/sql.rs @@ -27,6 +27,9 @@ use std::time::Duration; const BUCKET: &str = "test-sql-bucket"; const CSV_OBJECT: &str = "test-data.csv"; const JSON_OBJECT: &str = "test-data.json"; +const JSON_DOCUMENT_OBJECT: &str = "nested-data.json"; +const JSON_ROOT_ARRAY_OBJECT: &str = "root-array.json"; +const JSON_ROOT_SCALAR_ARRAY_OBJECT: &str = "root-scalars.json"; const SELECT_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30); type TestResult = Result>; @@ -74,6 +77,51 @@ async fn upload_test_json(client: &Client) -> TestResult<()> { Ok(()) } +async fn upload_nested_json_document(client: &Client) -> TestResult<()> { + let json_data = r#"{"departments":[{"employees":[{"name":"Alice","active":true},{"name":"Bob","active":false}]},{"employees":[{"name":"Charlie","active":true}]}]}"#; + + client + .put_object() + .bucket(BUCKET) + .key(JSON_DOCUMENT_OBJECT) + .body(Bytes::from_static(json_data.as_bytes()).into()) + .send() + .await?; + client + .put_object() + .bucket(BUCKET) + .key(JSON_ROOT_ARRAY_OBJECT) + .body(Bytes::from_static(br#"[{"name":"Alice"},{"name":"Bob"}]"#).into()) + .send() + .await?; + client + .put_object() + .bucket(BUCKET) + .key(JSON_ROOT_SCALAR_ARRAY_OBJECT) + .body(Bytes::from_static(b"[1,2]").into()) + .send() + .await?; + Ok(()) +} + +async fn select_json_document(client: &Client, key: &str, expression: &str) -> TestResult { + let response = client + .select_object_content() + .bucket(BUCKET) + .key(key) + .expression(expression) + .expression_type(ExpressionType::Sql) + .input_serialization( + InputSerialization::builder() + .json(JsonInput::builder().set_type(Some(JsonType::Document)).build()) + .build(), + ) + .output_serialization(OutputSerialization::builder().json(JsonOutput::builder().build()).build()) + .send() + .await?; + process_select_response(response).await +} + async fn process_select_response( mut event_stream: aws_sdk_s3::operation::select_object_content::SelectObjectContentOutput, ) -> TestResult { @@ -365,6 +413,107 @@ async fn test_select_object_content_json_basic() -> TestResult<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_select_object_content_nested_json_source_path() -> TestResult<()> { + let (_env, client) = create_test_environment().await?; + setup_test_bucket(&client).await?; + upload_nested_json_document(&client).await?; + + let result = select_json_document( + &client, + JSON_DOCUMENT_OBJECT, + "SELECT e.name FROM S3Object[*].departments[*].employees[*] AS e WHERE e.active = true", + ) + .await?; + let names: Vec = result + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["name"].as_str().ok_or("missing name field")?.to_string()) + }) + .collect::>()?; + + assert_eq!(names, vec!["Alice", "Charlie"]); + + let terminal_scalars = select_json_document( + &client, + JSON_DOCUMENT_OBJECT, + "SELECT NAME FROM S3Object[*].DEPARTMENTS[*].employees[*].NAME", + ) + .await?; + let scalar_names: Vec = terminal_scalars + .lines() + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["name"].as_str().ok_or("missing scalar name field")?.to_string()) + }) + .collect::>()?; + assert_eq!(scalar_names, vec!["Alice", "Bob", "Charlie"]); + + let aliased_scalars = select_json_document( + &client, + JSON_DOCUMENT_OBJECT, + "SELECT v FROM S3Object[*].departments[*].employees[*].name AS v", + ) + .await?; + let aliased_names: Vec = aliased_scalars + .lines() + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["v"].as_str().ok_or("missing aliased scalar field")?.to_string()) + }) + .collect::>()?; + assert_eq!(aliased_names, vec!["Alice", "Bob", "Charlie"]); + + let root_array = select_json_document(&client, JSON_ROOT_ARRAY_OBJECT, "SELECT c.name FROM S3Object[*][*] AS c").await?; + let root_names: Vec = root_array + .lines() + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["name"].as_str().ok_or("missing root-array name field")?.to_string()) + }) + .collect::>()?; + assert_eq!(root_names, vec!["Alice", "Bob"]); + + let root_index = select_json_document(&client, JSON_ROOT_ARRAY_OBJECT, "SELECT c.name FROM S3Object[*][0] AS c").await?; + let root_index_value: serde_json::Value = serde_json::from_str(root_index.trim())?; + assert_eq!(root_index_value["name"], "Alice"); + + let root_scalars = select_json_document(&client, JSON_ROOT_SCALAR_ARRAY_OBJECT, "SELECT V FROM S3Object AS V").await?; + let scalar_values: Vec = root_scalars + .lines() + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["v"].as_i64().ok_or("missing root scalar value")?) + }) + .collect::>()?; + assert_eq!(scalar_values, vec![1, 2]); + + let implicit_root_scalars = + select_json_document(&client, JSON_ROOT_SCALAR_ARRAY_OBJECT, "SELECT S3Object FROM S3Object").await?; + let implicit_scalar_values: Vec = implicit_root_scalars + .lines() + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["s3object"].as_i64().ok_or("missing implicit root scalar value")?) + }) + .collect::>()?; + assert_eq!(implicit_scalar_values, vec![1, 2]); + + let quoted_root_scalars = + select_json_document(&client, JSON_ROOT_SCALAR_ARRAY_OBJECT, "SELECT \"S3Object\" FROM \"S3Object\"").await?; + let quoted_scalar_values: Vec = quoted_root_scalars + .lines() + .map(|line| -> TestResult { + let value: serde_json::Value = serde_json::from_str(line)?; + Ok(value["S3Object"].as_i64().ok_or("missing quoted root scalar value")?) + }) + .collect::>()?; + assert_eq!(quoted_scalar_values, vec![1, 2]); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_select_object_content_csv_limit() -> TestResult<()> { let (_env, client) = create_test_environment().await?; diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 5134a43a7..20d165623 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -17,6 +17,7 @@ use crate::{ SelectObjectOptions, SelectObjectSnapshot, SelectObjectSnapshotReadError, SelectStorageError, SelectStore, SnapshotConsistencyError, query::{ + ast::{JsonPathSegment, JsonSource}, parser::RustFsDialect, session::{QueryExecutionGuard, QueryExecutionTracker}, }, @@ -28,14 +29,17 @@ use bytes::Bytes; use chrono::{DateTime, Utc}; use datafusion::{ common::{DataFusionError, runtime::SpawnedTask}, - execution::memory_pool::{MemoryConsumer, MemoryPool, UnboundedMemoryPool}, + execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool}, object_store::{ Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path, }, - sql::sqlparser::{ - ast::{ObjectNamePart, SetExpr, Statement, TableFactor}, - parser::Parser as SqlParser, + sql::{ + planner::IdentNormalizer, + sqlparser::{ + ast::{Expr, Ident, JsonPathElem, ObjectNamePart, SetExpr, Statement, TableFactor}, + parser::Parser as SqlParser, + }, }, }; use futures::pin_mut; @@ -78,11 +82,13 @@ fn select_default_read_buffer_size_u64() -> u64 { /// Default: 128 MiB. This matches the AWS S3 Select limit for JSON DOCUMENT /// 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. +/// limit, JSON DOCUMENT inputs larger than 1 MiB are rejected; scalar source +/// aliases reserve additional space for their maximum per-row expansion. +/// 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; +const JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER: usize = 14; 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."; const NORMALIZED_RECORD_DELIMITER: &[u8] = b"\r\n"; @@ -96,10 +102,8 @@ pub struct EcObjectStore { /// In that case the raw bytes are buffered and flattened to NDJSON before /// being handed to DataFusion's Arrow JSON reader. is_json_document: bool, - /// Optional JSON sub-path extracted from `FROM s3object.` in the SQL - /// expression. When set, `flatten_json_document_to_ndjson` navigates to - /// this key in the root JSON object before flattening. - json_sub_path: Option, + /// JSON source path produced by the SQL compatibility analyzer. + json_source: JsonSource, input_metrics: Arc, memory_pool: Arc, query_tracker: Option, @@ -174,48 +178,54 @@ pub struct InvalidScanRange; impl EcObjectStore { pub fn new(input: Arc) -> S3Result { + let source = legacy_json_source_from_input(&input); Self::build_lazy( input, Arc::new(UnboundedMemoryPool::default()), None, Arc::new(SelectInputMetrics::default()), + source, ) .map_err(map_build_error_to_s3) } pub fn new_with_snapshot(input: Arc, snapshot: Arc) -> S3Result { + let source = legacy_json_source_from_input(&input); Self::build_with_snapshot( input, Arc::new(UnboundedMemoryPool::default()), None, Arc::new(SelectInputMetrics::default()), snapshot, + source, ) .map_err(map_build_error_to_s3) } - pub(crate) fn new_with_memory_pool( + pub(crate) fn new_with_memory_pool_and_source( input: Arc, memory_pool: Arc, input_metrics: Arc, snapshot: Option>, + source: JsonSource, ) -> std::result::Result { match snapshot { - Some(snapshot) => Self::build_with_snapshot(input, memory_pool, None, input_metrics, snapshot), - None => Self::build_lazy(input, memory_pool, None, input_metrics), + Some(snapshot) => Self::build_with_snapshot(input, memory_pool, None, input_metrics, snapshot, source), + None => Self::build_lazy(input, memory_pool, None, input_metrics, source), } } - pub(crate) fn new_with_query_tracker( + pub(crate) fn new_with_query_tracker_and_source( input: Arc, memory_pool: Arc, query_tracker: QueryExecutionTracker, input_metrics: Arc, snapshot: Option>, + source: JsonSource, ) -> std::result::Result { match snapshot { - Some(snapshot) => Self::build_with_snapshot(input, memory_pool, Some(query_tracker), input_metrics, snapshot), - None => Self::build_lazy(input, memory_pool, Some(query_tracker), input_metrics), + Some(snapshot) => Self::build_with_snapshot(input, memory_pool, Some(query_tracker), input_metrics, snapshot, source), + None => Self::build_lazy(input, memory_pool, Some(query_tracker), input_metrics, source), } } @@ -224,9 +234,10 @@ impl EcObjectStore { memory_pool: Arc, query_tracker: Option, input_metrics: Arc, + source: JsonSource, ) -> std::result::Result { let store = resolve_select_object_store_handle().ok_or(EcObjectStoreBuildError::StoreUnavailable)?; - Ok(Self::build(input, memory_pool, query_tracker, input_metrics, Some(store), None)) + Ok(Self::build(input, memory_pool, query_tracker, input_metrics, Some(store), None, source)) } fn build_with_snapshot( @@ -235,11 +246,20 @@ impl EcObjectStore { query_tracker: Option, input_metrics: Arc, snapshot: Arc, + source: JsonSource, ) -> std::result::Result { if !snapshot.is_for(&input.bucket, &input.key) { return Err(EcObjectStoreBuildError::Snapshot(SnapshotConsistencyError::ObjectChanged)); } - Ok(Self::build(input, memory_pool, query_tracker, input_metrics, None, Some(snapshot))) + Ok(Self::build( + input, + memory_pool, + query_tracker, + input_metrics, + None, + Some(snapshot), + source, + )) } fn build( @@ -249,6 +269,7 @@ impl EcObjectStore { input_metrics: Arc, store: Option>, snapshot: Option>, + source: JsonSource, ) -> Self { let (need_convert, delimiter) = if let Some(csv) = input.request.input_serialization.csv.as_ref() { if let Some(delimiter) = csv.field_delimiter.as_ref() { @@ -266,29 +287,14 @@ impl EcObjectStore { // Detect JSON DOCUMENT type: the entire file is a single (possibly // multi-line) JSON object/array, NOT newline-delimited JSON. - let is_json_document = input - .request - .input_serialization - .json - .as_ref() - .and_then(|j| j.type_.as_ref()) - .map(|t| t.as_str() == "DOCUMENT") - .unwrap_or(false); - - // Extract the JSON sub-path from the SQL expression, e.g. - // `SELECT … FROM s3object.employees e` → `Some("employees")`. - let json_sub_path = if is_json_document { - extract_json_sub_path_from_expression(&input.request.expression) - } else { - None - }; + let is_json_document = is_json_document_input(&input); Self { input, need_convert, delimiter, is_json_document, - json_sub_path, + json_source: source, input_metrics, memory_pool, query_tracker, @@ -453,6 +459,86 @@ impl EcObjectStore { } } +pub(crate) fn is_json_document_input(input: &SelectObjectContentInput) -> bool { + input + .request + .input_serialization + .json + .as_ref() + .and_then(|json| json.type_.as_ref()) + .is_some_and(|json_type| json_type.as_str() == "DOCUMENT") +} + +/// Preserves the pre-typed-path single-key behavior of public legacy constructors. +pub(crate) fn legacy_json_source_from_input(input: &SelectObjectContentInput) -> JsonSource { + if !is_json_document_input(input) { + return JsonSource::default(); + } + let Ok(mut statements) = SqlParser::parse_sql(&RustFsDialect, &input.request.expression) else { + return JsonSource::default(); + }; + if statements.len() != 1 { + return JsonSource::default(); + } + let Some(Statement::Query(query)) = statements.pop() else { + return JsonSource::default(); + }; + let SetExpr::Select(select) = query.body.as_ref() else { + return JsonSource::default(); + }; + let [table] = select.from.as_slice() else { + return JsonSource::default(); + }; + let TableFactor::Table { + name, alias, json_path, .. + } = &table.relation + else { + return JsonSource::default(); + }; + let Some(ObjectNamePart::Identifier(table_name)) = name.0.first() else { + return JsonSource::default(); + }; + if name.0.len() > 2 { + return JsonSource::default(); + } + 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 JsonSource::default(); + } + let path = match (name.0.get(1), json_path.as_ref()) { + (Some(ObjectNamePart::Identifier(sub_path)), None) => vec![JsonPathSegment::Key { + name: sub_path.value.clone(), + quoted: sub_path.quote_style.is_some(), + }], + (None, None) => Vec::new(), + (None, Some(json_path)) if matches!(json_path.path.as_slice(), [JsonPathElem::Bracket { key: Expr::Wildcard(_) }]) => { + vec![JsonPathSegment::ArrayWildcard] + } + _ => return JsonSource::default(), + }; + let scalar_column = alias + .as_ref() + .map(|alias| IdentNormalizer::default().normalize(alias.name.clone())) + .or_else(|| match path.as_slice() { + [] => Some(IdentNormalizer::default().normalize(table_name.clone())), + [JsonPathSegment::Key { name, quoted }] => { + let alias = if *quoted { + Ident::with_quote('"', name) + } else { + Ident::new(name) + }; + Some(IdentNormalizer::default().normalize(alias)) + } + [JsonPathSegment::ArrayWildcard] => Some(IdentNormalizer::default().normalize(Ident::new("_1"))), + _ => None, + }); + JsonSource::new(path, scalar_column) +} + impl std::fmt::Debug for EcObjectStore { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("EcObjectStore") @@ -460,7 +546,7 @@ impl std::fmt::Debug for EcObjectStore { .field("object", &self.input.key) .field("need_convert", &self.need_convert) .field("is_json_document", &self.is_json_document) - .field("json_sub_path", &self.json_sub_path) + .field("json_source", &self.json_source) .finish_non_exhaustive() } } @@ -755,7 +841,7 @@ impl ObjectStore for EcObjectStore { let stream = json_document_ndjson_stream( reader.stream, original_size, - self.json_sub_path.clone(), + self.json_source.clone(), Arc::clone(&self.input_metrics), Arc::clone(&self.memory_pool), self.query_tracker.clone(), @@ -1097,34 +1183,6 @@ impl ScanRangeState { } } -fn extract_json_sub_path_from_expression(expression: &str) -> Option { - let mut statements = SqlParser::parse_sql(&RustFsDialect, expression).ok()?; - if statements.len() != 1 { - return None; - } - let Statement::Query(query) = statements.pop()? else { - return None; - }; - let SetExpr::Select(select) = query.body.as_ref() else { - return None; - }; - let [table] = select.from.as_slice() else { - return None; - }; - let TableFactor::Table { name, .. } = &table.relation else { - return None; - }; - let [ObjectNamePart::Identifier(table_name), ObjectNamePart::Identifier(sub_path)] = name.0.as_slice() else { - return None; - }; - let is_s3_object = if table_name.quote_style.is_some() { - table_name.value == "S3Object" - } else { - table_name.value.eq_ignore_ascii_case("S3Object") - }; - is_s3_object.then(|| sub_path.value.clone()) -} - /// Build a lazy NDJSON stream from a JSON DOCUMENT reader. /// /// `get_opts` calls this and returns immediately – no I/O is performed until @@ -1141,7 +1199,7 @@ fn extract_json_sub_path_from_expression(expression: &str) -> Option { fn json_document_ndjson_stream( stream: Box, original_size: u64, - json_sub_path: Option, + json_source: JsonSource, input_metrics: Arc, memory_pool: Arc, query_tracker: Option, @@ -1149,25 +1207,25 @@ fn json_document_ndjson_stream( json_document_ndjson_stream_with_parser( stream, original_size, - json_sub_path, + json_source, input_metrics, memory_pool, query_tracker, - |all_bytes, json_sub_path| parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()), + |all_bytes, json_source| parse_json_document_to_lines(&all_bytes, &json_source), ) } fn json_document_ndjson_stream_with_parser

( stream: Box, original_size: u64, - json_sub_path: Option, + json_source: JsonSource, input_metrics: Arc, memory_pool: Arc, query_tracker: Option, parser: P, ) -> futures_core::stream::BoxStream<'static, Result> where - P: FnOnce(Vec, Option) -> std::io::Result> + Send + 'static, + P: FnOnce(Vec, JsonSource) -> std::io::Result> + Send + 'static, { AsyncTryStream::::new(|mut y| async move { // Compact JSON can expand substantially into a serde_json DOM and @@ -1179,13 +1237,10 @@ where "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 { + let reservation_bytes = + json_document_memory_reservation_bytes(buffer_capacity, &json_source).map_err(|source| o_Error::Generic { store: "EcObjectStore", - source: Box::new(DataFusionError::ResourcesExhausted(format!( - "JSON DOCUMENT memory reservation overflow for {original_size} input bytes" - ))), + source: Box::new(source), })?; let reservation = MemoryConsumer::new("S3 Select JSON document").register(&memory_pool); reservation.try_resize(reservation_bytes).map_err(|err| o_Error::Generic { @@ -1207,11 +1262,24 @@ where } // ── 2. Parse phase (blocking thread pool, non-blocking runtime) ── + let queued_query_guard = match query_tracker.as_ref() { + Some(query_tracker) => Some(query_tracker.query_guard().ok_or_else(|| o_Error::Generic { + store: "EcObjectStore", + source: Box::new(json_document_parse_interrupted_error()), + })?), + None => None, + }; + let task_resources = JsonDocumentTaskResources { + _reservation: reservation, + query_guard: queued_query_guard, + }; 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 (lines, _task_resources) = SpawnedTask::spawn_blocking(move || { let query_guard = PendingQueryExecutionGuard::start(&task_query_guard)?; - parser(all_bytes, json_sub_path).map(|lines| (lines, reservation, query_guard)) + let mut task_resources = task_resources; + task_resources.query_guard = query_guard; + parser(all_bytes, json_source).map(|lines| (lines, task_resources)) }) .await .map_err(|e| o_Error::Generic { @@ -1220,11 +1288,7 @@ where })? .map_err(|e| o_Error::Generic { store: "EcObjectStore", - source: if e.kind() == std::io::ErrorKind::InvalidData { - Box::new(SelectError::JsonParsingError) - } else { - Box::new(e) - }, + source: classify_json_document_parse_error(e), })?; // ── 3. Yield phase (one Bytes per NDJSON line) ─────────────────── @@ -1236,6 +1300,45 @@ where .boxed() } +struct JsonDocumentTaskResources { + // Struct fields drop in declaration order, so admission covers the reservation through teardown. + _reservation: MemoryReservation, + query_guard: Option, +} + +fn json_document_memory_reservation_bytes(input_bytes: usize, json_source: &JsonSource) -> datafusion::common::Result { + let base = input_bytes + .checked_mul(JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER) + .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))?; + let scalar_column = json_source.scalar_column().unwrap_or_default(); + // A scalar row holds one key copy and its JSON encoding. One source byte + // can expand to six escaped bytes, and the serializer buffer can grow to + // twice its final length. + let scalar_column_per_row = scalar_column + .len() + .checked_mul(JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER) + .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))?; + let scalar_column_max = scalar_column_per_row + .checked_mul(input_bytes) + .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))?; + base.checked_add(scalar_column_max) + .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes)) +} + +fn json_document_memory_reservation_overflow(input_bytes: usize) -> DataFusionError { + DataFusionError::ResourcesExhausted(format!("JSON DOCUMENT memory reservation overflow for {input_bytes} input bytes")) +} + +fn classify_json_document_parse_error(error: std::io::Error) -> Box { + if let Some(select_error) = error.get_ref().and_then(|source| source.downcast_ref::()) { + Box::new(select_error.clone()) + } else if error.kind() == std::io::ErrorKind::InvalidData { + Box::new(SelectError::JsonParsingError) + } else { + Box::new(error) + } +} + struct PendingQueryExecutionGuard { state: Arc>, } @@ -1261,15 +1364,13 @@ impl PendingQueryExecutionGuard { 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::Pending(Some(query_tracker)) => query_tracker + .query_guard() + .map(Some) + .ok_or_else(json_document_parse_interrupted_error), QueryExecutionGuardState::Cancelled => { *state = QueryExecutionGuardState::Cancelled; - Err(std::io::Error::new( - std::io::ErrorKind::Interrupted, - "JSON DOCUMENT parse was cancelled before it started", - )) + Err(json_document_parse_interrupted_error()) } QueryExecutionGuardState::Started => { *state = QueryExecutionGuardState::Started; @@ -1279,6 +1380,10 @@ impl PendingQueryExecutionGuard { } } +fn json_document_parse_interrupted_error() -> std::io::Error { + std::io::Error::new(std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started") +} + impl Drop for PendingQueryExecutionGuard { fn drop(&mut self) { let query_guard = { @@ -1299,45 +1404,129 @@ impl Drop for PendingQueryExecutionGuard { /// Parse a JSON DOCUMENT (a single JSON value, possibly multi-line) into a /// list of NDJSON lines – one [`Bytes`] per record. /// -/// `json_sub_path` – when the SQL expression contains `FROM s3object.`, -/// pass `Some(key)` to navigate into that key before flattening. For -/// example, given `{"employees":[{…},{…}]}` and `json_sub_path = -/// Some("employees")`, each element of the `employees` array becomes one -/// NDJSON line. +/// `json_source` is produced from the SQL AST and expands nested source +/// arrays before DataFusion infers the table schema. /// /// - A JSON array → one line per element. -/// - A JSON object (no sub-path match, or scalar root) → one line. -fn parse_json_document_to_lines(bytes: &[u8], json_sub_path: Option<&str>) -> std::io::Result> { +/// - A JSON object or scalar root → one line. +fn parse_json_document_to_lines(bytes: &[u8], json_source: &JsonSource) -> std::io::Result> { let root: serde_json::Value = serde_json::from_slice(bytes).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - - // 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 = match (root, json_sub_path) { - (serde_json::Value::Object(mut object), Some(path)) => { - object.remove(path).unwrap_or_else(|| serde_json::Value::Object(object)) - } - (root, _) => root, - }; - + let json_source_path = json_source.path(); + let values = expand_json_source(root, json_source_path)?; + // Preserve the two pre-path-AST forms that flattened arrays implicitly. + // Explicit indexes and wildcards already identify the intended records + // and must not flatten an array-valued result a second time. + let implicitly_expand_arrays = matches!(json_source_path, [] | [JsonPathSegment::Key { .. }]); + let scalar_column = json_source.scalar_column().unwrap_or_else(|| match json_source_path.last() { + Some(JsonPathSegment::Key { name, .. }) => name, + Some(JsonPathSegment::Index(_) | JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard) | None => "_1", + }); let mut lines: Vec = Vec::new(); - match value { - serde_json::Value::Array(arr) => { - for item in arr { - let mut line = serde_json::to_vec(&item).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - line.push(b'\n'); - lines.push(Bytes::from(line)); + for value in values { + match value { + serde_json::Value::Array(array) if implicitly_expand_arrays => { + for item in array { + lines.push(json_value_to_line(item, scalar_column)?); + } } - } - other => { - let mut line = serde_json::to_vec(&other).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - line.push(b'\n'); - lines.push(Bytes::from(line)); + other => lines.push(json_value_to_line(other, scalar_column)?), } } Ok(lines) } +fn expand_json_source(root: serde_json::Value, json_source_path: &[JsonPathSegment]) -> std::io::Result> { + // S3Object[*] identifies the input record stream. JSON DOCUMENT already + // presents the root value as that stream, so the leading marker is not a + // lookup against the root object. + let (path, mut values) = match json_source_path.strip_prefix(&[JsonPathSegment::ArrayWildcard]) { + // Preserve RustFS's existing S3Object[*] root-array expansion while + // also allowing the AWS canonical S3Object[*][*] form. + Some(path) => match (root, path.first()) { + (root @ serde_json::Value::Array(_), Some(JsonPathSegment::ArrayWildcard | JsonPathSegment::Index(_))) => { + (path, vec![root]) + } + (serde_json::Value::Array(array), _) => (path, array), + (root, _) => (path, vec![root]), + }, + None => (json_source_path, vec![root]), + }; + + for segment in path { + let mut expanded = Vec::new(); + for value in values { + match (segment, value) { + (JsonPathSegment::Key { name, quoted }, serde_json::Value::Object(mut object)) => { + if let Some(value) = remove_json_source_key(&mut object, name, *quoted)? { + expanded.push(value); + } + } + (JsonPathSegment::Index(index), serde_json::Value::Array(array)) => { + if let Some(value) = array.into_iter().nth(*index) { + expanded.push(value); + } + } + (JsonPathSegment::ArrayWildcard, serde_json::Value::Array(mut array)) => { + if expanded.is_empty() { + expanded = array; + } else { + expanded.append(&mut array); + } + } + (JsonPathSegment::ObjectWildcard, serde_json::Value::Object(object)) => { + expanded.extend(object.into_values()); + } + (JsonPathSegment::Key { .. }, _) + | (JsonPathSegment::Index(_), _) + | (JsonPathSegment::ArrayWildcard, _) + | (JsonPathSegment::ObjectWildcard, _) => { + return Err(invalid_json_source_path("JSON source path segment does not match the input value")); + } + } + } + values = expanded; + } + + Ok(values) +} + +fn remove_json_source_key( + object: &mut serde_json::Map, + name: &str, + quoted: bool, +) -> std::io::Result> { + if quoted { + return Ok(object.remove(name)); + } + + let mut matches = object.keys().filter(|key| key.eq_ignore_ascii_case(name)); + let matched = matches.next().cloned(); + if matches.next().is_some() { + return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::AmbiguousFieldName)); + } + drop(matches); + Ok(matched.and_then(|key| object.remove(&key))) +} + +fn json_value_to_line(value: serde_json::Value, scalar_column: &str) -> std::io::Result { + let value = match value { + value @ serde_json::Value::Object(_) => value, + value => { + let mut row = serde_json::Map::new(); + row.insert(scalar_column.to_string(), value); + serde_json::Value::Object(row) + } + }; + let mut line = serde_json::to_vec(&value).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + line.push(b'\n'); + Ok(Bytes::from(line)) +} + +fn invalid_json_source_path(message: &'static str) -> std::io::Error { + std::io::Error::new(std::io::ErrorKind::InvalidData, message) +} + /// Convert a JSON DOCUMENT to a single concatenated NDJSON [`Bytes`] blob. /// /// This is a convenience wrapper around [`parse_json_document_to_lines`] used @@ -1345,8 +1534,8 @@ fn parse_json_document_to_lines(bytes: &[u8], json_sub_path: Option<&str>) -> st /// instead, which streams lines lazily without constructing this intermediate /// blob. #[cfg(test)] -fn flatten_json_document_to_ndjson(bytes: &[u8], json_sub_path: Option<&str>) -> std::io::Result { - let lines = parse_json_document_to_lines(bytes, json_sub_path)?; +fn flatten_json_document_to_ndjson(bytes: &[u8], json_source_path: &[JsonPathSegment]) -> std::io::Result { + let lines = parse_json_document_to_lines(bytes, &JsonSource::from_path(json_source_path.to_vec()))?; let total = lines.iter().map(|b| b.len()).sum(); let mut output = Vec::with_capacity(total); for line in lines { @@ -1425,10 +1614,12 @@ mod test { EcObjectStore, EcObjectStoreBuildError, JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER, OnceCell, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectObjectOptions, SelectObjectSnapshot, SelectScanRange, SnapshotConsistencyError, bytes_stream, convert_csv_delimiter_stream, convert_field_delimiter_stream, convert_record_delimiter_stream, - extract_json_sub_path_from_expression, find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, - json_document_ndjson_stream, json_document_ndjson_stream_with_parser, map_storage_error, meter_uncompressed_input_stream, - scan_range_from_bounds, scan_range_stream, select_read_headers, snapshot_last_modified, validate_json_document_size, + find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, json_document_ndjson_stream, + json_document_ndjson_stream_with_parser, legacy_json_source_from_input, map_storage_error, + meter_uncompressed_input_stream, scan_range_from_bounds, scan_range_stream, select_read_headers, snapshot_last_modified, + validate_json_document_size, }; + use crate::query::ast::{JsonPathSegment, JsonSource}; use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker}; use crate::storage_api::SelectPutObjReader; use crate::storage_api::object_store::ObjectIO as _; @@ -1436,7 +1627,7 @@ mod test { use bytes::Bytes; use datafusion::{ common::DataFusionError, - execution::memory_pool::{GreedyMemoryPool, MemoryPool}, + execution::memory_pool::{GreedyMemoryPool, MemoryLimit, MemoryPool, MemoryReservation}, execution::{config::SessionConfig, context::SessionContext}, object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path}, physical_plan::ExecutionPlanProperties, @@ -1444,6 +1635,7 @@ mod test { }; use futures::{StreamExt, TryStreamExt, stream}; use http::HeaderMap; + use parking_lot::Mutex; use rustfs_test_utils::PutObjectCommitBarrier; use s3s::S3ErrorCode; use s3s::dto::{ @@ -1461,6 +1653,78 @@ mod test { use tokio::{io::AsyncReadExt, sync::Semaphore}; + #[derive(Debug)] + struct AdmissionObservingMemoryPool { + inner: GreedyMemoryPool, + admission: Arc, + reservation_release: Mutex>>, + } + + impl AdmissionObservingMemoryPool { + fn new(pool_size: usize, admission: Arc) -> (Self, tokio::sync::oneshot::Receiver) { + let (reservation_release, reservation_released) = tokio::sync::oneshot::channel(); + ( + Self { + inner: GreedyMemoryPool::new(pool_size), + admission, + reservation_release: Mutex::new(Some(reservation_release)), + }, + reservation_released, + ) + } + } + + impl std::fmt::Display for AdmissionObservingMemoryPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + std::fmt::Display::fmt(&self.inner, f) + } + } + + impl MemoryPool for AdmissionObservingMemoryPool { + fn name(&self) -> &str { + self.inner.name() + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + if self.inner.reserved() == 0 + && let Some(reservation_release) = self.reservation_release.lock().take() + { + let _ = reservation_release.send(self.admission.available_permits() == 0); + } + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> datafusion::common::Result<()> { + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } + } + + fn source_key(name: &str) -> JsonPathSegment { + JsonPathSegment::Key { + name: name.to_string(), + quoted: false, + } + } + + fn quoted_source_key(name: &str) -> JsonPathSegment { + JsonPathSegment::Key { + name: name.to_string(), + quoted: true, + } + } + fn csv_input(bucket: &str, object: &str) -> Arc { Arc::new(SelectObjectContentInput { bucket: bucket.to_string(), @@ -1486,6 +1750,76 @@ mod test { }) } + fn json_document_input(bucket: &str, object: &str, expression: &str) -> Arc { + Arc::new(SelectObjectContentInput { + bucket: bucket.to_string(), + expected_bucket_owner: None, + key: object.to_string(), + sse_customer_algorithm: None, + sse_customer_key: None, + sse_customer_key_md5: None, + request: SelectObjectContentRequest { + expression: expression.to_string(), + expression_type: ExpressionType::from_static(ExpressionType::SQL), + input_serialization: InputSerialization { + json: Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::DOCUMENT)), + }), + ..Default::default() + }, + output_serialization: OutputSerialization { + csv: Some(CSVOutput::default()), + ..Default::default() + }, + request_progress: None, + scan_range: None, + }, + }) + } + + #[test] + fn legacy_source_adapter_normalizes_scalar_bindings() { + let input = json_document_input("bucket", "input.json", "SELECT V FROM S3Object AS V"); + let implicit = json_document_input("bucket", "input.json", "SELECT S3Object FROM S3Object"); + let wildcard = json_document_input("bucket", "input.json", "SELECT _1 FROM S3Object[*]"); + let key = json_document_input("bucket", "input.json", "SELECT * FROM S3Object.LongKey"); + let quoted_key = json_document_input("bucket", "input.json", "SELECT * FROM S3Object.\"LongKey\""); + + assert_eq!(legacy_json_source_from_input(&input), JsonSource::new(Vec::new(), Some("v".to_string()))); + assert_eq!( + legacy_json_source_from_input(&implicit), + JsonSource::new(Vec::new(), Some("s3object".to_string())) + ); + assert_eq!( + legacy_json_source_from_input(&wildcard), + JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some("_1".to_string())) + ); + assert_eq!(legacy_json_source_from_input(&key).scalar_column(), Some("longkey")); + assert_eq!(legacy_json_source_from_input("ed_key).scalar_column(), Some("LongKey")); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn legacy_public_constructors_preserve_single_key_json_source() { + const BUCKET: &str = "s3select-legacy-json-source"; + const OBJECT: &str = "input.json"; + + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + env.put_object_bytes(BUCKET, OBJECT, br#"{"employees":[{"name":"Alice"}]}"#.to_vec()) + .await; + let input = json_document_input(BUCKET, OBJECT, "SELECT e.name FROM S3Object.employees AS e"); + let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await; + + let lazy = EcObjectStore::new(Arc::clone(&input)).expect("legacy lazy constructor should resolve the global store"); + let pinned = + EcObjectStore::new_with_snapshot(input, snapshot).expect("legacy pinned constructor should accept the snapshot"); + let expected = JsonSource::new(vec![source_key("employees")], Some("e".to_string())); + + assert_eq!(lazy.json_source, expected); + assert_eq!(pinned.json_source, expected); + } + fn json_input(bucket: &str, object: &str, json_type: &'static str) -> Arc { let mut input = (*csv_input(bucket, object)).clone(); input.request.input_serialization = InputSerialization { @@ -2386,7 +2720,7 @@ mod test { need_convert: false, delimiter: String::new(), is_json_document: false, - json_sub_path: None, + json_source: JsonSource::default(), input_metrics: Arc::new(SelectInputMetrics::default()), memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, @@ -2573,7 +2907,7 @@ mod test { need_convert: false, delimiter: String::new(), is_json_document: false, - json_sub_path: None, + json_source: JsonSource::default(), input_metrics: Arc::new(SelectInputMetrics::default()), memory_pool: Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)), query_tracker: None, @@ -2671,7 +3005,7 @@ mod test { need_convert: true, delimiter: "\r\n".to_string(), is_json_document: false, - json_sub_path: None, + json_source: JsonSource::default(), input_metrics: Arc::clone(&input_metrics), memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, @@ -2795,6 +3129,7 @@ mod test { None, Arc::clone(&input_metrics), snapshot, + JsonSource::default(), ) .expect("build metrics-aware object store"); @@ -2861,6 +3196,7 @@ mod test { None, Arc::clone(&input_metrics), snapshot, + JsonSource::default(), ) .expect("build metrics-aware object store"); @@ -2911,6 +3247,7 @@ mod test { None, Arc::clone(&input_metrics), snapshot, + JsonSource::default(), ) .expect("build metrics-aware JSON object store"); @@ -2965,6 +3302,7 @@ mod test { None, Arc::clone(&input_metrics), snapshot, + JsonSource::default(), ) .expect("build ScanRange metrics-aware object store"); @@ -2995,7 +3333,7 @@ mod test { let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, - None, + JsonSource::default(), Arc::new(SelectInputMetrics::default()), memory_pool, None, @@ -3015,6 +3353,54 @@ mod test { )); } + #[tokio::test] + async fn scalar_alias_expansion_is_in_the_query_memory_reservation() { + let input = b"[0,0]".to_vec(); + let alias = "alias".repeat(128); + let source = JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some(alias.clone())); + // Keep this threshold independent from the production helper so a + // smaller scalar-alias multiplier cannot make the test self-validate. + let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER + alias.len() * 14 * input.len(); + assert!(required > input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER); + let memory_pool: Arc = 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, + source, + Arc::new(SelectInputMetrics::default()), + memory_pool, + None, + ); + + let err = output + .next() + .await + .expect("memory error") + .expect_err("scalar alias expansion must be reserved before parsing"); + let object_store::Error::Generic { source, .. } = err else { + panic!("expected generic object store error"); + }; + assert!(matches!( + source.downcast_ref::(), + Some(DataFusionError::ResourcesExhausted(_)) + )); + + let memory_pool = Arc::new(GreedyMemoryPool::new(required)); + let output: Vec = json_document_ndjson_stream( + Box::new(std::io::Cursor::new(input.clone())), + input.len() as u64, + JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some(alias)), + Arc::new(SelectInputMetrics::default()), + memory_pool.clone(), + None, + ) + .try_collect() + .await + .expect("scalar alias expansion should fit the exact reservation"); + assert_eq!(output.len(), 2); + assert_eq!(memory_pool.reserved(), 0); + } + #[tokio::test] async fn test_json_document_stream_releases_memory_reservation() { let input = b"[1,2]".to_vec(); @@ -3024,7 +3410,7 @@ mod test { let output: Vec = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, - None, + JsonSource::default(), Arc::clone(&input_metrics), memory_pool.clone(), None, @@ -3036,7 +3422,7 @@ mod test { assert_eq!(input_metrics.snapshot().bytes_scanned, input_len); assert_eq!(input_metrics.snapshot().bytes_processed, input_len); - assert_eq!(output, vec![Bytes::from_static(b"1\n"), Bytes::from_static(b"2\n")]); + assert_eq!(output, vec![Bytes::from_static(b"{\"_1\":1}\n"), Bytes::from_static(b"{\"_1\":2}\n")]); assert_eq!(memory_pool.reserved(), 0); } @@ -3049,7 +3435,7 @@ mod test { let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input)), 4, - None, + JsonSource::default(), Arc::clone(&input_metrics), memory_pool, None, @@ -3079,7 +3465,7 @@ mod test { let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, - None, + JsonSource::default(), Arc::new(SelectInputMetrics::default()), memory_pool, None, @@ -3096,6 +3482,32 @@ mod test { assert!(output.next().await.is_none()); } + #[tokio::test] + async fn json_document_stream_preserves_typed_parser_error() { + let input = b"{}".to_vec(); + let memory_pool: Arc = + Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let mut output = json_document_ndjson_stream_with_parser( + Box::new(std::io::Cursor::new(input.clone())), + input.len() as u64, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + memory_pool, + None, + |_, _| Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::AmbiguousFieldName)), + ); + + let source = output + .next() + .await + .expect("typed parser failure should produce one stream error") + .expect_err("typed parser failure must fail the stream"); + let error = QueryError::from(DataFusionError::ObjectStore(Box::new(source))); + + assert_eq!(error.select_error(), SelectError::AmbiguousFieldName); + assert!(output.next().await.is_none()); + } + #[test] fn storage_error_mapper_preserves_protocol_classification() { let classify = |source| QueryError::from(DataFusionError::ObjectStore(Box::new(source))).select_error(); @@ -3150,7 +3562,7 @@ mod test { } #[test] - fn test_json_document_queued_parse_releases_query_guard_when_cancelled() { + fn test_json_document_queued_parse_retains_query_guard_until_dequeued() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .max_blocking_threads(1) @@ -3180,14 +3592,18 @@ mod test { 30, ); let input = b"{}".to_vec(); - let memory_pool: Arc = - Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let (memory_pool, reservation_released) = AdmissionObservingMemoryPool::new( + input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER, + Arc::clone(&admission), + ); + let memory_pool = Arc::new(memory_pool); + let query_memory_pool: Arc = memory_pool.clone(); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, - None, + JsonSource::default(), Arc::new(SelectInputMetrics::default()), - memory_pool, + query_memory_pool, Some(query_tracker), ); @@ -3198,13 +3614,29 @@ mod test { } drop(output); + assert_eq!(admission.available_permits(), 0); + assert!(memory_pool.reserved() > 0); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), Arc::clone(&admission).acquire_owned()) + .await + .is_err(), + "queued parse resources must remain covered by admission" + ); + release_blocking_tx.send(()).expect("release blocking worker"); + blocker.await.expect("blocking worker should finish"); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(5), reservation_released) + .await + .expect("cancelled JSON parse should release its memory reservation") + .expect("memory reservation release observer should remain open"), + "query admission must cover the memory reservation through task teardown" + ); 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("cancelled JSON parse should be dequeued") .expect("query admission should remain open"); - release_blocking_tx.send(()).expect("release blocking worker"); - blocker.await.expect("blocking worker should finish"); + assert_eq!(memory_pool.reserved(), 0); drop(recovered_permit); assert_eq!(admission.available_permits(), 1); }); @@ -3248,9 +3680,9 @@ mod test { let mut output = json_document_ndjson_stream_with_parser( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, - None, + JsonSource::default(), Arc::new(SelectInputMetrics::default()), - memory_pool, + Arc::clone(&memory_pool), Some(query_tracker.clone()), move |_, _| { parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst); @@ -3264,7 +3696,8 @@ mod test { assert!(futures::poll!(next.as_mut()).is_pending()); } query_tracker.expire(&owner); - assert_eq!(admission.available_permits(), 1); + assert_eq!(admission.available_permits(), 0); + assert!(memory_pool.reserved() > 0); release_blocking_tx.send(()).expect("release blocking worker"); blocker.await.expect("blocking worker should finish"); @@ -3279,6 +3712,77 @@ mod test { let source = source.downcast_ref::().expect("I/O error source"); assert_eq!(source.kind(), std::io::ErrorKind::Interrupted); assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst)); + assert_eq!(memory_pool.reserved(), 0); + assert_eq!(admission.available_permits(), 1); + }); + } + + #[test] + fn test_json_document_expired_before_enqueue_releases_resources_without_blocking() { + 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, + ); + query_tracker.expire(&owner); + + let input = b"{}".to_vec(); + let memory_pool: Arc = + 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, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + Arc::clone(&memory_pool), + Some(query_tracker), + move |_, _| { + parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst); + Ok(vec![Bytes::from_static(b"{}\n")]) + }, + ); + + let err = tokio::time::timeout(std::time::Duration::from_millis(100), output.next()) + .await + .expect("expired parse must fail before entering the saturated blocking queue") + .expect("expired parse should return an error") + .expect_err("expired parse must not run"); + let object_store::Error::Generic { source, .. } = err else { + panic!("expected generic object store error"); + }; + let source = source.downcast_ref::().expect("I/O error source"); + assert_eq!(source.kind(), std::io::ErrorKind::Interrupted); + assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst)); + assert_eq!(memory_pool.reserved(), 0); + assert_eq!(admission.available_permits(), 1); + + release_blocking_tx.send(()).expect("release blocking worker"); + blocker.await.expect("blocking worker should finish"); }); } @@ -3312,7 +3816,7 @@ mod test { let mut output = json_document_ndjson_stream_with_parser( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, - None, + JsonSource::default(), Arc::new(SelectInputMetrics::default()), memory_pool, Some(query_tracker), @@ -3347,7 +3851,7 @@ mod test { #[test] fn test_flatten_array_produces_one_line_per_element() { let input = br#"[{"id":1,"name":"Alice"},{"id":2,"name":"Bob"}]"#; - let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed"); let text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 2); @@ -3365,7 +3869,7 @@ mod test { #[test] fn test_flatten_single_object_produces_one_line() { let input = br#"{"id":42,"value":"hello world"}"#; - let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed"); let text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 1); @@ -3378,7 +3882,7 @@ mod test { #[test] fn test_flatten_empty_array_produces_no_output() { let input = b"[]"; - let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed"); assert!(result.is_empty(), "empty array should yield zero bytes"); } @@ -3386,7 +3890,7 @@ mod test { #[test] fn test_flatten_pretty_printed_document() { let input = b"[\n {\"a\": 1},\n {\"a\": 2},\n {\"a\": 3}\n]"; - let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed"); let text = std::str::from_utf8(&result).unwrap(); assert_eq!(text.lines().count(), 3); } @@ -3395,7 +3899,7 @@ mod test { #[test] fn test_flatten_array_with_nested_objects() { let input = br#"[{"outer":{"inner":99}},{"outer":{"inner":100}}]"#; - let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed"); let text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 2); @@ -3411,7 +3915,7 @@ mod test { #[test] fn test_flatten_output_ends_with_newline_per_record() { let input = br#"[{"x":1},{"x":2}]"#; - let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed"); let text = std::str::from_utf8(&result).unwrap(); // Exactly 2 newlines for 2 records assert_eq!(text.chars().filter(|&c| c == '\n').count(), 2); @@ -3423,25 +3927,21 @@ mod test { #[test] fn test_flatten_invalid_json_returns_error() { let input = b"{ not valid json }"; - let err = flatten_json_document_to_ndjson(input, None).expect_err("should fail on invalid JSON"); + let err = flatten_json_document_to_ndjson(input, &[]).expect_err("should fail on invalid JSON"); assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); } /// Completely empty input returns an error (not valid JSON). #[test] fn test_flatten_empty_input_returns_error() { - let err = flatten_json_document_to_ndjson(b"", None).expect_err("empty bytes are not valid JSON"); + let err = flatten_json_document_to_ndjson(b"", &[]).expect_err("empty bytes are not valid JSON"); assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); } - // ── sub-path navigation tests ───────────────────────────────────────── - - /// `FROM s3object.employees` with a root JSON object navigates into the - /// `employees` array and emits one NDJSON line per element. #[test] - fn test_flatten_sub_path_object_with_array() { + fn key_path_expands_final_array_for_legacy_queries() { let input = br#"{"employees":[{"id":1,"name":"Alice","salary":75000},{"id":2,"name":"Bob","salary":65000}]}"#; - let result = flatten_json_document_to_ndjson(input, Some("employees")).expect("should succeed"); + let result = flatten_json_document_to_ndjson(input, &[source_key("employees")]).expect("key path should succeed"); let text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 2, "each employee should be its own NDJSON line"); @@ -3452,70 +3952,253 @@ mod test { assert_eq!(second["name"], "Bob"); } - /// Sub-path that does not exist in the root object falls back to emitting the - /// entire root object as one NDJSON line (graceful degradation). #[test] - fn test_flatten_sub_path_missing_key_falls_back() { - let input = br#"{"employees":[]}"#; - let result = flatten_json_document_to_ndjson(input, Some("nonexistent")).expect("should succeed"); - let text = std::str::from_utf8(&result).unwrap(); - // Falls back to emitting the whole root object. - assert_eq!(text.lines().count(), 1); - let parsed: serde_json::Value = serde_json::from_str(text.trim_end()).unwrap(); - assert!(parsed.get("employees").is_some(), "root object preserved"); + fn nested_array_wildcards_expand_in_source_order() { + let input = br#"{"departments":[{"employees":[{"id":1},{"id":2}]},{"employees":[{"id":3}]}]}"#; + let path = [ + JsonPathSegment::ArrayWildcard, + source_key("departments"), + JsonPathSegment::ArrayWildcard, + source_key("employees"), + JsonPathSegment::ArrayWildcard, + ]; + + let result = flatten_json_document_to_ndjson(input, &path).expect("nested wildcard path should succeed"); + let ids: Vec = std::str::from_utf8(&result) + .unwrap() + .lines() + .map(|line| { + serde_json::from_str::(line).unwrap()["id"] + .as_i64() + .unwrap() + }) + .collect(); + + assert_eq!(ids, vec![1, 2, 3]); } - /// Sub-path is ignored when the root is already an array. #[test] - fn test_flatten_sub_path_ignored_for_root_array() { + fn explicit_array_wildcard_does_not_expand_nested_array_records_twice() { + let input = br#"{"groups":[[1,2],[3,4]]}"#; + let source = JsonSource::new(vec![source_key("groups"), JsonPathSegment::ArrayWildcard], Some("g".to_string())); + + let result = super::parse_json_document_to_lines(input, &source).expect("explicit wildcard path should succeed"); + + assert_eq!( + result, + vec![Bytes::from_static(b"{\"g\":[1,2]}\n"), Bytes::from_static(b"{\"g\":[3,4]}\n")] + ); + } + + #[test] + fn leading_array_wildcard_expands_a_root_array_before_nested_keys() { + let input = br#"[{"employees":[{"id":1}]},{"employees":[{"id":2}]}]"#; + let path = [ + JsonPathSegment::ArrayWildcard, + source_key("employees"), + JsonPathSegment::ArrayWildcard, + ]; + + let result = flatten_json_document_to_ndjson(input, &path).expect("root array path should succeed"); + let ids: Vec = std::str::from_utf8(&result) + .unwrap() + .lines() + .map(|line| { + serde_json::from_str::(line).unwrap()["id"] + .as_i64() + .unwrap() + }) + .collect(); + + assert_eq!(ids, vec![1, 2]); + } + + #[test] + fn root_array_index_selects_one_record() { let input = br#"[{"id":1},{"id":2}]"#; - let result = flatten_json_document_to_ndjson(input, Some("employees")).expect("should succeed"); - let text = std::str::from_utf8(&result).unwrap(); - // The root array is flattened directly regardless of the sub-path hint. - assert_eq!(text.lines().count(), 2); - } - // ── SQL path extraction tests ───────────────────────────────────────── + let result = flatten_json_document_to_ndjson(input, &[JsonPathSegment::Index(1)]).expect("array index should succeed"); - #[test] - fn test_extract_json_sub_path_basic() { - let sql = "SELECT e.name FROM s3object.employees e WHERE e.salary > 70000"; - assert_eq!(extract_json_sub_path_from_expression(sql), Some("employees".to_string())); + assert_eq!(result, Bytes::from_static(b"{\"id\":2}\n")); } #[test] - fn test_extract_json_sub_path_uppercase() { - let sql = "SELECT s.name FROM S3Object.records s"; - assert_eq!(extract_json_sub_path_from_expression(sql), Some("records".to_string())); + fn canonical_root_array_index_selects_one_record() { + let input = br#"[{"id":1},{"id":2}]"#; + let path = [JsonPathSegment::ArrayWildcard, JsonPathSegment::Index(0)]; + + let result = flatten_json_document_to_ndjson(input, &path).expect("canonical array index should succeed"); + + assert_eq!(result, Bytes::from_static(b"{\"id\":1}\n")); } #[test] - fn test_extract_json_sub_path_no_sub_path() { - let sql = "SELECT * FROM s3object WHERE s3object.age > 30"; - assert_eq!(extract_json_sub_path_from_expression(sql), None); + fn out_of_range_root_array_indexes_produce_no_records() { + let input = br#"[{"id":1}]"#; + + for path in [ + vec![JsonPathSegment::Index(1)], + vec![JsonPathSegment::ArrayWildcard, JsonPathSegment::Index(1)], + ] { + let result = flatten_json_document_to_ndjson(input, &path).expect("out-of-range index should not fail"); + assert!(result.is_empty()); + } } #[test] - fn test_extract_json_sub_path_rejects_unsupported_bracket_path() { - let sql = "SELECT e.name FROM s3object.employees[*] e"; - assert_eq!(extract_json_sub_path_from_expression(sql), None); + fn canonical_key_path_does_not_expand_an_array_without_a_wildcard() { + let input = br#"{"rules":[{"id":1},{"id":2}]}"#; + let source = JsonSource::new(vec![JsonPathSegment::ArrayWildcard, source_key("rules")], Some("r".to_string())); + + let result = super::parse_json_document_to_lines(input, &source).expect("array-valued source path should remain one row"); + + assert_eq!(result, vec![Bytes::from_static(b"{\"r\":[{\"id\":1},{\"id\":2}]}\n")]); } #[test] - fn test_extract_json_sub_path_ignores_from_in_string_literal() { - let sql = "SELECT ' from ' AS marker FROM S3Object.employees"; - assert_eq!(extract_json_sub_path_from_expression(sql), Some("employees".to_string())); + fn object_wildcard_expands_values_and_allows_continuation() { + let input = br#"{"groups":{"first":{"id":1},"second":{"id":2}}}"#; + let path = [ + JsonPathSegment::ArrayWildcard, + source_key("groups"), + JsonPathSegment::ObjectWildcard, + source_key("id"), + ]; + + let result = flatten_json_document_to_ndjson(input, &path).expect("object wildcard should succeed"); + let values: Vec = std::str::from_utf8(&result) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + + assert_eq!(values, vec![serde_json::json!({"id": 1}), serde_json::json!({"id": 2})]); } #[test] - fn test_extract_json_sub_path_ignores_from_in_comment() { - let sql = "SELECT /* from S3Object.wrong */ e.name FROM S3Object.employees AS e"; - assert_eq!(extract_json_sub_path_from_expression(sql), Some("employees".to_string())); + fn canonical_double_wildcard_expands_a_root_array() { + let input = br#"[{"id":1},{"id":2}]"#; + let path = [JsonPathSegment::ArrayWildcard, JsonPathSegment::ArrayWildcard]; + + let result = flatten_json_document_to_ndjson(input, &path).expect("root wildcard should succeed"); + + assert_eq!(result, Bytes::from_static(b"{\"id\":1}\n{\"id\":2}\n")); } #[test] - fn test_extract_json_sub_path_supports_quoted_identifier() { - let sql = "SELECT \" from \" FROM S3Object.\"employee data\""; - assert_eq!(extract_json_sub_path_from_expression(sql), Some("employee data".to_string())); + fn terminal_scalar_path_uses_the_terminal_key_as_its_column() { + let input = br#"{"rules":[{"id":"one"},{"id":"two"}]}"#; + let path = [ + JsonPathSegment::ArrayWildcard, + source_key("rules"), + JsonPathSegment::ArrayWildcard, + source_key("id"), + ]; + + let result = flatten_json_document_to_ndjson(input, &path).expect("scalar source path should succeed"); + + assert_eq!(result, Bytes::from_static(b"{\"id\":\"one\"}\n{\"id\":\"two\"}\n")); + } + + #[test] + fn terminal_scalar_path_uses_the_explicit_source_alias() { + let input = br#"{"rules":[{"id":"one"},{"id":"two"}]}"#; + let source = JsonSource::new( + vec![ + JsonPathSegment::ArrayWildcard, + source_key("rules"), + JsonPathSegment::ArrayWildcard, + source_key("id"), + ], + Some("v".to_string()), + ); + + let result = + super::parse_json_document_to_lines(input, &source).expect("explicit scalar source alias should be preserved"); + + assert_eq!( + result, + vec![ + Bytes::from_static(b"{\"v\":\"one\"}\n"), + Bytes::from_static(b"{\"v\":\"two\"}\n") + ] + ); + } + + #[test] + fn source_keys_follow_s3_case_sensitivity_rules() { + let input = br#"{"Employees":[{"id":1}]}"#; + + let unquoted = flatten_json_document_to_ndjson(input, &[source_key("employees")]) + .expect("unquoted source key should be case insensitive"); + let quoted_exact = flatten_json_document_to_ndjson(input, &[quoted_source_key("Employees")]) + .expect("exact quoted source key should match"); + let quoted = flatten_json_document_to_ndjson(input, &[quoted_source_key("employees")]) + .expect("missing quoted source key should not fail"); + + assert_eq!(unquoted, Bytes::from_static(b"{\"id\":1}\n")); + assert_eq!(quoted_exact, Bytes::from_static(b"{\"id\":1}\n")); + assert!(quoted.is_empty()); + } + + #[test] + fn ambiguous_unquoted_source_key_is_rejected() { + let input = br#"{"Employees":[],"employees":[]}"#; + + let error = + flatten_json_document_to_ndjson(input, &[source_key("EMPLOYEES")]).expect_err("ambiguous source key should fail"); + + assert!( + error + .get_ref() + .and_then(|source| source.downcast_ref::()) + .is_some_and(|error| *error == SelectError::AmbiguousFieldName) + ); + } + + #[test] + fn missing_source_key_produces_no_records() { + let input = br#"{"employees":[]}"#; + + let result = flatten_json_document_to_ndjson(input, &[source_key("nonexistent")]) + .expect("missing source key should not fail the query"); + + assert!(result.is_empty()); + } + + #[tokio::test] + async fn source_path_type_mismatches_are_json_parsing_errors() { + let cases: [(&str, &[u8], Vec); 4] = [ + ("key on array", b"[1]", vec![source_key("id")]), + ("index on object", b"{}", vec![JsonPathSegment::Index(0)]), + ( + "array wildcard on object", + b"{}", + vec![JsonPathSegment::ArrayWildcard, JsonPathSegment::ArrayWildcard], + ), + ("object wildcard on array", b"[1]", vec![JsonPathSegment::ObjectWildcard]), + ]; + + for (case, input, path) in cases { + let memory_pool: Arc = + Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let mut output = json_document_ndjson_stream( + Box::new(std::io::Cursor::new(input.to_vec())), + input.len() as u64, + JsonSource::from_path(path), + Arc::new(SelectInputMetrics::default()), + memory_pool, + None, + ); + let source = output + .next() + .await + .unwrap_or_else(|| panic!("{case} should produce one stream error")) + .unwrap_err(); + let error = QueryError::from(DataFusionError::ObjectStore(Box::new(source))); + + assert_eq!(error.select_error(), SelectError::JsonParsingError, "{case}"); + assert!(output.next().await.is_none(), "{case}"); + } } } diff --git a/crates/s3select-api/src/query/ast.rs b/crates/s3select-api/src/query/ast.rs index af234ed75..d37bf40bd 100644 --- a/crates/s3select-api/src/query/ast.rs +++ b/crates/s3select-api/src/query/ast.rs @@ -13,6 +13,43 @@ // limitations under the License. use datafusion::sql::sqlparser::ast::Statement; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum JsonPathSegment { + Key { name: String, quoted: bool }, + Index(usize), + ArrayWildcard, + ObjectWildcard, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct JsonSource { + path: Arc<[JsonPathSegment]>, + scalar_column: Option, +} + +impl JsonSource { + pub fn new(path: Vec, scalar_column: Option) -> Self { + Self { + path: path.into(), + scalar_column, + } + } + + #[cfg(test)] + pub(crate) fn from_path(path: Vec) -> Self { + Self::new(path, None) + } + + pub fn path(&self) -> &[JsonPathSegment] { + &self.path + } + + pub fn scalar_column(&self) -> Option<&str> { + self.scalar_column.as_deref() + } +} #[derive(Debug, Clone, PartialEq, Eq)] pub enum ExtStatement { diff --git a/crates/s3select-api/src/query/execution.rs b/crates/s3select-api/src/query/execution.rs index 935b415d2..85b5dea89 100644 --- a/crates/s3select-api/src/query/execution.rs +++ b/crates/s3select-api/src/query/execution.rs @@ -29,6 +29,7 @@ use tracing::debug; use crate::{QueryError, QueryResult}; use super::Query; +use super::ast::ExtStatement; use super::logical_planner::Plan; use super::session::{QueryExecutionTracker, SessionCtx}; @@ -172,6 +173,7 @@ pub struct QueryStateMachine { pub session: SessionCtx, pub query: Query, + prepared_statement: Option, query_tracker: Option, state: RwLock, start: Instant, @@ -196,6 +198,7 @@ impl QueryStateMachine { Self { session, query, + prepared_statement: None, query_tracker: None, state: RwLock::new(QueryState::ACCEPTING), start: Instant::now(), @@ -211,6 +214,21 @@ impl QueryStateMachine { Ok(state_machine) } + pub fn begin_tracked_prepared( + query: Query, + session: SessionCtx, + query_tracker: QueryExecutionTracker, + prepared_statement: ExtStatement, + ) -> QueryResult { + let mut state_machine = Self::begin_tracked(query, session, query_tracker)?; + state_machine.prepared_statement = Some(prepared_statement); + Ok(state_machine) + } + + pub fn prepared_statement(&self) -> Option<&ExtStatement> { + self.prepared_statement.as_ref() + } + pub fn query_tracker(&self) -> Option<&QueryExecutionTracker> { self.query_tracker.as_ref() } diff --git a/crates/s3select-api/src/query/parser.rs b/crates/s3select-api/src/query/parser.rs index e4781ca22..1df24b093 100644 --- a/crates/s3select-api/src/query/parser.rs +++ b/crates/s3select-api/src/query/parser.rs @@ -34,6 +34,10 @@ impl Dialect for RustFsDialect { fn supports_group_by_expr(&self) -> bool { true } + + fn supports_partiql(&self) -> bool { + true + } } pub trait Parser { diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index dc5ca8d25..dc3a66dc0 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -12,9 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::query::{Context, Query}; -use crate::{QueryError, QueryResult, object_store::EcObjectStore}; -use crate::{SelectInputMetrics, SelectObjectSnapshot}; +use crate::query::{Context, Query, ast::JsonSource}; +use crate::{ + QueryError, QueryResult, SelectInputMetrics, SelectObjectSnapshot, + object_store::{EcObjectStore, is_json_document_input, legacy_json_source_from_input}, +}; use datafusion::{ arrow::{ array::{Int32Array, StringArray}, @@ -314,8 +316,15 @@ impl SessionCtxFactory { } pub async fn create_session_ctx(&self, context: &Context) -> QueryResult { - self.create_session_ctx_inner(context, None, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES) - .await + self.create_session_ctx_inner( + context, + None, + legacy_json_source_from_input(&context.input), + None, + None, + DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, + ) + .await } pub async fn create_session_ctx_with_tracker_and_memory_limit( @@ -324,8 +333,15 @@ impl SessionCtxFactory { query_tracker: QueryExecutionTracker, memory_limit_bytes: usize, ) -> QueryResult { - self.create_session_ctx_inner(context, None, Some(query_tracker), None, memory_limit_bytes) - .await + self.create_session_ctx_inner( + context, + None, + legacy_json_source_from_input(&context.input), + Some(query_tracker), + None, + memory_limit_bytes, + ) + .await } pub async fn create_session_ctx_with_snapshot_and_tracker_and_memory_limit( @@ -335,8 +351,33 @@ impl SessionCtxFactory { query_tracker: QueryExecutionTracker, memory_limit_bytes: usize, ) -> QueryResult { - self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), None, memory_limit_bytes) - .await + self.create_session_ctx_inner( + context, + Some(snapshot), + legacy_json_source_from_input(&context.input), + Some(query_tracker), + None, + memory_limit_bytes, + ) + .await + } + + pub async fn create_session_ctx_for_query_with_source_and_tracker_and_memory_limit( + &self, + query: &Query, + source: JsonSource, + query_tracker: QueryExecutionTracker, + memory_limit_bytes: usize, + ) -> QueryResult { + self.create_session_ctx_inner( + query.context(), + query.snapshot().cloned(), + source, + Some(query_tracker), + Some(Arc::clone(query.input_metrics())), + memory_limit_bytes, + ) + .await } pub async fn create_session_ctx_for_query_with_tracker_and_memory_limit( @@ -348,6 +389,7 @@ impl SessionCtxFactory { self.create_session_ctx_inner( query.context(), query.snapshot().cloned(), + legacy_json_source_from_input(&query.context().input), Some(query_tracker), Some(Arc::clone(query.input_metrics())), memory_limit_bytes, @@ -359,12 +401,13 @@ impl SessionCtxFactory { &self, context: &Context, snapshot: Option>, + source: JsonSource, query_tracker: Option, input_metrics: Option>, memory_limit_bytes: usize, ) -> QueryResult { let df_session_ctx = self - .build_df_session_context(context, snapshot, query_tracker.clone(), input_metrics, memory_limit_bytes) + .build_df_session_context(context, snapshot, source, query_tracker.clone(), input_metrics, memory_limit_bytes) .await?; Ok(SessionCtx { @@ -378,6 +421,7 @@ impl SessionCtxFactory { &self, context: &Context, snapshot: Option>, + source: JsonSource, query_tracker: Option, input_metrics: Option>, memory_limit_bytes: usize, @@ -401,10 +445,12 @@ impl SessionCtxFactory { .is_some_and(|delimiter| delimiter.len() == 2 && delimiter.as_bytes() != b"\r\n"); let scan_range_requires_single_file_scan = context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none(); + let json_document_requires_single_file_scan = is_json_document_input(&context.input); let metered_input_requires_single_file_scan = input_metrics.is_some() && context.input.request.input_serialization.parquet.is_none(); let config = if custom_two_byte_record_delimiter || scan_range_requires_single_file_scan + || json_document_requires_single_file_scan || metered_input_requires_single_file_scan { config.with_repartition_file_scans(false) @@ -463,14 +509,21 @@ impl SessionCtxFactory { } else { let input_metrics = input_metrics.unwrap_or_else(|| Arc::new(SelectInputMetrics::default())); let store: EcObjectStore = match query_tracker { - Some(query_tracker) => EcObjectStore::new_with_query_tracker( + Some(query_tracker) => EcObjectStore::new_with_query_tracker_and_source( context.input.clone(), memory_pool, query_tracker, input_metrics, snapshot, + source, + ), + None => EcObjectStore::new_with_memory_pool_and_source( + context.input.clone(), + memory_pool, + input_metrics, + snapshot, + source, ), - None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, input_metrics, snapshot), } .map_err(|err| QueryError::Datafusion { source: Box::new(DataFusionError::External(Box::new(err))), @@ -543,15 +596,15 @@ mod tests { use crate::storage_api::object_store::ObjectIO as _; use datafusion::{ datasource::{ - file_format::csv::CsvFormat, + file_format::{csv::CsvFormat, json::JsonFormat}, listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, }, execution::memory_pool::MemoryLimit, }; use http::HeaderMap; use s3s::dto::{ - CSVInput, CSVOutput, ExpressionType, InputSerialization, JSONInput, OutputSerialization, ParquetInput, ScanRange, - SelectObjectContentInput, SelectObjectContentRequest, + CSVInput, CSVOutput, ExpressionType, InputSerialization, JSONInput, JSONType, OutputSerialization, ParquetInput, + ScanRange, SelectObjectContentInput, SelectObjectContentRequest, }; use std::io::Write as _; @@ -592,6 +645,103 @@ mod tests { ) } + async fn test_query_tracker() -> QueryExecutionTracker { + let permit = Arc::new(tokio::sync::Semaphore::new(1)) + .acquire_owned() + .await + .expect("query permit should be available"); + QueryExecutionTracker::new( + &QueryExecutionOwner::new(), + Arc::new(permit), + Instant::now() + std::time::Duration::from_secs(300), + 300, + ) + } + + async fn assert_legacy_json_column(bucket: &str, expression: &str, document: &[u8], column_name: &str, expected: &[&str]) { + const OBJECT: &str = "input.json"; + + let env = crate::storage_api::select_test_ecstore_env().await; + let mut context = test_context(); + { + let input = Arc::make_mut(&mut context.input); + input.bucket = bucket.to_string(); + input.key = OBJECT.to_string(); + input.request.expression = expression.to_string(); + input.request.input_serialization.csv = None; + input.request.input_serialization.json = Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::DOCUMENT)), + }); + } + env.make_bucket(bucket, false).await; + env.put_object_bytes(bucket, OBJECT, document.to_vec()).await; + let factory = SessionCtxFactory::new(false); + let lazy_session = factory + .create_session_ctx(&context) + .await + .expect("legacy lazy session should preserve the JSON source"); + let tracked_lazy_session = factory + .create_session_ctx_with_tracker_and_memory_limit( + &context, + test_query_tracker().await, + DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, + ) + .await + .expect("legacy tracked lazy session should preserve the JSON source"); + let snapshot = prepare_test_snapshot(&context).await; + let snapshot_session = factory + .create_session_ctx_with_snapshot_and_tracker_and_memory_limit( + &context, + snapshot, + test_query_tracker().await, + DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, + ) + .await + .expect("legacy snapshot session should preserve the JSON source"); + + for (kind, session) in [ + ("lazy", lazy_session), + ("tracked lazy", tracked_lazy_session), + ("tracked snapshot", snapshot_session), + ] { + let table_path = ListingTableUrl::parse(format!("s3://{bucket}/{OBJECT}")).expect("parse JSON table URL"); + let listing_options = ListingOptions::new(Arc::new(JsonFormat::default())).with_file_extension(".json"); + let schema = listing_options + .infer_schema(session.inner(), &table_path) + .await + .expect("infer expanded JSON schema"); + let table = ListingTable::try_new( + ListingTableConfig::new(table_path) + .with_listing_options(listing_options) + .with_schema(schema), + ) + .expect("build expanded JSON table"); + let query_context = SessionContext::new_with_state(session.inner().clone()); + query_context + .register_table("legacy_input", Arc::new(table)) + .expect("register expanded JSON table"); + let batches = query_context + .sql(&format!("SELECT {column_name} FROM legacy_input")) + .await + .expect("plan expanded JSON query") + .collect() + .await + .expect("execute expanded JSON query"); + let mut values = Vec::new(); + for batch in batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("expanded column should be Utf8"); + for row in 0..batch.num_rows() { + values.push(column.value(row).to_string()); + } + } + assert_eq!(values, expected, "{kind} legacy constructor"); + } + } + #[test] fn session_factory_fields_remain_source_compatible() { let factory = SessionCtxFactory { @@ -629,6 +779,7 @@ mod tests { .create_session_ctx_inner( context, None, + JsonSource::default(), None, Some(Arc::new(SelectInputMetrics::default())), DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, @@ -680,6 +831,40 @@ mod tests { assert!(!session.inner().config().options().optimizer.repartition_file_scans); } + #[tokio::test] + async fn json_lines_without_scan_range_keeps_file_repartitioning() { + let mut context = test_context(); + let request = &mut Arc::make_mut(&mut context.input).request; + request.input_serialization.csv = None; + request.input_serialization.json = Some(JSONInput::default()); + + let session = SessionCtxFactory::new(true) + .with_target_partitions(2) + .create_session_ctx(&context) + .await + .expect("JSON LINES session should be created"); + + assert!(session.inner().config().options().optimizer.repartition_file_scans); + } + + #[tokio::test] + async fn json_document_disables_file_repartitioning() { + let mut context = test_context(); + let request = &mut Arc::make_mut(&mut context.input).request; + request.input_serialization.csv = None; + request.input_serialization.json = Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::DOCUMENT)), + }); + + let session = SessionCtxFactory::new(true) + .with_target_partitions(2) + .create_session_ctx(&context) + .await + .expect("JSON DOCUMENT session should be created"); + + assert!(!session.inner().config().options().optimizer.repartition_file_scans); + } + #[tokio::test] async fn csv_scan_range_disables_file_repartitioning() { let mut context = test_context(); @@ -755,7 +940,7 @@ mod tests { async fn session_factory_applies_memory_limit() { let factory = SessionCtxFactory::new(true); let session = factory - .create_session_ctx_inner(&test_context(), None, None, None, 1024) + .create_session_ctx_inner(&test_context(), None, JsonSource::default(), None, None, 1024) .await .expect("session should be created with a bounded memory pool"); @@ -803,6 +988,45 @@ mod tests { assert!(session.is_bound_to(&tracker)); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn legacy_session_factory_preserves_single_key_json_source() { + assert_legacy_json_column( + "s3select-legacy-session-json-source", + "SELECT e.name FROM S3Object.employees AS e", + br#"{"employees":[{"name":"Alice"}]}"#, + "name", + &["Alice"], + ) + .await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn legacy_session_factory_preserves_implicit_root_alias() { + assert_legacy_json_column( + "s3select-legacy-session-root-alias", + "SELECT S3Object FROM S3Object", + br#"["one","two"]"#, + "s3object", + &["one", "two"], + ) + .await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn legacy_session_factory_preserves_quoted_root_alias() { + assert_legacy_json_column( + "s3select-legacy-session-quoted-root-alias", + "SELECT \"V\" FROM S3Object AS \"V\"", + br#"["one","two"]"#, + "\"V\"", + &["one", "two"], + ) + .await; + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial] async fn session_factory_propagates_query_guard_to_ec_store() { diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 77a989da3..72117ce61 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -41,11 +41,11 @@ use rustfs_s3select_api::{ QueryError, QueryResult, SelectError, query::{ Query, - ast::ExtStatement, + ast::{ExtStatement, JsonPathSegment, JsonSource}, dispatcher::{DispatchedQuery, QueryDispatcher}, execution::{Output, QueryStateMachine}, function::FuncMetaManagerRef, - logical_planner::{LogicalPlanner, Plan}, + logical_planner::Plan, parser::Parser, session::{ DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, QueryAdmission, QueryExecutionOwner, QueryExecutionStatus, @@ -53,7 +53,7 @@ use rustfs_s3select_api::{ }, }, }; -use s3s::dto::{FileHeaderInfo, SelectObjectContentInput}; +use s3s::dto::{FileHeaderInfo, JSONType, SelectObjectContentInput}; use std::sync::LazyLock; use tokio::{ sync::Semaphore, @@ -66,6 +66,7 @@ use crate::{ instance::{DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_QUERY_TIMEOUT_SECS}, metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider}, sql::logical::planner::DefaultLogicalPlanner, + sql::planner::prepare_s3_select_statement, }; static IGNORE: LazyLock = LazyLock::new(|| FileHeaderInfo::from_static(FileHeaderInfo::IGNORE)); @@ -160,25 +161,18 @@ impl QueryDispatcher for SimpleQueryDispatcher { let session = &query_state_machine.session; let query = &query_state_machine.query; - let scheme_provider = self.build_scheme_provider(session).await?; - let logical_planner = DefaultLogicalPlanner::new(&scheme_provider); - let statements = self.parser.parse(query.content())?; - - if statements.len() > 1 { - return Err(QueryError::MultiStatement { - num: statements.len(), - sql: query_state_machine.query.content().to_string(), - }); - } - - let stmt = match statements.front() { - Some(stmt) => stmt.clone(), + let stmt = match query_state_machine.prepared_statement() { + Some(statement) => statement.clone(), None => { - return Err(QueryError::Parser { - source: ParserError::ParserError("empty SQL expression".to_string()), - }); + let (statement, source) = self.prepare_query_statement(query.content())?; + if source_path_requires_expansion(source.path()) { + return Err(SelectError::DataSourcePathUnsupported.into()); + } + statement } }; + let scheme_provider = self.build_scheme_provider(session).await?; + let logical_planner = DefaultLogicalPlanner::new(&scheme_provider); let logical_plan = self .statement_to_logical_plan(stmt, &logical_planner, Arc::clone(&query_state_machine)) @@ -269,12 +263,20 @@ impl SimpleQueryDispatcher { self.query_timeout.as_secs(), ); let phase_guard = QueryPhaseGuard::new(&query_tracker, &self.query_execution_owner); + // Keep parser and analyzer errors in the planning phase. Successful + // preparation is cached here because the source path configures the + // object store before schema inference starts. + let (prepared_statement, source) = match self.prepare_query_statement(query.content()) { + Ok((statement, source)) => (Some(statement), source), + Err(_) => (None, JsonSource::default()), + }; let session = self .run_with_query_deadline( &query_tracker, self.session_factory - .create_session_ctx_for_query_with_tracker_and_memory_limit( + .create_session_ctx_for_query_with_source_and_tracker_and_memory_limit( &query, + source, query_tracker.clone(), self.memory_limit_bytes, ), @@ -285,7 +287,28 @@ impl SimpleQueryDispatcher { return Err(self.query_tracker_error(&query_tracker)); } phase_guard.disarm(); - Ok(Arc::new(QueryStateMachine::begin_tracked(query, session, query_tracker)?)) + let state_machine = match prepared_statement { + Some(statement) => QueryStateMachine::begin_tracked_prepared(query, session, query_tracker, statement)?, + None => QueryStateMachine::begin_tracked(query, session, query_tracker)?, + }; + Ok(Arc::new(state_machine)) + } + + fn prepare_query_statement(&self, sql: &str) -> QueryResult<(ExtStatement, JsonSource)> { + let mut statements = self.parser.parse(sql)?; + if statements.len() > 1 { + return Err(QueryError::MultiStatement { + num: statements.len(), + sql: sql.to_string(), + }); + } + let mut statement = statements.pop_front().ok_or_else(|| QueryError::Parser { + source: ParserError::ParserError("empty SQL expression".to_string()), + })?; + let ExtStatement::SqlStatement(sql_statement) = &mut statement; + let source = prepare_s3_select_statement(sql_statement)?; + validate_json_source_path_input(&self.input, source.path())?; + Ok((statement, source)) } async fn run_with_query_deadline( &self, @@ -365,7 +388,7 @@ impl SimpleQueryDispatcher { // begin analyze query_state_machine.begin_analyze(); let logical_plan = logical_planner - .create_logical_plan(stmt, &query_state_machine.session) + .prepared_statement_to_plan(stmt, &query_state_machine.session) .await?; query_state_machine.end_analyze(); @@ -503,6 +526,31 @@ impl SimpleQueryDispatcher { } } +fn validate_json_source_path_input(input: &SelectObjectContentInput, source_path: &[JsonPathSegment]) -> QueryResult<()> { + if source_path.is_empty() { + return Ok(()); + } + let Some(json) = input.request.input_serialization.json.as_ref() else { + return Err(SelectError::DataSourcePathUnsupported.into()); + }; + if !source_path_requires_expansion(source_path) + || json + .type_ + .as_ref() + .is_some_and(|json_type| json_type.as_str() == JSONType::DOCUMENT) + { + return Ok(()); + } + Err(SelectError::DataSourcePathUnsupported.into()) +} + +fn source_path_requires_expansion(source_path: &[JsonPathSegment]) -> bool { + !source_path + .strip_prefix(&[JsonPathSegment::ArrayWildcard]) + .unwrap_or(source_path) + .is_empty() +} + pub struct TrackedRecordBatchStream { state: Arc, schema: SchemaRef, @@ -760,7 +808,10 @@ impl SimpleQueryDispatcherBuilder { #[cfg(test)] mod tests { - use super::{QueryPhaseGuard, SimpleQueryDispatcher, SimpleQueryDispatcherBuilder, TrackedRecordBatchStream}; + use super::{ + QueryPhaseGuard, SimpleQueryDispatcher, SimpleQueryDispatcherBuilder, TrackedRecordBatchStream, + validate_json_source_path_input, + }; use crate::{ execution::{ factory::{QueryExecutionFactoryRef, SqlQueryExecutionFactory}, @@ -789,6 +840,7 @@ mod tests { QueryError, QueryResult, SelectError, query::{ Context as QueryContext, Query, + ast::JsonPathSegment, dispatcher::QueryDispatcher, execution::{ Output, QueryExecution, QueryExecutionFactory, QueryExecutionRef, QueryStateMachine, QueryStateMachineRef, @@ -1031,7 +1083,17 @@ mod tests { query_execution_factory: QueryExecutionFactoryRef, ) -> (Arc, Arc) { let input = Arc::new(test_input()); - let dispatcher = SimpleQueryDispatcherBuilder::default() + let dispatcher = test_dispatcher_for_input(Arc::clone(&input), admission, query_timeout, query_execution_factory); + (dispatcher, input) + } + + fn test_dispatcher_for_input( + input: Arc, + admission: Arc, + query_timeout: Duration, + query_execution_factory: QueryExecutionFactoryRef, + ) -> Arc { + SimpleQueryDispatcherBuilder::default() .with_input(Arc::clone(&input)) .with_default_table_provider(Arc::new(BaseTableProvider::default())) .with_session_factory(Arc::new(SessionCtxFactory::new(true))) @@ -1041,8 +1103,7 @@ mod tests { .with_query_admission(admission) .with_query_timeout(query_timeout) .build() - .expect("query dispatcher should build"); - (dispatcher, input) + .expect("query dispatcher should build") } async fn snapshot_test_env() -> &'static TestECStoreEnv { @@ -1171,6 +1232,68 @@ mod tests { }) } + #[test] + fn nested_source_paths_require_json_document_input() { + let lines_input = json_snapshot_input(); + let nested_path = [JsonPathSegment::Key { + name: "employees".to_string(), + quoted: false, + }]; + + assert!(matches!( + validate_json_source_path_input(&lines_input, &nested_path), + Err(ref error) if matches!(error.s3_select_policy_error(), Some(SelectError::DataSourcePathUnsupported)) + )); + assert!(validate_json_source_path_input(&lines_input, &[JsonPathSegment::ArrayWildcard]).is_ok()); + + let csv_input = test_input(); + let parquet_input = parquet_snapshot_input(); + for input in [&csv_input, parquet_input.as_ref()] { + assert!(matches!( + validate_json_source_path_input(input, &nested_path), + Err(ref error) if matches!(error.s3_select_policy_error(), Some(SelectError::DataSourcePathUnsupported)) + )); + } + + let mut document_input = (*lines_input).clone(); + document_input.request.input_serialization.json.as_mut().unwrap().type_ = Some(JSONType::from_static(JSONType::DOCUMENT)); + assert!(validate_json_source_path_input(&document_input, &nested_path).is_ok()); + } + + #[tokio::test] + async fn normal_planning_pipeline_rejects_json_lines_source_expansion() { + let mut input = (*json_snapshot_input()).clone(); + input.request.expression = "SELECT * FROM S3Object.employees".to_string(); + let input = Arc::new(input); + let admission = Arc::new(Semaphore::new(1)); + let optimizer = Arc::new(CascadeOptimizerBuilder::default().build()); + let scheduler = Arc::new(LocalScheduler {}); + let dispatcher = test_dispatcher_for_input( + Arc::clone(&input), + Arc::clone(&admission), + Duration::from_secs(300), + Arc::new(SqlQueryExecutionFactory::new(optimizer, scheduler)), + ); + let query = Query::new( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + ); + let query_state_machine = dispatcher + .build_query_state_machine(query) + .await + .expect("validation errors should remain in the planning phase"); + + let result = dispatcher.build_logical_plan(query_state_machine).await; + + assert!(matches!( + result, + Err(ref error) if matches!(error.s3_select_policy_error(), Some(SelectError::DataSourcePathUnsupported)) + )); + assert_eq!(admission.available_permits(), 1); + } + fn parquet_snapshot_input() -> Arc { Arc::new(SelectObjectContentInput { bucket: "s3select-parquet-snapshot-race".to_string(), @@ -1530,6 +1653,140 @@ mod tests { assert!(matches!(result, Err(QueryError::Cancel))); } + #[tokio::test] + async fn unprepared_tracked_session_rejects_source_path_expansion() { + let mut input = test_input(); + input.key = "test.json".to_string(); + input.request.expression = "SELECT * FROM S3Object.employees".to_string(); + input.request.input_serialization = InputSerialization { + json: Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::DOCUMENT)), + }), + ..Default::default() + }; + input.request.output_serialization = OutputSerialization { + json: Some(JSONOutput::default()), + ..Default::default() + }; + let input = Arc::new(input); + let admission = Arc::new(Semaphore::new(1)); + let optimizer = Arc::new(CascadeOptimizerBuilder::default().build()); + let scheduler = Arc::new(LocalScheduler {}); + let dispatcher = test_dispatcher_for_input( + Arc::clone(&input), + Arc::clone(&admission), + Duration::from_secs(300), + Arc::new(SqlQueryExecutionFactory::new(optimizer, scheduler)), + ); + let query = Query::new( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + ); + let permit = Arc::clone(&admission).acquire_owned().await.expect("admission permit"); + let tracker = QueryExecutionTracker::new( + &dispatcher.query_execution_owner, + Arc::new(permit), + Instant::now() + Duration::from_secs(300), + 300, + ); + let session = SessionCtxFactory::new(true) + .create_session_ctx_with_tracker_and_memory_limit( + query.context(), + tracker.clone(), + DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, + ) + .await + .expect("test session"); + assert!(tracker.mark_admitted(&dispatcher.query_execution_owner)); + let state_machine = + Arc::new(QueryStateMachine::begin_tracked(query, session, tracker).expect("tracked state machine should be valid")); + + let result = dispatcher.build_logical_plan(state_machine).await; + + assert!(matches!( + result, + Err(ref error) if matches!(error.s3_select_policy_error(), Some(SelectError::DataSourcePathUnsupported)) + )); + assert_eq!(admission.available_permits(), 1); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn unprepared_tracked_session_preserves_root_scalar_bindings() { + for (bucket, expression) in [ + ("s3select-unprepared-root-scalar-alias", "SELECT V FROM S3Object AS V"), + ("s3select-unprepared-root-wildcard", "SELECT _1 FROM S3Object[*]"), + ] { + let mut input = test_input(); + input.bucket = bucket.to_string(); + input.key = "input.json".to_string(); + input.request.expression = expression.to_string(); + input.request.input_serialization = InputSerialization { + json: Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::DOCUMENT)), + }), + ..Default::default() + }; + input.request.output_serialization = OutputSerialization { + json: Some(JSONOutput::default()), + ..Default::default() + }; + let input = Arc::new(input); + let env = snapshot_test_env().await; + env.make_bucket(&input.bucket, false).await; + env.put_object_bytes(&input.bucket, &input.key, br#"["one","two"]"#.to_vec()) + .await; + let snapshot = env.prepare_select_object_snapshot(&input.bucket, &input.key).await; + let dispatcher = production_dispatcher(Arc::clone(&input)); + let query = Query::new_with_snapshot( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + Arc::clone(&snapshot), + ); + let permit = Arc::clone(&dispatcher.query_admission) + .acquire_owned() + .await + .expect("query permit should be available"); + let tracker = QueryExecutionTracker::new( + &dispatcher.query_execution_owner, + Arc::new(permit), + Instant::now() + Duration::from_secs(300), + 300, + ); + let session = SessionCtxFactory::new(false) + .create_session_ctx_with_snapshot_and_tracker_and_memory_limit( + query.context(), + snapshot, + tracker.clone(), + DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, + ) + .await + .expect("legacy session should preserve the root scalar binding"); + assert!(tracker.mark_admitted(&dispatcher.query_execution_owner)); + let state_machine = Arc::new( + QueryStateMachine::begin_tracked(query, session, tracker).expect("tracked state machine should be valid"), + ); + + let logical_plan = dispatcher + .build_logical_plan(Arc::clone(&state_machine)) + .await + .expect("unprepared query should plan") + .expect("SELECT should produce a logical plan"); + let values = collect_utf8_output( + dispatcher + .execute_logical_plan(logical_plan, state_machine) + .await + .expect("unprepared query should execute"), + ) + .await; + + assert_eq!(values, ["one", "two"]); + } + } + #[tokio::test] async fn staged_query_rejects_unbound_session() { let admission = Arc::new(Semaphore::new(1)); @@ -1955,6 +2212,38 @@ mod tests { )); } + #[tokio::test] + async fn invalid_sql_precedes_malformed_json_snapshot_read() { + let mut input = json_snapshot_input(); + let input_mut = Arc::make_mut(&mut input); + input_mut.bucket = "s3select-invalid-sql-precedence".to_string(); + input_mut.key = "malformed.json".to_string(); + input_mut.request.expression = "SELECT * FROM".to_string(); + input_mut.request.input_serialization.json.as_mut().unwrap().type_ = Some(JSONType::from_static(JSONType::DOCUMENT)); + + let env = snapshot_test_env().await; + env.make_bucket(&input.bucket, false).await; + env.put_object_bytes(&input.bucket, &input.key, b"{bad".to_vec()).await; + let snapshot = env.prepare_select_object_snapshot(&input.bucket, &input.key).await; + let dispatcher = production_dispatcher(Arc::clone(&input)); + let query = Query::new_with_snapshot( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + snapshot, + ); + let query_state_machine = dispatcher + .build_query_state_machine(query) + .await + .expect("invalid SQL should remain a planning-phase error"); + + assert!(matches!( + dispatcher.build_logical_plan(query_state_machine).await, + Err(QueryError::Parser { .. }) + )); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn cancelled_execution_start_drops_future_before_releasing_admission() { let admission = Arc::new(Semaphore::new(1)); diff --git a/crates/s3select-query/src/sql/dialect.rs b/crates/s3select-query/src/sql/dialect.rs index 686d24621..6a59fa764 100644 --- a/crates/s3select-query/src/sql/dialect.rs +++ b/crates/s3select-query/src/sql/dialect.rs @@ -184,6 +184,13 @@ mod tests { assert!(dialect.supports_group_by_expr(), "RustFsDialect should support GROUP BY expressions"); } + #[test] + fn test_supports_partiql_paths() { + let dialect = RustFsDialect; + + assert!(dialect.supports_partiql(), "RustFsDialect should support JSON source paths"); + } + #[test] fn test_identifier_validation_comprehensive() { let dialect = RustFsDialect; diff --git a/crates/s3select-query/src/sql/parser.rs b/crates/s3select-query/src/sql/parser.rs index 83eb40f03..302646324 100644 --- a/crates/s3select-query/src/sql/parser.rs +++ b/crates/s3select-query/src/sql/parser.rs @@ -16,6 +16,7 @@ use std::{collections::VecDeque, fmt::Display}; use datafusion::sql::sqlparser::{ dialect::Dialect, + keywords::{Keyword, RESERVED_FOR_TABLE_ALIAS}, parser::{Parser, ParserError}, tokenizer::{Token, Tokenizer}, }; @@ -53,7 +54,8 @@ impl<'a> ExtParser<'a> { /// Parse the specified tokens with dialect fn new_with_dialect(sql: &str, dialect: &'a dyn Dialect) -> Result { let mut tokenizer = Tokenizer::new(dialect, sql); - let tokens = tokenizer.tokenize()?; + let mut tokens = tokenizer.tokenize()?; + rewrite_source_object_wildcards(&mut tokens); Ok(ExtParser { parser: Parser::new(dialect).with_tokens(tokens), }) @@ -104,6 +106,41 @@ impl<'a> ExtParser<'a> { } } +fn rewrite_source_object_wildcards(tokens: &mut [Token]) { + let mut paren_depth = 0usize; + let mut in_from = false; + let mut index = 0usize; + + while index < tokens.len() { + match &tokens[index] { + Token::Word(word) if paren_depth == 0 && word.keyword == Keyword::FROM => { + in_from = true; + } + Token::Word(word) if in_from && paren_depth == 0 && ends_from_source(word.keyword) => { + in_from = false; + } + Token::SemiColon if paren_depth == 0 => in_from = false, + Token::LParen => paren_depth = paren_depth.saturating_add(1), + Token::RParen => paren_depth = paren_depth.saturating_sub(1), + Token::Period if in_from => { + if let Some(next) = tokens[index + 1..] + .iter_mut() + .find(|token| !matches!(token, Token::Whitespace(_))) + && matches!(next, Token::Mul) + { + *next = Token::make_word("*", None); + } + } + _ => {} + } + index += 1; + } +} + +fn ends_from_source(keyword: Keyword) -> bool { + RESERVED_FOR_TABLE_ALIAS.contains(&keyword) || matches!(keyword, Keyword::PREWHERE | Keyword::SETTINGS | Keyword::FORMAT) +} + #[cfg(test)] mod tests { use super::*; @@ -172,6 +209,23 @@ mod tests { } } + #[test] + fn parses_source_object_wildcard_without_rewriting_projection_wildcard() { + let mut statements = ExtParser::parse_sql("SELECT e.* FROM S3Object[*].* AS e").expect("query should parse"); + let ExtStatement::SqlStatement(statement) = statements.pop_front().expect("one statement"); + + assert_eq!(statement.to_string(), "SELECT e.* FROM S3Object[*].* AS e"); + } + + #[test] + fn from_tokens_in_literals_and_comments_do_not_change_wildcard_scope() { + let sql = "SELECT 'FROM x.*' AS marker /* FROM y.* */ FROM S3Object.*"; + let mut statements = ExtParser::parse_sql(sql).expect("query should parse"); + let ExtStatement::SqlStatement(statement) = statements.pop_front().expect("one statement"); + + assert_eq!(statement.to_string(), "SELECT 'FROM x.*' AS marker FROM S3Object.*"); + } + #[test] fn test_default_parser_multiple_statements() { let parser = DefaultParser::default(); diff --git a/crates/s3select-query/src/sql/planner.rs b/crates/s3select-query/src/sql/planner.rs index a5bbb2881..b847c07d9 100644 --- a/crates/s3select-query/src/sql/planner.rs +++ b/crates/s3select-query/src/sql/planner.rs @@ -12,20 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::ops::ControlFlow; +use std::{convert::Infallible, ops::ControlFlow}; use async_recursion::async_recursion; use async_trait::async_trait; use datafusion::sql::{ - planner::SqlToRel, + planner::{IdentNormalizer, SqlToRel}, sqlparser::ast::{ - GroupByExpr, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor, SetExpr, Statement, TableFactor, Visit, Visitor, + AccessExpr, Expr, GroupByExpr, Ident, JsonPath, JsonPathElem, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor, + SetExpr, Statement, Subscript, TableAlias, TableFactor, Value, Visit, VisitMut, Visitor, VisitorMut, }, }; use rustfs_s3select_api::{ QueryError, QueryResult, SelectError, query::{ - ast::ExtStatement, + ast::{ExtStatement, JsonPathSegment, JsonSource}, logical_planner::{LogicalPlanner, Plan, QueryPlan}, session::SessionCtx, }, @@ -64,21 +65,24 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> { } } - async fn df_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult { - match stmt { - Statement::Query(_) => { - validate_s3_select_statement(&stmt)?; - let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?; - let plan = Plan::Query(QueryPlan { - df_plan, - is_tag_scan: false, - }); - - Ok(plan) - } - _ => Err(unsupported_structure("only SELECT queries are supported")), + pub(crate) async fn prepared_statement_to_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult { + match statement { + ExtStatement::SqlStatement(stmt) => self.df_prepared_sql_to_plan(*stmt, session).await, } } + + async fn df_sql_to_plan(&self, mut stmt: Statement, session: &SessionCtx) -> QueryResult { + prepare_s3_select_statement(&mut stmt)?; + self.df_prepared_sql_to_plan(stmt, session).await + } + + async fn df_prepared_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult { + let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?; + Ok(Plan::Query(QueryPlan { + df_plan, + is_tag_scan: false, + })) + } } fn classify_planner_error(error: datafusion::common::DataFusionError) -> QueryError { @@ -95,11 +99,10 @@ fn classify_planner_error(error: datafusion::common::DataFusionError) -> QueryEr error.into() } -fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> { +pub(crate) fn prepare_s3_select_statement(statement: &mut 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() @@ -137,17 +140,81 @@ fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> { } let mut detector = SubqueryDetector { visited_root: false }; - if query.visit(&mut detector).is_break() { + if Visit::visit(&*query, &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")); + let source = { + let SetExpr::Select(select) = query.body.as_mut() else { + return Err(unsupported_structure("set operations and nested queries are not supported")); + }; + prepare_select(select)? }; - validate_select(select) + let mut normalizer = PartiQlSubscriptNormalizer; + let _ = VisitMut::visit(query, &mut normalizer); + Ok(source) } -fn validate_select(select: &Select) -> QueryResult<()> { +struct PartiQlSubscriptNormalizer; + +impl VisitorMut for PartiQlSubscriptNormalizer { + type Break = Infallible; + + fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { + let Expr::JsonAccess { value, path } = expr else { + return ControlFlow::Continue(()); + }; + if !matches!(path.path.first(), Some(JsonPathElem::Bracket { .. })) + || path + .path + .iter() + .any(|element| matches!(element, JsonPathElem::ColonBracket { .. })) + { + return ControlFlow::Continue(()); + } + + let mut appended_access = Vec::with_capacity(path.path.len()); + for element in std::mem::take(&mut path.path) { + match element { + JsonPathElem::Dot { key, quoted } => { + let identifier = if quoted { + Ident::with_quote('"', key) + } else { + Ident::new(key) + }; + appended_access.push(AccessExpr::Dot(Expr::Identifier(identifier))); + } + JsonPathElem::Bracket { key } => { + appended_access.push(AccessExpr::Subscript(Subscript::Index { index: key })); + } + JsonPathElem::ColonBracket { key } => { + appended_access.push(AccessExpr::Subscript(Subscript::Index { index: key })); + } + } + } + + let value = std::mem::replace(value, Box::new(Expr::Identifier(Ident::new("")))); + let (root, mut access_chain) = match *value { + Expr::CompoundFieldAccess { root, access_chain } => (root, access_chain), + root => (Box::new(root), Vec::new()), + }; + access_chain.extend(appended_access); + *expr = Expr::CompoundFieldAccess { root, access_chain }; + ControlFlow::Continue(()) + } +} + +fn implicit_source_alias(source_path: &[JsonPathSegment]) -> Ident { + match source_path.last() { + Some(JsonPathSegment::Key { name, quoted: true }) => Ident::with_quote('"', name), + Some(JsonPathSegment::Key { name, quoted: false }) => Ident::new(name), + Some(JsonPathSegment::Index(_) | JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard) | None => { + Ident::new("_1") + } + } +} + +fn prepare_select(select: &mut Select) -> QueryResult { if !select.optimizer_hints.is_empty() || select.distinct.is_some() || select.select_modifiers.is_some() @@ -170,7 +237,7 @@ fn validate_select(select: &Select) -> QueryResult<()> { return Err(unsupported_structure("the SELECT contains an unsupported clause")); } - let [table] = select.from.as_slice() else { + let [table] = select.from.as_mut_slice() else { return Err(unsupported_structure("exactly one S3Object source is required")); }; if !table.joins.is_empty() { @@ -186,8 +253,8 @@ fn validate_select(select: &Select) -> QueryResult<()> { partitions, sample, index_hints, - .. - } = &table.relation + json_path, + } = &mut table.relation else { return Err(unsupported_structure("subqueries and table functions are not supported")); }; @@ -202,9 +269,7 @@ fn validate_select(select: &Select) -> QueryResult<()> { { 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 { + let Some(ObjectNamePart::Identifier(table_name)) = name.0.first() else { return Err(SelectError::DataSourcePathUnsupported.into()); }; let is_s3_object = if table_name.quote_style.is_some() { @@ -216,6 +281,72 @@ fn validate_select(select: &Select) -> QueryResult<()> { return Err(SelectError::DataSourcePathUnsupported.into()); } + let mut source_path = Vec::new(); + for part in &name.0[1..] { + let ObjectNamePart::Identifier(identifier) = part else { + return Err(SelectError::DataSourcePathUnsupported.into()); + }; + if identifier.quote_style.is_none() && identifier.value == "*" { + source_path.push(JsonPathSegment::ObjectWildcard); + } else { + source_path.push(JsonPathSegment::Key { + name: identifier.value.clone(), + quoted: identifier.quote_style.is_some(), + }); + } + } + if let Some(json_path) = json_path.as_ref() { + append_json_path_segments(&mut source_path, json_path)?; + } + if alias.is_none() && !source_path.is_empty() { + *alias = Some(TableAlias { + explicit: true, + name: implicit_source_alias(&source_path), + columns: Vec::new(), + at: None, + }); + } + let scalar_column = alias + .as_ref() + .map(|alias| IdentNormalizer::default().normalize(alias.name.clone())) + .or_else(|| { + source_path + .is_empty() + .then(|| IdentNormalizer::default().normalize(table_name.clone())) + }); + name.0.truncate(1); + *json_path = None; + Ok(JsonSource::new(source_path, scalar_column)) +} + +fn append_json_path_segments(source_path: &mut Vec, json_path: &JsonPath) -> QueryResult<()> { + for element in &json_path.path { + let segment = match element { + JsonPathElem::Dot { key, quoted } if key == "*" && !quoted => JsonPathSegment::ObjectWildcard, + JsonPathElem::Dot { key, quoted } => JsonPathSegment::Key { + name: key.clone(), + quoted: *quoted, + }, + JsonPathElem::Bracket { key: Expr::Wildcard(_) } => JsonPathSegment::ArrayWildcard, + JsonPathElem::Bracket { key: Expr::Value(value) } => match &value.value { + Value::Number(number, false) => JsonPathSegment::Index( + number + .to_string() + .parse() + .map_err(|_| QueryError::from(SelectError::DataSourcePathUnsupported))?, + ), + Value::SingleQuotedString(key) => JsonPathSegment::Key { + name: key.clone(), + quoted: true, + }, + _ => return Err(SelectError::DataSourcePathUnsupported.into()), + }, + JsonPathElem::Bracket { .. } | JsonPathElem::ColonBracket { .. } => { + return Err(SelectError::DataSourcePathUnsupported.into()); + } + }; + source_path.push(segment); + } Ok(()) } @@ -245,10 +376,14 @@ impl Visitor for SubqueryDetector { #[cfg(test)] mod tests { - use super::validate_s3_select_statement; + use super::prepare_s3_select_statement; use crate::sql::parser::ExtParser; - use datafusion::sql::sqlparser::ast::Statement; - use rustfs_s3select_api::{SelectError, query::ast::ExtStatement}; + use datafusion::sql::sqlparser::ast::{AccessExpr, Expr, Statement, Visit, Visitor}; + use rustfs_s3select_api::{ + QueryResult, SelectError, + query::ast::{ExtStatement, JsonPathSegment, JsonSource}, + }; + use std::ops::ControlFlow; fn parse_statement(sql: &str) -> Statement { let mut statements = ExtParser::parse_sql(sql).expect("SQL should parse"); @@ -256,6 +391,10 @@ mod tests { *statement } + fn validate_s3_select_statement(statement: &Statement) -> QueryResult { + prepare_s3_select_statement(&mut statement.clone()) + } + #[test] fn accepts_s3_select_query_shape() { let statement = parse_statement("SELECT s.id FROM S3Object AS s WHERE s.id = '1' LIMIT 10"); @@ -270,6 +409,195 @@ mod tests { assert!(validate_s3_select_statement(&statement).is_ok()); } + #[test] + fn prepares_nested_json_source_path_and_normalizes_table() { + let mut statement = parse_statement("SELECT e.name FROM S3Object[*].employees[*] AS e"); + + let source = prepare_s3_select_statement(&mut statement).expect("JSON source path should be supported"); + + assert_eq!( + source.path(), + &[ + JsonPathSegment::ArrayWildcard, + JsonPathSegment::Key { + name: "employees".to_string(), + quoted: false, + }, + JsonPathSegment::ArrayWildcard, + ] + ); + assert_eq!(statement.to_string(), "SELECT e.name FROM S3Object AS e"); + } + + #[test] + fn partiql_source_support_preserves_projection_and_filter_subscripts() { + let mut statement = parse_statement("SELECT s.tags[1] FROM S3Object AS s WHERE s.values[0] = 1"); + + prepare_s3_select_statement(&mut statement).expect("array expressions should remain supported"); + let mut counter = FieldAccessCounter::default(); + let _ = Visit::visit(&statement, &mut counter); + + assert_eq!(counter.json_accesses, 0); + assert_eq!(counter.subscripts, 2); + } + + #[derive(Default)] + struct FieldAccessCounter { + json_accesses: usize, + subscripts: usize, + } + + impl Visitor for FieldAccessCounter { + type Break = (); + + fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow { + match expr { + Expr::JsonAccess { .. } => self.json_accesses += 1, + Expr::CompoundFieldAccess { access_chain, .. } => { + self.subscripts += access_chain + .iter() + .filter(|access| matches!(access, AccessExpr::Subscript(_))) + .count(); + } + _ => {} + } + ControlFlow::Continue(()) + } + } + + #[test] + fn prepares_array_index_and_object_wildcard_paths() { + let mut index_statement = parse_statement("SELECT * FROM S3Object[0]"); + let mut wildcard_statement = parse_statement("SELECT * FROM S3Object[*].*"); + + assert_eq!( + prepare_s3_select_statement(&mut index_statement) + .expect("array index should be supported") + .path(), + &[JsonPathSegment::Index(0)] + ); + assert_eq!( + prepare_s3_select_statement(&mut wildcard_statement) + .expect("object wildcard should be supported") + .path(), + &[JsonPathSegment::ArrayWildcard, JsonPathSegment::ObjectWildcard] + ); + } + + #[test] + fn quoted_star_remains_an_object_key() { + let mut statement = parse_statement("SELECT * FROM S3Object.\"*\""); + + let source = prepare_s3_select_statement(&mut statement).expect("quoted key should be supported"); + + assert_eq!( + source.path(), + &[JsonPathSegment::Key { + name: "*".to_string(), + quoted: true, + }] + ); + } + + #[test] + fn preserves_quoted_keys_and_adds_implicit_source_aliases() { + let mut key_statement = parse_statement("SELECT employee.name FROM S3Object[*].department.employee"); + let mut wildcard_statement = parse_statement("SELECT _1.name FROM S3Object[*].employees[*]"); + + prepare_s3_select_statement(&mut key_statement).expect("named source path should be supported"); + prepare_s3_select_statement(&mut wildcard_statement).expect("wildcard source path should be supported"); + + assert_eq!(key_statement.to_string(), "SELECT employee.name FROM S3Object AS employee"); + assert_eq!(wildcard_statement.to_string(), "SELECT _1.name FROM S3Object AS _1"); + } + + #[test] + fn root_scalar_aliases_are_preserved_and_unquoted_aliases_are_normalized() { + let mut implicit = parse_statement("SELECT S3Object FROM S3Object"); + let mut unquoted = parse_statement("SELECT V FROM S3Object AS V"); + let mut quoted = parse_statement("SELECT \"V\" FROM S3Object AS \"V\""); + + let implicit_source = prepare_s3_select_statement(&mut implicit).expect("implicit root alias should be supported"); + let unquoted_source = prepare_s3_select_statement(&mut unquoted).expect("unquoted root alias should be supported"); + let quoted_source = prepare_s3_select_statement(&mut quoted).expect("quoted root alias should be supported"); + + assert!(implicit_source.path().is_empty()); + assert_eq!(implicit_source.scalar_column(), Some("s3object")); + assert!(unquoted_source.path().is_empty()); + assert_eq!(unquoted_source.scalar_column(), Some("v")); + assert!(quoted_source.path().is_empty()); + assert_eq!(quoted_source.scalar_column(), Some("V")); + } + + #[test] + fn unquoted_terminal_scalar_alias_uses_datafusion_identifier_case() { + let mut statement = parse_statement("SELECT NAME FROM S3Object[*].NAME"); + + let source = prepare_s3_select_statement(&mut statement).expect("terminal scalar source should be supported"); + + assert_eq!(source.scalar_column(), Some("name")); + assert_eq!(statement.to_string(), "SELECT NAME FROM S3Object AS NAME"); + } + + #[test] + fn single_quoted_source_key_adds_a_quoted_implicit_alias() { + let mut statement = parse_statement("SELECT \"Employee Data\".id FROM S3Object['Employee Data']"); + + let source = prepare_s3_select_statement(&mut statement).expect("single-quoted source key should be supported"); + + assert_eq!( + source.path(), + &[JsonPathSegment::Key { + name: "Employee Data".to_string(), + quoted: true, + }] + ); + assert_eq!(source.scalar_column(), Some("Employee Data")); + assert_eq!(statement.to_string(), "SELECT \"Employee Data\".id FROM S3Object AS \"Employee Data\""); + } + + #[test] + fn accepts_object_wildcard_continuation() { + let mut statement = parse_statement("SELECT * FROM S3Object[*].groups.*.id"); + + assert_eq!( + prepare_s3_select_statement(&mut statement) + .expect("object wildcard continuation should be supported") + .path(), + &[ + JsonPathSegment::ArrayWildcard, + JsonPathSegment::Key { + name: "groups".to_string(), + quoted: false, + }, + JsonPathSegment::ObjectWildcard, + JsonPathSegment::Key { + name: "id".to_string(), + quoted: false, + }, + ] + ); + } + + #[test] + fn rejects_non_literal_or_out_of_range_array_indexes() { + for sql in [ + "SELECT * FROM S3Object[-1]", + "SELECT * FROM S3Object[1 + 1]", + "SELECT * FROM S3Object[999999999999999999999999999999999999]", + ] { + let statement = parse_statement(sql); + assert!( + matches!( + validate_s3_select_statement(&statement), + Err(ref error) + if matches!(error.s3_select_policy_error(), Some(SelectError::DataSourcePathUnsupported)) + ), + "query should reject an unsafe array index: {sql}" + ); + } + } + #[test] fn accepts_group_by_and_order_by() { let statement = parse_statement("SELECT department, COUNT(*) FROM S3Object GROUP BY department ORDER BY department");