// Copyright 2024 RustFS Team // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. use crate::{ SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectInfo, SelectObjectOptions, SelectStorageError, SelectStore, query::{ parser::RustFsDialect, session::{QueryExecutionGuard, QueryExecutionTracker}, }, resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found, select_is_err_version_not_found, }; use async_trait::async_trait; use bytes::Bytes; use chrono::Utc; use datafusion::{ common::{DataFusionError, runtime::SpawnedTask}, execution::memory_pool::{MemoryConsumer, MemoryPool, UnboundedMemoryPool}, object_store::{ Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path, }, sql::sqlparser::{ ast::{ObjectNamePart, SetExpr, Statement, TableFactor}, parser::Parser as SqlParser, }, }; use futures::pin_mut; use futures::{Stream, StreamExt, future::ready, stream}; use futures_core::stream::BoxStream; use http::{HeaderMap, HeaderValue, header::HeaderName}; use parking_lot::Mutex; use rustfs_common::DEFAULT_DELIMITER; use s3s::S3Result; use s3s::dto::SelectObjectContentInput; use s3s::header::{ X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5, }; use s3s::s3_error; use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; use tokio::io::AsyncReadExt; use tokio_util::io::ReaderStream; use transform_stream::AsyncTryStream; use crate::storage_api::object_store::{HTTPRangeSpec, ObjectIO as _, ObjectOperations as _}; fn select_default_read_buffer_size_u64() -> u64 { u64::try_from(SELECT_DEFAULT_READ_BUFFER_SIZE).unwrap_or(u64::MAX) } fn validated_object_size(size: i64) -> Result { u64::try_from(size).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), }) } /// Maximum allowed object size for JSON DOCUMENT mode. /// /// JSON DOCUMENT format requires loading the entire file into memory for DOM /// parsing, so memory consumption grows linearly with file size. Objects /// larger than this threshold are rejected with an error rather than risking /// an OOM condition. /// /// To process larger JSON files, convert the input to **JSON LINES** (NDJSON, /// `type = LINES`), which supports line-by-line streaming with no memory /// size limit. /// /// 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. pub const MAX_JSON_DOCUMENT_BYTES: u64 = 128 * 1024 * 1024; const JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER: usize = 64; pub const INVALID_SCAN_RANGE_MESSAGE: &str = "The value of a parameter in ScanRange element is invalid. Check the service API documentation and try again."; const NORMALIZED_RECORD_DELIMITER: &[u8] = b"\r\n"; const NORMALIZED_FIELD_DELIMITER: &[u8] = &[DEFAULT_DELIMITER]; #[derive(Debug)] pub struct EcObjectStore { input: Arc, need_convert: bool, delimiter: String, /// True when the JSON input type is DOCUMENT (multi-line formatted JSON). /// 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, memory_pool: Arc, query_tracker: Option, store: Arc, } #[cfg(test)] struct ScanRangeBeforeMainHook { bucket: String, object: String, reached: tokio::sync::oneshot::Sender<()>, resume: tokio::sync::oneshot::Receiver<()>, } #[cfg(test)] static SCAN_RANGE_BEFORE_MAIN_HOOK: tokio::sync::Mutex> = tokio::sync::Mutex::const_new(None); #[cfg(test)] async fn run_scan_range_before_main_hook(bucket: &str, object: &str) { let hook = { let mut hook = SCAN_RANGE_BEFORE_MAIN_HOOK.lock().await; if hook .as_ref() .is_some_and(|hook| hook.bucket == bucket && hook.object == object) { hook.take() } else { None } }; if let Some(hook) = hook { let _ = hook.reached.send(()); let _ = hook.resume.await; } } #[derive(Clone, Copy, Debug)] pub struct SelectScanRange { start: u64, end: u64, } impl SelectScanRange { pub const fn new(start: u64, end: u64) -> Self { Self { start, end } } pub const fn start(&self) -> u64 { self.start } pub const fn end(&self) -> u64 { self.end } } #[derive(Clone, Copy, Debug)] pub struct InvalidScanRange; impl EcObjectStore { pub fn new(input: Arc) -> S3Result { Self::build(input, Arc::new(UnboundedMemoryPool::default()), None, None) } pub(crate) fn new_with_memory_pool(input: Arc, memory_pool: Arc) -> S3Result { Self::build(input, memory_pool, None, None) } pub(crate) fn new_with_query_tracker( input: Arc, memory_pool: Arc, query_tracker: QueryExecutionTracker, store: Option>, ) -> S3Result { Self::build(input, memory_pool, Some(query_tracker), store) } fn build( input: Arc, memory_pool: Arc, query_tracker: Option, store: Option>, ) -> S3Result { let Some(store) = store.or_else(resolve_select_object_store_handle) else { return Err(s3_error!(InternalError, "ec store not inited")); }; let (need_convert, delimiter) = if let Some(csv) = input.request.input_serialization.csv.as_ref() { if let Some(delimiter) = csv.field_delimiter.as_ref() { if delimiter.len() > 1 { (true, delimiter.to_owned()) } else { (false, String::new()) } } else { (false, String::new()) } } else { (false, String::new()) }; // 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 }; Ok(Self { input, need_convert, delimiter, is_json_document, json_sub_path, memory_pool, query_tracker, store, }) } fn object_options(&self, options: &GetOptions) -> SelectObjectOptions { SelectObjectOptions { version_id: options.version.clone(), ..Default::default() } } fn read_headers(&self) -> HeaderMap { select_read_headers(&self.input) } fn scan_range(&self, object_size: u64) -> Result> { let Some(scan_range) = self.input.request.scan_range.as_ref() else { return Ok(None); }; scan_range_from_bounds(scan_range.start, scan_range.end, object_size) } fn record_delimiter(&self) -> Vec { self.input .request .input_serialization .csv .as_ref() .and_then(|csv| csv.record_delimiter.as_ref()) .map(|delimiter| delimiter.as_bytes().to_vec()) .unwrap_or_else(|| b"\n".to_vec()) } fn record_delimiter_for_conversion(&self) -> Option> { let delimiter = self.record_delimiter(); (self.need_convert || (delimiter.len() == 2 && delimiter != NORMALIZED_RECORD_DELIMITER)).then_some(delimiter) } fn csv_has_header(&self) -> bool { self.input .request .input_serialization .csv .as_ref() .and_then(|csv| csv.file_header_info.as_ref()) .is_some_and(|info| matches!(info.as_str(), "USE" | "IGNORE")) } async fn object_info(&self, opts: &SelectObjectOptions) -> Result { self.store .get_object_info(&self.input.bucket, &self.input.key, opts) .await .map_err(|err| map_storage_error(&self.input.bucket, &self.input.key, err)) } async fn object_reader(&self, range: Option, opts: &SelectObjectOptions) -> Result { let h = self.read_headers(); self.store .get_object_reader(&self.input.bucket, &self.input.key, range, h, opts) .await .map_err(|err| map_storage_error(&self.input.bucket, &self.input.key, err)) } async fn read_raw_range_with_opts( &self, range: Range, opts: &SelectObjectOptions, expected_snapshot: Option<&SelectObjectInfo>, ) -> Result { if range.is_empty() { return Ok(Bytes::new()); } let reader = self .object_reader(Some(http_range_spec_from_range(range.clone())), opts) .await?; if let Some(expected_snapshot) = expected_snapshot { validate_object_snapshot(expected_snapshot, &reader.object_info)?; } let object_size = validated_object_size(reader.object_info.size)?; let resolved_range = GetRange::Bounded(range) .as_range(object_size) .map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; let expected_size = usize::try_from(resolved_range.end - resolved_range.start).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; let mut reader = reader.stream; let mut bytes = Vec::new(); reader.read_to_end(&mut bytes).await.map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; if bytes.len() < expected_size { return Err(incomplete_object_stream_error(expected_size - bytes.len())); } bytes.truncate(expected_size); Ok(Bytes::from(bytes)) } async fn read_raw_range(&self, range: Range) -> Result { self.read_raw_range_with_opts(range, &self.object_options(&GetOptions::new()), None) .await } async fn read_header_record( &self, object_size: u64, delimiter: &[u8], opts: &SelectObjectOptions, expected_snapshot: &SelectObjectInfo, ) -> Result { if object_size == 0 { return Ok(Bytes::new()); } let mut end = select_default_read_buffer_size_u64().min(object_size); loop { let bytes = self.read_raw_range_with_opts(0..end, opts, Some(expected_snapshot)).await?; if let Some(pos) = find_delimiter(&bytes, delimiter) { return Ok(bytes.slice(0..pos + delimiter.len())); } if end == object_size { return Ok(bytes); } end = end.saturating_mul(2).min(object_size); } } async fn scan_range_read_start( &self, scan_range: SelectScanRange, delimiter: &[u8], opts: &SelectObjectOptions, expected_snapshot: &SelectObjectInfo, ) -> Result { let delimiter_len = u64::try_from(delimiter.len()).unwrap_or(u64::MAX); let fallback_start = scan_range.start().saturating_sub(delimiter_len); if delimiter.len() != 2 || delimiter[0] != delimiter[1] || scan_range.start() == 0 { return Ok(fallback_start); } let context_start = scan_range.start().saturating_sub(select_default_read_buffer_size_u64()); let context = self .read_raw_range_with_opts(context_start..scan_range.start(), opts, Some(expected_snapshot)) .await?; let suffix_len = context.iter().rev().take_while(|byte| **byte == delimiter[0]).count(); if suffix_len == context.len() && context_start > 0 { return Err(o_Error::Generic { store: "EcObjectStore", source: "self-overlapping CSV record delimiter exceeds the bounded ScanRange context".into(), }); } if suffix_len == 0 { return Ok(fallback_start); } Ok(scan_range.start().saturating_sub(u64::from(suffix_len % 2 != 0))) } } impl std::fmt::Display for EcObjectStore { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str("EcObjectStore") } } fn unsupported_store_error(op: &str) -> o_Error { o_Error::Generic { store: "s3select-api", source: Box::new(std::io::Error::new( std::io::ErrorKind::Unsupported, format!("operation {op} is not supported in EcObjectStore"), )), } } fn insert_header(headers: &mut HeaderMap, name: HeaderName, value: Option<&str>) { if let Some(value) = value && let Ok(value) = HeaderValue::from_str(value) { headers.insert(name, value); } } fn select_read_headers(input: &SelectObjectContentInput) -> HeaderMap { let mut headers = HeaderMap::new(); insert_header( &mut headers, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM, input.sse_customer_algorithm.as_deref(), ); insert_header(&mut headers, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY, input.sse_customer_key.as_deref()); insert_header( &mut headers, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5, input.sse_customer_key_md5.as_deref(), ); headers } fn http_range_spec_from_get_range(range: &GetRange) -> HTTPRangeSpec { match range { GetRange::Bounded(range) => http_range_spec_from_range(range.clone()), GetRange::Offset(start) => HTTPRangeSpec { is_suffix_length: false, start: *start as i64, end: -1, }, GetRange::Suffix(length) => HTTPRangeSpec { is_suffix_length: true, start: *length as i64, end: -1, }, } } fn http_range_spec_from_range(range: Range) -> HTTPRangeSpec { HTTPRangeSpec { is_suffix_length: false, start: range.start as i64, end: range.end.saturating_sub(1) as i64, } } fn http_range_spec_from_start(start: u64) -> HTTPRangeSpec { HTTPRangeSpec { is_suffix_length: false, start: start as i64, end: -1, } } fn validate_object_snapshot(expected: &SelectObjectInfo, actual: &SelectObjectInfo) -> Result<()> { if expected.size != actual.size || expected.version_id != actual.version_id || expected.data_dir != actual.data_dir || expected.etag != actual.etag || expected.mod_time != actual.mod_time { return Err(o_Error::Generic { store: "EcObjectStore", source: "object changed while preparing SelectObjectContent ScanRange".into(), }); } Ok(()) } fn find_delimiter(bytes: &[u8], delimiter: &[u8]) -> Option { if delimiter.is_empty() { return None; } bytes.windows(delimiter.len()).position(|window| window == delimiter) } fn map_storage_error(bucket: &str, object: &str, err: SelectStorageError) -> o_Error { if select_is_err_bucket_not_found(&err) || select_is_err_object_not_found(&err) || select_is_err_version_not_found(&err) { return o_Error::NotFound { path: format!("{bucket}/{object}"), source: err.to_string().into(), }; } o_Error::Generic { store: "EcObjectStore", source: Box::new(err), } } pub fn scan_range_from_bounds(start: Option, end: Option, object_size: u64) -> Result> { parse_scan_range_from_bounds(start, end, object_size).map_err(|_| invalid_scan_range_store_error()) } pub fn validate_scan_range_bounds( start: Option, end: Option, object_size: u64, ) -> std::result::Result<(), InvalidScanRange> { parse_scan_range_from_bounds(start, end, object_size).map(|_| ()) } fn parse_scan_range_from_bounds( start: Option, end: Option, object_size: u64, ) -> std::result::Result, InvalidScanRange> { if start.is_none() && end.is_none() { return Ok(None); } if start.is_some_and(|value| value < 0) || end.is_some_and(|value| value < 0) { return Err(InvalidScanRange); } if let (Some(start), Some(end)) = (start, end) && start > end { return Err(InvalidScanRange); } if let Some(start) = start { let start = start as u64; if object_size == 0 { if start > 0 { return Err(InvalidScanRange); } return Ok(Some(SelectScanRange::new(0, 0))); } if start >= object_size { return Err(InvalidScanRange); } } if object_size == 0 { return Ok(Some(SelectScanRange::new(0, 0))); } let last_byte = object_size - 1; let (start, end) = match (start, end) { (Some(start), Some(end)) => (start as u64, (end as u64).min(last_byte)), (Some(start), None) => (start as u64, last_byte), (None, Some(suffix_len)) => { let suffix_len = suffix_len as u64; (object_size.saturating_sub(suffix_len), last_byte) } (None, None) => return Ok(None), }; Ok(Some(SelectScanRange::new(start, end))) } fn invalid_scan_range_store_error() -> o_Error { o_Error::Generic { store: "EcObjectStore", source: format!("ScanRange: {INVALID_SCAN_RANGE_MESSAGE}").into(), } } #[async_trait] impl ObjectStore for EcObjectStore { async fn put_opts(&self, _location: &Path, _payload: PutPayload, _opts: PutOptions) -> Result { Err(unsupported_store_error("put_opts")) } async fn put_multipart_opts(&self, _location: &Path, _opts: PutMultipartOptions) -> Result> { Err(unsupported_store_error("put_multipart_opts")) } async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { let opts = self.object_options(&options); let record_delimiter = if options.head { None } else { self.record_delimiter_for_conversion() }; let needs_scan_context = options.range.is_none() && !options.head && self.input.request.scan_range.is_some(); let scan_context = if needs_scan_context { let source_snapshot = self.object_info(&opts).await?; let original_size = validated_object_size(source_snapshot.size)?; if let Some(scan_range) = self.scan_range(original_size)? { let delimiter = self.record_delimiter(); let read_start = self .scan_range_read_start(scan_range, &delimiter, &opts, &source_snapshot) .await?; Some((source_snapshot, scan_range, read_start)) } else { None } } else { None }; #[cfg(test)] if scan_context.is_some() { run_scan_range_before_main_hook(&self.input.bucket, &self.input.key).await; } let range = options.range.as_ref().map(http_range_spec_from_get_range); let reader = if let Some((source_snapshot, _, read_start)) = scan_context.as_ref() { let range = (source_snapshot.size > 0).then(|| http_range_spec_from_start(*read_start)); self.object_reader(range, &opts).await? } else { self.object_reader(range, &opts).await? }; if let Some((source_snapshot, _, _)) = scan_context.as_ref() { validate_object_snapshot(source_snapshot, &reader.object_info)?; } let original_size = match scan_context.as_ref() { Some((source_snapshot, _, _)) => validated_object_size(source_snapshot.size)?, None => validated_object_size(reader.object_info.size)?, }; let etag = reader.object_info.etag; let version = reader.object_info.version_id.map(|version| version.to_string()); let attributes = Attributes::default(); let result_range = match options.range.as_ref() { Some(range) => range.as_range(original_size).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?, None => 0..original_size, }; let payload = if options.head { GetResultPayload::Stream(stream::empty().boxed()) } else if options.range.is_some() { let size = usize::try_from(result_range.end - result_range.start).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; GetResultPayload::Stream( bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), size).boxed(), ) } else if self.is_json_document { // JSON DOCUMENT mode: gate on object size before doing any I/O. // // Small files (<= MAX_JSON_DOCUMENT_BYTES): build a lazy stream // that defers all I/O and JSON parsing until DataFusion first // polls it. Parsing runs inside spawn_blocking so the async // runtime thread is never blocked. // // Large files (> MAX_JSON_DOCUMENT_BYTES): return an error // immediately. JSON DOCUMENT relies on serde_json DOM parsing // which must load the whole file into memory; rejecting oversized // files upfront is safer than risking OOM. Users should convert // their data to JSON LINES (NDJSON) format for large files. validate_json_document_size(original_size)?; let stream = json_document_ndjson_stream( reader.stream, original_size, self.json_sub_path.clone(), Arc::clone(&self.memory_pool), self.query_tracker.clone(), ); GetResultPayload::Stream(stream) } else if let Some((source_snapshot, scan_range, read_start)) = scan_context { let delimiter = self.record_delimiter(); let include_header = self.csv_has_header(); let header = if include_header && read_start > 0 { Some( self.read_header_record(original_size, &delimiter, &opts, &source_snapshot) .await?, ) } else { None }; let stream = scan_range_stream( ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), delimiter, scan_range, include_header && header.is_none(), read_start, original_size, ) .boxed(); let stream = if let Some(header) = header { stream::once(ready(Ok(header))).chain(stream).boxed() } else { stream }; GetResultPayload::Stream(convert_csv_delimiter_stream( stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone()), )) } else { let stream = bytes_stream( ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), original_size as usize, ); GetResultPayload::Stream(convert_csv_delimiter_stream( stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone()), )) }; let meta = ObjectMeta { location: location.clone(), last_modified: Utc::now(), size: original_size, e_tag: etag, version, }; Ok(GetResult { payload, meta, range: result_range, attributes, }) } async fn get_ranges(&self, _location: &Path, ranges: &[Range]) -> Result> { let mut out = Vec::with_capacity(ranges.len()); for range in ranges { out.push(self.read_raw_range(range.clone()).await?); } Ok(out) } fn delete_stream(&self, _locations: BoxStream<'static, Result>) -> BoxStream<'static, Result> { stream::once(ready(Err(unsupported_store_error("delete_stream")))).boxed() } fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, Result> { stream::once(ready(Err(unsupported_store_error("list")))).boxed() } async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> Result { Err(unsupported_store_error("list_with_delimiter")) } async fn copy_opts(&self, _from: &Path, _to: &Path, _options: CopyOptions) -> Result<()> { Err(unsupported_store_error("copy_opts")) } } struct CsvDelimiterConverter { record_delimiter: Option>, field_delimiter: Option>, carry: Vec, } impl CsvDelimiterConverter { fn new(record_delimiter: Option>, field_delimiter: Option>) -> Self { Self { record_delimiter: record_delimiter.filter(|delimiter| !delimiter.is_empty()), field_delimiter: field_delimiter.filter(|delimiter| !delimiter.is_empty()), carry: Vec::new(), } } fn max_delimiter_len(&self) -> usize { self.record_delimiter .as_ref() .into_iter() .chain(self.field_delimiter.as_ref()) .map(Vec::len) .max() .unwrap_or(1) } fn convert_prefix(&self, bytes: &[u8], end: usize) -> (Vec, usize) { let mut converted = Vec::with_capacity(bytes.len()); let mut pos = 0; while pos < end { let record_match = self .record_delimiter .as_ref() .filter(|delimiter| bytes[pos..].starts_with(delimiter)); let field_match = self .field_delimiter .as_ref() .filter(|delimiter| bytes[pos..].starts_with(delimiter)); if let Some(delimiter) = field_match && record_match.is_none_or(|record_delimiter| delimiter.len() > record_delimiter.len()) { converted.extend_from_slice(NORMALIZED_FIELD_DELIMITER); pos += delimiter.len(); } else if let Some(delimiter) = record_match { if delimiter.len() == 2 && delimiter != NORMALIZED_RECORD_DELIMITER { converted.extend_from_slice(NORMALIZED_RECORD_DELIMITER); } else { converted.extend_from_slice(delimiter); } pos += delimiter.len(); } else { converted.push(bytes[pos]); pos += 1; } } (converted, pos) } fn convert_chunk(&mut self, chunk: &[u8]) -> Vec { let mut combined = Vec::with_capacity(self.carry.len() + chunk.len()); combined.extend_from_slice(&self.carry); combined.extend_from_slice(chunk); let safe_end = combined.len().saturating_sub(self.max_delimiter_len().saturating_sub(1)); let (converted, pos) = self.convert_prefix(&combined, safe_end); self.carry.clear(); self.carry.extend_from_slice(&combined[pos..]); converted } fn finish(&mut self) -> Vec { let (converted, _) = self.convert_prefix(&self.carry, self.carry.len()); self.carry.clear(); converted } } fn convert_delimiter_stream( stream: S, record_delimiter: Option>, field_delimiter: Option>, ) -> BoxStream<'static, Result> where S: Stream> + Send + 'static, { AsyncTryStream::::new(|mut y| async move { let mut converter = CsvDelimiterConverter::new(record_delimiter, field_delimiter); pin_mut!(stream); while let Some(result) = stream.next().await { let bytes = result?; let converted = converter.convert_chunk(&bytes); if !converted.is_empty() { y.yield_ok(Bytes::from(converted)).await; } } let converted = converter.finish(); if !converted.is_empty() { y.yield_ok(Bytes::from(converted)).await; } Ok(()) }) .boxed() } #[cfg(test)] fn convert_record_delimiter_stream(stream: S, delimiter: Vec) -> BoxStream<'static, Result> where S: Stream> + Send + 'static, { // DataFusion's CSV reader treats CRLF as a record terminator. convert_delimiter_stream(stream, Some(delimiter), None) } #[cfg(test)] fn convert_field_delimiter_stream(stream: S, delimiter: String) -> BoxStream<'static, Result> where S: Stream> + Send + 'static, { convert_delimiter_stream(stream, None, Some(delimiter.into_bytes())) } fn convert_csv_delimiter_stream( stream: S, record_delimiter: Option>, field_delimiter: Option, ) -> BoxStream<'static, Result> where S: Stream> + Send + 'static, { match (record_delimiter, field_delimiter) { (None, None) => stream.boxed(), (record, field) => convert_delimiter_stream(stream, record, field.map(String::into_bytes)), } } struct ScanRangeState { stream: S, delimiter: Vec, range: SelectScanRange, include_header: bool, offset: u64, record_start: u64, record: Vec, pending: VecDeque, done: bool, expected_end: u64, } fn scan_range_stream( stream: S, delimiter: Vec, range: SelectScanRange, include_header: bool, base_offset: u64, expected_end: u64, ) -> BoxStream<'static, Result> where S: Stream> + Send + Unpin + 'static, { let state = ScanRangeState { stream, delimiter, range, include_header, offset: base_offset, record_start: base_offset, record: Vec::new(), pending: VecDeque::new(), done: false, expected_end, }; stream::unfold(state, |mut state| async move { loop { if let Some(bytes) = state.pending.pop_front() { return Some((Ok(bytes), state)); } if state.done { return None; } match state.stream.next().await { Some(Ok(bytes)) => state.push_chunk(&bytes), Some(Err(err)) => { state.done = true; return Some(( Err(o_Error::Generic { store: "EcObjectStore", source: Box::new(err), }), state, )); } None => { if state.offset < state.expected_end { state.done = true; return Some((Err(incomplete_object_stream_error(state.expected_end - state.offset)), state)); } state.finish_pending_record(); state.done = true; } } } }) .boxed() } impl ScanRangeState { fn push_chunk(&mut self, bytes: &[u8]) { if bytes.is_empty() || self.done { return; } if self.record.is_empty() { self.record_start = self.offset; } let search_start = self.record.len().saturating_sub(self.delimiter.len().saturating_sub(1)); self.record.extend_from_slice(bytes); self.offset = self.offset.saturating_add(bytes.len() as u64); let mut search_start = search_start; while let Some(pos) = find_delimiter(&self.record[search_start..], &self.delimiter) { let record_end = search_start + pos + self.delimiter.len(); self.finish_record(record_end); if self.done { break; } search_start = 0; } } fn finish_record(&mut self, record_end: usize) { let record = self.record.drain(..record_end).collect::>(); let record_start = self.record_start; self.record_start = self.record_start.saturating_add(record_end as u64); self.push_record(record, record_start); } fn finish_pending_record(&mut self) { if self.record.is_empty() { return; } let record = std::mem::take(&mut self.record); let record_start = self.record_start; self.push_record(record, record_start); } fn push_record(&mut self, record: Vec, record_start: u64) { let include_header = self.include_header && record_start == 0; let include_record = record_start >= self.range.start() && record_start <= self.range.end(); if include_header || include_record { self.pending.push_back(Bytes::from(record)); } else { if record_start > self.range.end() { self.done = true; } } } } 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 /// DataFusion begins polling the returned stream. The pipeline is: /// /// 1. **Read** – the object bytes are read asynchronously from `stream` only /// when the returned stream is first polled. /// 2. **Parse** – JSON deserialization runs inside /// `tokio::task::spawn_blocking` so the async runtime is never blocked by /// CPU-bound work, even for very large documents. /// 3. **Yield** – each NDJSON line (one per array element, or one line for a /// scalar/object root) is yielded as a separate [`Bytes`] chunk, so /// DataFusion can pipeline row processing as lines arrive. fn json_document_ndjson_stream( stream: Box, original_size: u64, json_sub_path: Option, memory_pool: Arc, query_tracker: Option, ) -> futures_core::stream::BoxStream<'static, Result> { json_document_ndjson_stream_with_parser( stream, original_size, json_sub_path, memory_pool, query_tracker, |all_bytes, json_sub_path| parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()), ) } fn json_document_ndjson_stream_with_parser

( stream: Box, original_size: u64, json_sub_path: Option, memory_pool: Arc, query_tracker: Option, parser: P, ) -> futures_core::stream::BoxStream<'static, Result> where P: FnOnce(Vec, Option) -> std::io::Result> + Send + 'static, { AsyncTryStream::::new(|mut y| async move { // Compact JSON can expand substantially into a serde_json DOM and // per-record output buffers, so reserve a conservative upper bound // before the source buffer is allocated. let buffer_capacity = usize::try_from(original_size).map_err(|_| o_Error::Generic { store: "EcObjectStore", source: Box::new(DataFusionError::ResourcesExhausted(format!( "JSON DOCUMENT input size {original_size} does not fit in memory" ))), })?; let reservation_bytes = buffer_capacity .checked_mul(JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER) .ok_or_else(|| o_Error::Generic { store: "EcObjectStore", source: Box::new(DataFusionError::ResourcesExhausted(format!( "JSON DOCUMENT memory reservation overflow for {original_size} input bytes" ))), })?; let reservation = MemoryConsumer::new("S3 Select JSON document").register(&memory_pool); reservation.try_resize(reservation_bytes).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; pin_mut!(stream); // ── 1. Read phase (lazy: only runs when the stream is polled) ──── let mut all_bytes = Vec::with_capacity(buffer_capacity); stream .take(original_size) .read_to_end(&mut all_bytes) .await .map_err(|e| o_Error::Generic { store: "EcObjectStore", source: Box::new(e), })?; if all_bytes.len() != buffer_capacity { return Err(incomplete_object_stream_error(buffer_capacity - all_bytes.len())); } // ── 2. Parse phase (blocking thread pool, non-blocking runtime) ── let pending_query_guard = PendingQueryExecutionGuard::new(query_tracker); let task_query_guard = pending_query_guard.task_state(); let (lines, _reservation, _query_guard) = SpawnedTask::spawn_blocking(move || { let query_guard = PendingQueryExecutionGuard::start(&task_query_guard)?; parser(all_bytes, json_sub_path).map(|lines| (lines, reservation, query_guard)) }) .await .map_err(|e| o_Error::Generic { store: "EcObjectStore", source: e.to_string().into(), })? .map_err(|e| o_Error::Generic { store: "EcObjectStore", source: Box::new(e), })?; // ── 3. Yield phase (one Bytes per NDJSON line) ─────────────────── for line in lines { y.yield_ok(line).await; } Ok(()) }) .boxed() } struct PendingQueryExecutionGuard { state: Arc>, } enum QueryExecutionGuardState { Pending(Option), Started, Cancelled, } impl PendingQueryExecutionGuard { fn new(query_tracker: Option) -> Self { Self { state: Arc::new(Mutex::new(QueryExecutionGuardState::Pending(query_tracker))), } } fn task_state(&self) -> Arc> { Arc::clone(&self.state) } fn start(state: &Mutex) -> std::io::Result> { let mut state = state.lock(); match std::mem::replace(&mut *state, QueryExecutionGuardState::Started) { QueryExecutionGuardState::Pending(None) => Ok(None), QueryExecutionGuardState::Pending(Some(query_tracker)) => query_tracker.query_guard().map(Some).ok_or_else(|| { std::io::Error::new(std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started") }), QueryExecutionGuardState::Cancelled => { *state = QueryExecutionGuardState::Cancelled; Err(std::io::Error::new( std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started", )) } QueryExecutionGuardState::Started => { *state = QueryExecutionGuardState::Started; Err(std::io::Error::other("JSON DOCUMENT parse started more than once")) } } } } impl Drop for PendingQueryExecutionGuard { fn drop(&mut self) { let query_guard = { let mut state = self.state.lock(); match std::mem::replace(&mut *state, QueryExecutionGuardState::Cancelled) { QueryExecutionGuardState::Pending(query_guard) => query_guard, QueryExecutionGuardState::Started => { *state = QueryExecutionGuardState::Started; None } QueryExecutionGuardState::Cancelled => None, } }; drop(query_guard); } } /// Parse a JSON DOCUMENT (a single JSON value, possibly multi-line) into a /// list of NDJSON lines – one [`Bytes`] per record. /// /// `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. /// /// - 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> { 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 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)); } } 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)); } } Ok(lines) } /// Convert a JSON DOCUMENT to a single concatenated NDJSON [`Bytes`] blob. /// /// This is a convenience wrapper around [`parse_json_document_to_lines`] used /// by the unit tests. Production code uses `json_document_ndjson_stream` /// 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)?; let total = lines.iter().map(|b| b.len()).sum(); let mut output = Vec::with_capacity(total); for line in lines { output.extend_from_slice(&line); } Ok(Bytes::from(output)) } pub fn bytes_stream(stream: S, content_length: usize) -> impl Stream> + Send + 'static where S: Stream> + Send + 'static, { AsyncTryStream::::new(|mut y| async move { pin_mut!(stream); let mut remaining: usize = content_length; while remaining > 0 { let Some(result) = stream.next().await else { break; }; let mut bytes = result.map_err(|e| o_Error::Generic { store: "", source: Box::new(e), })?; if bytes.len() > remaining { bytes.truncate(remaining); } remaining -= bytes.len(); y.yield_ok(bytes).await; } if remaining > 0 { return Err(incomplete_object_stream_error(remaining)); } Ok(()) }) } fn validate_json_document_size(original_size: u64) -> Result<()> { if original_size <= MAX_JSON_DOCUMENT_BYTES { return Ok(()); } Err(o_Error::Generic { store: "EcObjectStore", source: Box::new(DataFusionError::ResourcesExhausted(format!( "JSON DOCUMENT object is {original_size} bytes, which exceeds the maximum allowed size of \ {MAX_JSON_DOCUMENT_BYTES} bytes ({} MiB). Convert the input to JSON LINES (NDJSON) to process large files.", MAX_JSON_DOCUMENT_BYTES / (1024 * 1024) ))), }) } fn incomplete_object_stream_error(remaining: impl std::fmt::Display) -> o_Error { o_Error::Generic { store: "EcObjectStore", source: Box::new(std::io::Error::new( std::io::ErrorKind::UnexpectedEof, format!("object stream ended with {remaining} bytes remaining"), )), } } #[cfg(test)] mod test { use super::{ EcObjectStore, JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER, SCAN_RANGE_BEFORE_MAIN_HOOK, SELECT_DEFAULT_READ_BUFFER_SIZE, ScanRangeBeforeMainHook, SelectScanRange, bytes_stream, convert_csv_delimiter_stream, convert_field_delimiter_stream, convert_record_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, json_document_ndjson_stream, json_document_ndjson_stream_with_parser, scan_range_from_bounds, scan_range_stream, select_read_headers, validate_json_document_size, validate_object_snapshot, validated_object_size, }; use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker}; use crate::storage_api::SelectPutObjReader; use crate::storage_api::object_store::ObjectIO as _; use bytes::Bytes; use datafusion::{ common::DataFusionError, execution::memory_pool::{GreedyMemoryPool, MemoryPool}, execution::{config::SessionConfig, context::SessionContext}, object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path}, physical_plan::ExecutionPlanProperties, prelude::CsvReadOptions, }; use futures::{StreamExt, TryStreamExt, stream}; use s3s::dto::{ CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, OutputSerialization, ScanRange, SelectObjectContentInput, SelectObjectContentRequest, }; use s3s::header::{ X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5, }; use std::sync::{ Arc, atomic::{AtomicUsize, Ordering}, }; #[test] fn ec_object_store_constructor_remains_source_compatible() { let _constructor: fn(Arc) -> s3s::S3Result = EcObjectStore::new; } use tokio::sync::Semaphore; #[test] fn test_validated_object_size_rejects_negative_metadata() { assert_eq!(validated_object_size(0).expect("zero object size should be valid"), 0); assert!(validated_object_size(-1).is_err()); } #[test] fn test_scan_range_snapshot_validation_rejects_changed_object() { let expected = crate::SelectObjectInfo::default(); let mut actual = expected.clone(); assert!(validate_object_snapshot(&expected, &actual).is_ok()); actual.size = 1; assert!(validate_object_snapshot(&expected, &actual).is_err()); actual = expected.clone(); actual.version_id = Some("00000000-0000-0000-0000-000000000001".parse().expect("valid version UUID")); assert!(validate_object_snapshot(&expected, &actual).is_err()); actual = expected.clone(); actual.data_dir = Some("00000000-0000-0000-0000-000000000002".parse().expect("valid data-dir UUID")); assert!(validate_object_snapshot(&expected, &actual).is_err()); actual = expected.clone(); actual.etag = Some("changed".to_string()); assert!(validate_object_snapshot(&expected, &actual).is_err()); actual = expected.clone(); actual.mod_time = Some(std::time::SystemTime::UNIX_EPOCH.into()); assert!(validate_object_snapshot(&expected, &actual).is_err()); } #[tokio::test] async fn test_scan_range_stream_keeps_header_and_selected_record() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]); let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0, 18); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"h1,h2\n2,b\n"); } #[tokio::test] async fn test_scan_range_stream_skips_record_when_start_is_in_middle() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]); let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0, 12); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"2,b\n"); } #[tokio::test] async fn test_scan_range_stream_keeps_record_when_end_is_in_middle() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]); let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0, 12); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"1,a\n2,b\n"); } #[tokio::test] async fn test_scan_range_stream_uses_base_offset_for_range_reader() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"\n2,b\n3,c\n"))]); let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3, 12); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"2,b\n"); } #[tokio::test] async fn test_scan_range_stream_handles_delimiter_split_across_chunks() { let chunks = stream::iter(vec![ Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\r")), Ok::<_, std::io::Error>(Bytes::from_static(b"\n1,a\r\n2,b\r")), Ok::<_, std::io::Error>(Bytes::from_static(b"\n3,c\r\n")), ]); let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0, 22); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"h1,h2\r\n2,b\r\n"); } #[tokio::test] async fn test_scan_range_stream_converts_custom_delimiter_split_across_chunks() { let chunks = stream::iter(vec![ Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2^")), Ok::<_, std::io::Error>(Bytes::from_static(b"Y1,a^Y2,b^")), Ok::<_, std::io::Error>(Bytes::from_static(b"Y3,c^Y")), ]); let stream = scan_range_stream(chunks, b"^Y".to_vec(), SelectScanRange::new(12, 14), true, 0, 22); let mut stream = convert_record_delimiter_stream(stream, b"^Y".to_vec()); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.expect("custom-delimiter ScanRange chunk should be valid")); } assert_eq!(output, b"h1,h2\r\n2,b\r\n"); } #[tokio::test] async fn test_record_delimiter_conversion_carries_partial_delimiter() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a,1^")), Ok::<_, object_store::Error>(Bytes::from_static(b"Yb,2^Y")), ]); let output = convert_record_delimiter_stream(chunks, b"^Y".to_vec()) .try_collect::>() .await .expect("convert record delimiter") .concat(); assert_eq!(output, b"a,1\r\nb,2\r\n"); } #[tokio::test] async fn test_record_delimiter_conversion_preserves_overlapping_match_order() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a")), Ok::<_, object_store::Error>(Bytes::from_static(b"aa")), Ok::<_, object_store::Error>(Bytes::from_static(b"a")), ]); let output = convert_record_delimiter_stream(chunks, b"aa".to_vec()) .try_collect::>() .await .expect("convert overlapping record delimiter") .concat(); assert_eq!(output, b"\r\n\r\n"); } #[tokio::test] async fn test_record_and_field_delimiter_conversion_order() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a\r")), Ok::<_, object_store::Error>(Bytes::from_static(b"\n1^")), Ok::<_, object_store::Error>(Bytes::from_static(b"Yb\r\n2^Y")), ]); let output = convert_csv_delimiter_stream(chunks, Some(b"^Y".to_vec()), Some("\r\n".to_string())) .try_collect::>() .await .expect("convert record and field delimiters") .concat(); assert_eq!(output, b"a,1\r\nb,2\r\n"); } #[tokio::test] async fn test_record_delimiter_takes_precedence_when_delimiters_match() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a^")), Ok::<_, object_store::Error>(Bytes::from_static(b"Yb^Y")), ]); let output = convert_csv_delimiter_stream(chunks, Some(b"^Y".to_vec()), Some("^Y".to_string())) .try_collect::>() .await .expect("convert matching record and field delimiters") .concat(); assert_eq!(output, b"a\r\nb\r\n"); let chunks = stream::iter(vec![Ok::<_, object_store::Error>(Bytes::from_static(b"a\r\nb\r\n"))]); let output = convert_csv_delimiter_stream(chunks, Some(b"\r\n".to_vec()), Some("\r\n".to_string())) .try_collect::>() .await .expect("preserve matching normalized record delimiter") .concat(); assert_eq!(output, b"a\r\nb\r\n"); } #[tokio::test] async fn test_longer_field_delimiter_takes_precedence_over_record_prefix() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a^")), Ok::<_, object_store::Error>(Bytes::from_static(b"YQb^Y")), ]); let output = convert_csv_delimiter_stream(chunks, Some(b"^Y".to_vec()), Some("^YQ".to_string())) .try_collect::>() .await .expect("convert record delimiter that prefixes field delimiter") .concat(); assert_eq!(output, b"a,b\r\n"); let chunks = stream::iter(vec![Ok::<_, object_store::Error>(Bytes::from_static(b"a\nXb\nc\nXd\n"))]); let output = convert_csv_delimiter_stream(chunks, Some(b"\n".to_vec()), Some("\nX".to_string())) .try_collect::>() .await .expect("preserve longer field delimiter with native record prefix") .concat(); assert_eq!(output, b"a,b\nc,d\n"); } #[tokio::test] async fn test_scan_range_stream_rejects_early_eof() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n"))]); let mut output = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 7), false, 0, 8); assert_eq!(output.next().await.expect("first stream item").expect("first record"), b"1,a\n"[..]); let err = output .next() .await .expect("early EOF error") .expect_err("short ScanRange stream must fail"); let object_store::Error::Generic { source, .. } = err else { panic!("expected generic object store error"); }; let source = source.downcast_ref::().expect("I/O error source"); assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof); assert!(source.to_string().contains("4 bytes remaining")); assert!(output.next().await.is_none()); } #[test] fn test_find_delimiter_handles_multi_byte_delimiter() { assert_eq!(find_delimiter(b"one\r\ntwo", b"\r\n"), Some(3)); assert_eq!(find_delimiter(b"one\ntwo", b"\r\n"), None); } #[test] fn test_scan_range_end_only_uses_aws_suffix_semantics() { let range = scan_range_from_bounds(None, Some(35), 100).unwrap().unwrap(); assert_eq!(range.start(), 65); assert_eq!(range.end(), 99); } #[test] fn test_scan_range_start_after_object_is_rejected_before_reader() { let err = scan_range_from_bounds(Some(100), None, 100).unwrap_err(); assert!(err.to_string().contains("ScanRange")); } #[test] fn test_scan_range_start_after_end_is_rejected() { let err = scan_range_from_bounds(Some(20), Some(10), 100).unwrap_err(); assert!(err.to_string().contains("ScanRange")); } #[test] fn test_get_range_conversion_for_parquet_bounded_ranges() { let range = http_range_spec_from_get_range(&GetRange::Bounded(10..20)); assert!(!range.is_suffix_length); assert_eq!(range.start, 10); assert_eq!(range.end, 19); } #[test] fn test_select_read_headers_preserves_ssec_context() { let input = SelectObjectContentInput { bucket: "bucket".to_string(), expected_bucket_owner: None, key: "object.csv".to_string(), sse_customer_algorithm: Some("AES256".to_string()), sse_customer_key: Some("customer-key".to_string()), sse_customer_key_md5: Some("customer-key-md5".to_string()), request: SelectObjectContentRequest { expression: "SELECT * FROM s3object".to_string(), expression_type: ExpressionType::from_static(ExpressionType::SQL), input_serialization: InputSerialization { csv: Some(CSVInput::default()), ..Default::default() }, output_serialization: OutputSerialization { csv: Some(CSVOutput::default()), ..Default::default() }, request_progress: None, scan_range: None, }, }; let headers = select_read_headers(&input); assert_eq!(headers.get(X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM).unwrap(), "AES256"); assert_eq!(headers.get(X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY).unwrap(), "customer-key"); assert_eq!(headers.get(X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5).unwrap(), "customer-key-md5"); } #[tokio::test] async fn test_scan_range_output_can_convert_field_delimiter() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"a&&1\nb&&2\n"))]); let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0, 10); let mut stream = convert_field_delimiter_stream(stream, "&&".to_string()); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"a,1\nb,2\n"); } #[tokio::test] async fn test_field_delimiter_stream_converts_delimiter_split_across_chunks() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a&")), Ok::<_, object_store::Error>(Bytes::from_static(b"&1\nb&&2\n")), ]); let mut stream = convert_field_delimiter_stream(chunks, "&&".to_string()); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"a,1\nb,2\n"); } #[tokio::test] async fn test_field_delimiter_stream_converts_delimiter_at_stream_end() { let chunks = stream::iter(vec![ Ok::<_, object_store::Error>(Bytes::from_static(b"a&")), Ok::<_, object_store::Error>(Bytes::from_static(b"&")), ]); let mut stream = convert_field_delimiter_stream(chunks, "&&".to_string()); let mut output = Vec::new(); while let Some(bytes) = stream.next().await { output.extend_from_slice(&bytes.unwrap()); } assert_eq!(output, b"a,"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial] async fn test_scan_range_self_overlapping_delimiter_retains_record_context() { let env = crate::storage_api::select_test_ecstore_env().await; let bucket = "s3select-scan-range-record-context"; let object = "input.csv"; env.make_bucket(bucket, false).await; let mut reader = SelectPutObjReader::from_vec(b"111aaa222aa333aa".to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put self-overlapping delimiter ScanRange fixture"); let make_store = |start, end, file_header_info| EcObjectStore { input: 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: "SELECT * FROM s3object".to_string(), expression_type: ExpressionType::from_static(ExpressionType::SQL), input_serialization: InputSerialization { csv: Some(CSVInput { record_delimiter: Some("aa".to_string()), file_header_info, ..Default::default() }), ..Default::default() }, output_serialization: OutputSerialization { csv: Some(CSVOutput::default()), ..Default::default() }, request_progress: None, scan_range: Some(ScanRange { start: Some(start), end: Some(end), }), }, }), need_convert: false, delimiter: String::new(), is_json_document: false, json_sub_path: None, memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, store: Arc::clone(&env.ecstore), }; let store = make_store(6, 6, None); let result = store .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read ScanRange starting inside record data"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming ScanRange payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect record-data ScanRange output"); assert!(chunks.concat().is_empty()); let store = make_store(4, 5, None); let result = store .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read ScanRange starting inside overlapping delimiter"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming overlapping-delimiter payload"); }; let chunks: Vec = stream .try_collect() .await .expect("collect overlapping-delimiter ScanRange output"); assert_eq!(chunks.concat(), b"a222\r\n"); let mut reader = SelectPutObjReader::from_vec(b"111aa222aa333aa".to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put exact review ScanRange fixture"); let result = make_store(6, 6, None) .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read exact review ScanRange fixture"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming exact review ScanRange payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect exact review ScanRange output"); assert!(chunks.concat().is_empty()); let result = make_store(5, 5, None) .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read ScanRange starting after an even delimiter run"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming even-run ScanRange payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect even-run ScanRange output"); assert_eq!(chunks.concat(), b"222\r\n"); let mut reader = SelectPutObjReader::from_vec(b"h1aav1aav2aa".to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put ScanRange header snapshot fixture"); let result = make_store(8, 8, Some(FileHeaderInfo::from_static(FileHeaderInfo::USE))) .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read ScanRange with a separate header read"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming ScanRange header payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect ScanRange header output"); assert_eq!(chunks.concat(), b"h1\r\nv2\r\n"); let header_store = make_store(8, 8, Some(FileHeaderInfo::from_static(FileHeaderInfo::USE))); let header_opts = header_store.object_options(&GetOptions::new()); let header_snapshot = header_store .object_info(&header_opts) .await .expect("read header snapshot before overwrite"); let header_size = validated_object_size(header_snapshot.size).expect("header fixture size should be valid"); let mut reader = SelectPutObjReader::from_vec(b"q9aaz8aaz7aa".to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("overwrite header snapshot fixture"); let err = header_store .read_header_record(header_size, b"aa", &header_opts, &header_snapshot) .await .expect_err("stale header snapshot must fail closed"); assert!(err.to_string().contains("object changed")); let mut reader = SelectPutObjReader::from_vec(b"aaa222aa".to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put object-start delimiter context fixture"); let result = make_store(3, 3, None) .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read delimiter context that reaches the object start"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming object-start context payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect object-start context output"); assert!(chunks.concat().is_empty()); let run_start = SELECT_DEFAULT_READ_BUFFER_SIZE + 7; let mut large_fixture = vec![b'b'; run_start]; large_fixture.extend_from_slice(b"aaa222aa"); let mut reader = SelectPutObjReader::from_vec(large_fixture); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put large self-overlapping delimiter ScanRange fixture"); let scan_start = i64::try_from(run_start + 3).expect("fixture offset should fit in i64"); let result = make_store(scan_start, scan_start, None) .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read large ScanRange with bounded delimiter context"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming large ScanRange payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect large ScanRange output"); assert!(chunks.concat().is_empty()); let mut oversized_run = vec![b'b']; oversized_run.resize(SELECT_DEFAULT_READ_BUFFER_SIZE + 2, b'a'); let scan_start = i64::try_from(oversized_run.len()).expect("fixture offset should fit in i64"); oversized_run.extend_from_slice(b"222aa"); let mut reader = SelectPutObjReader::from_vec(oversized_run); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put oversized delimiter context fixture"); let err = make_store(scan_start, scan_start, None) .get_opts(&Path::from(object), GetOptions::default()) .await .expect_err("oversized self-overlapping delimiter context must fail closed"); assert!(err.to_string().contains("bounded ScanRange context")); let store = make_store(0, 0, None); let opts = store.object_options(&GetOptions::new()); let snapshot = store.object_info(&opts).await.expect("read snapshot before overwrite"); let snapshot_size = usize::try_from(snapshot.size).expect("fixture size should fit in usize"); let mut reader = SelectPutObjReader::from_vec(vec![b'x'; snapshot_size]); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("overwrite snapshot fixture"); let err = store .read_raw_range_with_opts(0..1, &opts, Some(&snapshot)) .await .expect_err("stale ScanRange snapshot must fail closed"); assert!(err.to_string().contains("object changed")); let original = b"111aa222aa333aa"; let mut reader = SelectPutObjReader::from_vec(original.to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("restore context-to-main race fixture"); let (reached_tx, reached_rx) = tokio::sync::oneshot::channel(); let (resume_tx, resume_rx) = tokio::sync::oneshot::channel(); *SCAN_RANGE_BEFORE_MAIN_HOOK.lock().await = Some(ScanRangeBeforeMainHook { bucket: bucket.to_string(), object: object.to_string(), reached: reached_tx, resume: resume_rx, }); let store = make_store(6, 6, None); let read_task = tokio::spawn(async move { store.get_opts(&Path::from("input.csv"), GetOptions::default()).await }); reached_rx .await .expect("ScanRange read should pause before opening its main reader"); let mut reader = SelectPutObjReader::from_vec(b"999aa888aa777aa".to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("overwrite between ScanRange context and main reads"); resume_tx.send(()).expect("resume ScanRange main read"); let err = read_task .await .expect("ScanRange read task should join") .expect_err("context-to-main overwrite must fail closed"); assert!(err.to_string().contains("object changed")); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial] async fn test_self_overlapping_record_delimiter_uses_single_full_file_partition() { const TARGET_PARTITIONS: usize = 4; let mut input_bytes = Vec::with_capacity(SELECT_DEFAULT_READ_BUFFER_SIZE + 8); input_bytes.extend_from_slice(b"0,"); input_bytes.resize(SELECT_DEFAULT_READ_BUFFER_SIZE - 1, b'b'); input_bytes.extend_from_slice(b"aaaX,caa"); assert!(input_bytes.len() > 1024 * 1024); assert_eq!( &input_bytes[SELECT_DEFAULT_READ_BUFFER_SIZE - 1..SELECT_DEFAULT_READ_BUFFER_SIZE + 2], b"aaa" ); let env = crate::storage_api::select_test_ecstore_env().await; let bucket = "s3select-record-delimiter"; let object = "input.csv"; env.make_bucket(bucket, false).await; let mut reader = SelectPutObjReader::from_vec(input_bytes); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put multi-byte record-delimited test object"); let input = 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: "SELECT * FROM s3object".to_string(), expression_type: ExpressionType::from_static(ExpressionType::SQL), input_serialization: InputSerialization { csv: Some(CSVInput { record_delimiter: Some("aa".to_string()), ..Default::default() }), ..Default::default() }, output_serialization: OutputSerialization { csv: Some(CSVOutput::default()), ..Default::default() }, request_progress: None, scan_range: None, }, }); let store = Arc::new(EcObjectStore { input, need_convert: false, delimiter: String::new(), is_json_document: false, json_sub_path: None, memory_pool: Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)), query_tracker: None, store: Arc::clone(&env.ecstore), }); let config = SessionConfig::new() .with_repartition_file_scans(false) .with_repartition_file_min_size(0) .with_target_partitions(TARGET_PARTITIONS); let context = SessionContext::new_with_config(config); let store_url = url::Url::parse(&format!("s3://{bucket}")).expect("valid object store URL"); context.runtime_env().register_object_store(&store_url, store); context .register_csv("records", &format!("s3://{bucket}/{object}"), CsvReadOptions::new().has_header(false)) .await .expect("register partitioned CSV"); let scan_plan = context .sql("SELECT * FROM records") .await .expect("plan partitioned CSV") .create_physical_plan() .await .expect("create partitioned CSV physical plan"); assert_eq!(scan_plan.output_partitioning().partition_count(), 1); let batches = context .sql("SELECT column_1 FROM records") .await .expect("plan exact-result query") .collect() .await .expect("query self-overlapping record-delimited CSV"); let values = batches .iter() .flat_map(|batch| { let column = batch .column(0) .as_any() .downcast_ref::() .expect("mixed first column should be Utf8"); column.iter().map(|value| value.map(str::to_string)).collect::>() }) .collect::>(); assert_eq!(values, vec![Some("0".to_string()), Some("aX".to_string())]); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial] async fn test_get_opts_validates_raw_length_before_delimiter_conversion() { let env = crate::storage_api::select_test_ecstore_env().await; let bucket = "s3select-multi-byte-delimiter"; let object = "input.csv"; let input_bytes = b"a\r\n1^Y"; env.make_bucket(bucket, false).await; let mut reader = SelectPutObjReader::from_vec(input_bytes.to_vec()); env.ecstore .put_object(bucket, object, &mut reader, &Default::default()) .await .expect("put multi-byte-delimited test object"); let input = 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: "SELECT * FROM s3object".to_string(), expression_type: ExpressionType::from_static(ExpressionType::SQL), input_serialization: InputSerialization { csv: Some(CSVInput { field_delimiter: Some("\r\n".to_string()), record_delimiter: Some("^Y".to_string()), ..Default::default() }), ..Default::default() }, output_serialization: OutputSerialization { csv: Some(CSVOutput::default()), ..Default::default() }, request_progress: None, scan_range: None, }, }); let store = super::EcObjectStore { input, need_convert: true, delimiter: "\r\n".to_string(), is_json_document: false, json_sub_path: None, memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, store: Arc::clone(&env.ecstore), }; let result = store .get_opts(&Path::from(object), GetOptions::default()) .await .expect("read multi-byte-delimited test object"); let GetResultPayload::Stream(stream) = result.payload else { panic!("expected streaming object payload"); }; let chunks: Vec = stream.try_collect().await.expect("collect converted object stream"); assert_eq!(chunks.concat(), b"a,1\r\n"); let requested_range = 3..10; let ranges = store .get_ranges(&Path::from(object), std::slice::from_ref(&requested_range)) .await .expect("bounded range past EOF should return the object remainder"); assert_eq!(ranges, vec![Bytes::from_static(b"1^Y")]); } #[tokio::test] async fn test_bytes_stream_stops_at_content_length() { let poll_count = Arc::new(AtomicUsize::new(0)); let stream_poll_count = Arc::clone(&poll_count); let source = stream::unfold(0, move |index| { let stream_poll_count = Arc::clone(&stream_poll_count); async move { stream_poll_count.fetch_add(1, Ordering::SeqCst); let bytes = match index { 0 => Bytes::from_static(b"abcd"), 1 => Bytes::from_static(b"efgh"), _ => return None, }; Some((Ok::<_, std::io::Error>(bytes), index + 1)) } }); let chunks: Vec = bytes_stream(source, 4).try_collect().await.unwrap(); assert_eq!(chunks, vec![Bytes::from_static(b"abcd")]); assert_eq!(poll_count.load(Ordering::SeqCst), 1); } #[tokio::test] async fn test_bytes_stream_rejects_early_eof() { let source = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"ab"))]); let output = bytes_stream(source, 4); futures::pin_mut!(output); assert_eq!(output.next().await.expect("first stream item").expect("first chunk"), b"ab"[..]); let err = output .next() .await .expect("early EOF error") .expect_err("short stream must fail"); let object_store::Error::Generic { store, source } = err else { panic!("expected generic object store error"); }; assert_eq!(store, "EcObjectStore"); let source = source.downcast_ref::().expect("I/O error source"); assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof); assert!(source.to_string().contains("2 bytes remaining")); assert!(output.next().await.is_none()); } #[tokio::test] async fn test_json_document_stream_respects_query_memory_pool() { let input = b"{}".to_vec(); let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER; let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(required - 1)); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, memory_pool, None, ); let err = output .next() .await .expect("memory error") .expect_err("reservation should exceed the pool"); let object_store::Error::Generic { source, .. } = err else { panic!("expected generic object store error"); }; assert!(matches!( source.downcast_ref::(), Some(DataFusionError::ResourcesExhausted(_)) )); } #[tokio::test] async fn test_json_document_stream_releases_memory_reservation() { let input = b"[1,2]".to_vec(); let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER; let memory_pool = Arc::new(GreedyMemoryPool::new(required)); let output: Vec = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, memory_pool.clone(), None, ) .try_collect() .await .expect("JSON conversion should fit the pool"); assert_eq!(output, vec![Bytes::from_static(b"1\n"), Bytes::from_static(b"2\n")]); assert_eq!(memory_pool.reserved(), 0); } #[tokio::test] async fn test_json_document_stream_rejects_early_eof() { let input = b"{}".to_vec(); let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); let mut output = json_document_ndjson_stream(Box::new(std::io::Cursor::new(input)), 4, None, memory_pool, None); let err = output .next() .await .expect("early EOF error") .expect_err("short JSON document must fail"); let object_store::Error::Generic { source, .. } = err else { panic!("expected generic object store error"); }; let source = source.downcast_ref::().expect("I/O error source"); assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof); assert!(source.to_string().contains("2 bytes remaining")); assert!(output.next().await.is_none()); } #[test] fn test_json_document_size_error_is_resource_exhausted() { assert!(validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES).is_ok()); let err = validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES + 1).expect_err("oversized JSON document must fail"); let object_store::Error::Generic { source, .. } = err else { panic!("expected generic object store error"); }; assert!(matches!( source.downcast_ref::(), Some(DataFusionError::ResourcesExhausted(_)) )); } #[test] fn test_json_document_queued_parse_releases_query_guard_when_cancelled() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .max_blocking_threads(1) .enable_all() .build() .expect("build test runtime"); runtime.block_on(async { let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel(); let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel(); let blocker = tokio::task::spawn_blocking(move || { let _ = blocking_started_tx.send(()); release_blocking_rx.recv().expect("release blocking worker"); }); blocking_started_rx.await.expect("blocking worker should start"); let admission = Arc::new(Semaphore::new(1)); let permit = Arc::clone(&admission) .acquire_owned() .await .expect("query permit should be available"); let query_guard: QueryExecutionGuard = Arc::new(permit); let query_tracker = QueryExecutionTracker::new( &QueryExecutionOwner::new(), query_guard, tokio::time::Instant::now() + std::time::Duration::from_secs(30), 30, ); let input = b"{}".to_vec(); let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, memory_pool, Some(query_tracker), ); { let next = output.next(); futures::pin_mut!(next); assert!(futures::poll!(next.as_mut()).is_pending()); } drop(output); let recovered_permit = tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned()) .await .expect("queued JSON parse should be cancelled") .expect("query admission should remain open"); release_blocking_tx.send(()).expect("release blocking worker"); blocker.await.expect("blocking worker should finish"); drop(recovered_permit); assert_eq!(admission.available_permits(), 1); }); } #[test] fn test_json_document_expired_queued_parse_does_not_start() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .max_blocking_threads(1) .enable_all() .build() .expect("build test runtime"); runtime.block_on(async { let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel(); let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel(); let blocker = tokio::task::spawn_blocking(move || { let _ = blocking_started_tx.send(()); release_blocking_rx.recv().expect("release blocking worker"); }); blocking_started_rx.await.expect("blocking worker should start"); let admission = Arc::new(Semaphore::new(1)); let permit = Arc::clone(&admission) .acquire_owned() .await .expect("query permit should be available"); let owner = QueryExecutionOwner::new(); let query_tracker = QueryExecutionTracker::new( &owner, Arc::new(permit), tokio::time::Instant::now() + std::time::Duration::from_secs(30), 30, ); let input = b"{}".to_vec(); let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); let parser_started = Arc::new(std::sync::atomic::AtomicBool::new(false)); let parser_started_in_task = Arc::clone(&parser_started); let mut output = json_document_ndjson_stream_with_parser( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, memory_pool, Some(query_tracker.clone()), move |_, _| { parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst); Ok(vec![Bytes::from_static(b"{}\n")]) }, ); { let next = output.next(); futures::pin_mut!(next); assert!(futures::poll!(next.as_mut()).is_pending()); } query_tracker.expire(&owner); assert_eq!(admission.available_permits(), 1); release_blocking_tx.send(()).expect("release blocking worker"); blocker.await.expect("blocking worker should finish"); let err = tokio::time::timeout(std::time::Duration::from_secs(5), output.next()) .await .expect("queued parser should resume") .expect("queued parser should return an error") .expect_err("expired queued parser must not run"); let object_store::Error::Generic { source, .. } = err else { panic!("expected generic object store error"); }; let source = source.downcast_ref::().expect("I/O error source"); assert_eq!(source.kind(), std::io::ErrorKind::Interrupted); assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst)); }); } #[test] fn test_json_document_started_parse_retains_query_guard_when_cancelled() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .max_blocking_threads(1) .enable_all() .build() .expect("build test runtime"); runtime.block_on(async { let admission = Arc::new(Semaphore::new(1)); let permit = Arc::clone(&admission) .acquire_owned() .await .expect("query permit should be available"); let query_guard: QueryExecutionGuard = Arc::new(permit); let query_tracker = QueryExecutionTracker::new( &QueryExecutionOwner::new(), query_guard, tokio::time::Instant::now() + std::time::Duration::from_secs(30), 30, ); let input = b"{}".to_vec(); let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); let (parse_started_tx, parse_started_rx) = tokio::sync::oneshot::channel(); let (release_parse_tx, release_parse_rx) = std::sync::mpsc::channel(); let mut output = json_document_ndjson_stream_with_parser( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, memory_pool, Some(query_tracker), move |_, _| { let _ = parse_started_tx.send(()); release_parse_rx.recv().expect("release JSON parser"); Ok(vec![Bytes::from_static(b"{}\n")]) }, ); { let next = output.next(); futures::pin_mut!(next); assert!(futures::poll!(next.as_mut()).is_pending()); } parse_started_rx.await.expect("JSON parser should start"); drop(output); assert!(Arc::clone(&admission).try_acquire_owned().is_err()); release_parse_tx.send(()).expect("release JSON parser"); let recovered_permit = tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned()) .await .expect("started JSON parse should release the query guard") .expect("query admission should remain open"); drop(recovered_permit); assert_eq!(admission.available_permits(), 1); }); } /// A JSON array is split into one NDJSON line per element. #[test] fn test_flatten_array_produces_one_line_per_element() { let input = br#"[{"id":1,"name":"Alice"},{"id":2,"name":"Bob"}]"#; let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); let text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 2); // Each line must be valid JSON for line in &lines { serde_json::from_str::(line).expect("each line must be valid JSON"); } // Spot-check field values let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); assert_eq!(first["id"], 1); assert_eq!(first["name"], "Alice"); } /// A single JSON object emits exactly one NDJSON line. #[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 text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 1); let parsed: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); assert_eq!(parsed["id"], 42); assert_eq!(parsed["value"], "hello world"); } /// An empty JSON array produces empty output (zero bytes). #[test] fn test_flatten_empty_array_produces_no_output() { let input = b"[]"; let result = flatten_json_document_to_ndjson(input, None).expect("should succeed"); assert!(result.is_empty(), "empty array should yield zero bytes"); } /// A multi-line (pretty-printed) JSON document is flattened correctly. #[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 text = std::str::from_utf8(&result).unwrap(); assert_eq!(text.lines().count(), 3); } /// Nested objects inside array elements are preserved as compact single-line JSON. #[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 text = std::str::from_utf8(&result).unwrap(); let lines: Vec<&str> = text.lines().collect(); assert_eq!(lines.len(), 2); // Each line must not contain a newline mid-value for line in &lines { assert!(!line.is_empty()); let v: serde_json::Value = serde_json::from_str(line).unwrap(); assert!(v["outer"]["inner"].as_i64().unwrap() >= 99); } } /// Each output line ends with exactly one newline (no blank lines between records). #[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 text = std::str::from_utf8(&result).unwrap(); // Exactly 2 newlines for 2 records assert_eq!(text.chars().filter(|&c| c == '\n').count(), 2); // No leading blank line assert!(!text.starts_with('\n')); } /// Invalid JSON returns an `InvalidData` IO error. #[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"); 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"); 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() { 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 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"); let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); assert_eq!(first["name"], "Alice"); assert_eq!(first["salary"], 75000); let second: serde_json::Value = serde_json::from_str(lines[1]).unwrap(); 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"); } /// Sub-path is ignored when the root is already an array. #[test] fn test_flatten_sub_path_ignored_for_root_array() { 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 ───────────────────────────────────────── #[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())); } #[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())); } #[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); } #[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); } #[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())); } #[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())); } #[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())); } }