// 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::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,
    },
};
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> {
    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.";

#[derive(Debug)]
pub struct EcObjectStore {
    input: Arc<SelectObjectContentInput>,
    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.<path>` 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<String>,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,

    store: Arc<SelectStore>,
}

#[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<SelectObjectContentInput>) -> S3Result<Self> {
        Self::build(input, Arc::new(UnboundedMemoryPool::default()), None, None)
    }

    pub(crate) fn new_with_memory_pool(input: Arc<SelectObjectContentInput>, memory_pool: Arc<dyn MemoryPool>) -> S3Result<Self> {
        Self::build(input, memory_pool, None, None)
    }

    pub(crate) fn new_with_query_tracker(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        query_tracker: QueryExecutionTracker,
        store: Option<Arc<SelectStore>>,
    ) -> S3Result<Self> {
        Self::build(input, memory_pool, Some(query_tracker), store)
    }

    fn build(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        query_tracker: Option<QueryExecutionTracker>,
        store: Option<Arc<SelectStore>>,
    ) -> S3Result<Self> {
        let Some(store) = store.or_else(resolve_select_object_store_handle) else {
            return Err(s3_error!(InternalError, "ec store not inited"));
        };

        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<Option<SelectScanRange>> {
        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<u8> {
        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 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<SelectObjectInfo> {
        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<HTTPRangeSpec>, opts: &SelectObjectOptions) -> Result<SelectGetObjectReader> {
        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<u64>, opts: &SelectObjectOptions) -> Result<Bytes> {
        if range.is_empty() {
            return Ok(Bytes::new());
        }
        let reader = self
            .object_reader(Some(http_range_spec_from_range(range.clone())), opts)
            .await?;
        let object_size = validated_object_size(reader.object_info.size)?;
        let resolved_range = GetRange::Bounded(range)
            .as_range(object_size)
            .map_err(|err| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(err),
            })?;
        let expected_size = usize::try_from(resolved_range.end - resolved_range.start).map_err(|err| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        })?;
        let mut reader = reader.stream;
        let mut bytes = Vec::new();
        reader.read_to_end(&mut bytes).await.map_err(|err| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        })?;
        if bytes.len() < expected_size {
            return Err(incomplete_object_stream_error(expected_size - bytes.len()));
        }
        bytes.truncate(expected_size);
        Ok(Bytes::from(bytes))
    }

    async fn read_raw_range(&self, range: Range<u64>) -> Result<Bytes> {
        self.read_raw_range_with_opts(range, &self.object_options(&GetOptions::new()))
            .await
    }

    async fn read_header_record(&self, object_size: u64, delimiter: &[u8], opts: &SelectObjectOptions) -> Result<Bytes> {
        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).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);
        }
    }
}

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<u64>) -> 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 scan_range_read_start(scan_range: SelectScanRange, delimiter: &[u8]) -> u64 {
    scan_range.start().saturating_sub(delimiter.len() as u64)
}

fn find_delimiter(bytes: &[u8], delimiter: &[u8]) -> Option<usize> {
    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<i64>, end: Option<i64>, object_size: u64) -> Result<Option<SelectScanRange>> {
    parse_scan_range_from_bounds(start, end, object_size).map_err(|_| invalid_scan_range_store_error())
}

pub fn validate_scan_range_bounds(
    start: Option<i64>,
    end: Option<i64>,
    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<i64>,
    end: Option<i64>,
    object_size: u64,
) -> std::result::Result<Option<SelectScanRange>, 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<PutResult> {
        Err(unsupported_store_error("put_opts"))
    }

    async fn put_multipart_opts(&self, _location: &Path, _opts: PutMultipartOptions) -> Result<Box<dyn MultipartUpload>> {
        Err(unsupported_store_error("put_multipart_opts"))
    }

    async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
        let opts = self.object_options(&options);
        let needs_scan_context = options.range.is_none() && !options.head && self.input.request.scan_range.is_some();
        let source_size = if needs_scan_context {
            Some(validated_object_size(self.object_info(&opts).await?.size)?)
        } else {
            None
        };
        let scan_context = if needs_scan_context {
            let original_size = source_size.expect("source size is loaded when scan range is present");
            self.scan_range(original_size)?.map(|scan_range| (original_size, scan_range))
        } else {
            None
        };

        let range = options.range.as_ref().map(http_range_spec_from_get_range);
        let reader = if let Some((original_size, scan_range)) = scan_context.as_ref() {
            let delimiter = self.record_delimiter();
            let read_start = scan_range_read_start(*scan_range, &delimiter);
            let range = (*original_size > 0).then(|| http_range_spec_from_start(read_start));
            self.object_reader(range, &opts).await?
        } else {
            self.object_reader(range, &opts).await?
        };

        let original_size = match source_size {
            Some(source_size) => source_size,
            None => validated_object_size(reader.object_info.size)?,
        };
        let etag = reader.object_info.etag;
        let version = reader.object_info.version_id.map(|version| version.to_string());
        let attributes = Attributes::default();
        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 = (result_range.end - result_range.start) as usize;
            let stream = bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), size).boxed();
            GetResultPayload::Stream(stream)
        } 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((_, scan_range)) = scan_context {
            let delimiter = self.record_delimiter();
            let include_header = self.csv_has_header();
            let read_start = scan_range_read_start(scan_range, &delimiter);
            let header = if include_header && scan_range.start() > 0 {
                Some(self.read_header_record(original_size, &delimiter, &opts).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_field_delimiter_stream(stream, 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_field_delimiter_stream(stream, 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<u64>]) -> Result<Vec<Bytes>> {
        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<Path>>) -> BoxStream<'static, Result<Path>> {
        stream::once(ready(Err(unsupported_store_error("delete_stream")))).boxed()
    }

    fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
        stream::once(ready(Err(unsupported_store_error("list")))).boxed()
    }

    async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> Result<ListResult> {
        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 DelimiterConverter {
    delimiter: Vec<u8>,
    carry: Vec<u8>,
}

impl DelimiterConverter {
    fn new(delimiter: Vec<u8>) -> Self {
        Self {
            delimiter,
            carry: Vec::new(),
        }
    }

    fn convert_chunk(&mut self, chunk: &[u8]) -> Vec<u8> {
        if self.delimiter.is_empty() {
            return chunk.to_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.delimiter.len().saturating_sub(1));
        let mut converted = Vec::with_capacity(combined.len());
        let mut pos = 0;
        while pos < safe_end {
            if combined[pos..].starts_with(&self.delimiter) {
                converted.push(DEFAULT_DELIMITER);
                pos += self.delimiter.len();
            } else {
                converted.push(combined[pos]);
                pos += 1;
            }
        }
        self.carry.clear();
        self.carry.extend_from_slice(&combined[pos..]);
        converted
    }

    fn finish(&mut self) -> Vec<u8> {
        if self.delimiter.is_empty() {
            return std::mem::take(&mut self.carry);
        }
        let converted = replace_symbol(&self.delimiter, &self.carry);
        self.carry.clear();
        converted
    }
}

fn replace_symbol(delimiter: &[u8], slice: &[u8]) -> Vec<u8> {
    if delimiter.is_empty() {
        return slice.to_vec();
    }

    let mut result = Vec::with_capacity(slice.len());
    let mut i = 0;
    while i < slice.len() {
        if slice[i..].starts_with(delimiter) {
            result.push(DEFAULT_DELIMITER);
            i += delimiter.len();
        } else {
            result.push(slice[i]);
            i += 1;
        }
    }
    result
}

fn convert_field_delimiter_stream<S>(stream: S, delimiter: Option<String>) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = Result<Bytes>> + Send + 'static,
{
    let Some(delimiter) = delimiter else {
        return stream.boxed();
    };
    AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
        let mut converter = DelimiterConverter::new(delimiter.into_bytes());
        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()
}

struct ScanRangeState<S> {
    stream: S,
    delimiter: Vec<u8>,
    range: SelectScanRange,
    include_header: bool,
    offset: u64,
    record_start: u64,
    record: Vec<u8>,
    pending: VecDeque<Bytes>,
    done: bool,
    expected_end: u64,
}

fn scan_range_stream<S>(
    stream: S,
    delimiter: Vec<u8>,
    range: SelectScanRange,
    include_header: bool,
    base_offset: u64,
    expected_end: u64,
) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = std::io::Result<Bytes>> + Send + Unpin + 'static,
{
    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<S> ScanRangeState<S> {
    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::<Vec<_>>();
        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<u8>, 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;
            }
        }
    }
}

/// Extract the JSON sub-path from a SQL expression's FROM clause.
///
/// Given `SELECT e.name FROM s3object.employees e WHERE …` this returns
/// `Some("employees")`.  Returns `None` when the FROM target is plain
/// `s3object` (no sub-path) or when the expression cannot be parsed.
fn extract_json_sub_path_from_expression(expression: &str) -> Option<String> {
    // Find " FROM " (case-insensitive).
    let lower = expression.to_lowercase();
    let from_pos = lower.find(" from ")?;
    let after_from = expression[from_pos + 6..].trim_start();

    // Must start with "s3object" (case-insensitive, ASCII-only for the prefix).
    const S3OBJECT_LOWER: &str = "s3object";
    let mut chars = after_from.char_indices();
    for expected in S3OBJECT_LOWER.chars() {
        let (idx, actual) = chars.next()?;
        if actual.to_ascii_lowercase() != expected {
            return None;
        }
        // When we have consumed the full prefix, `idx` is the byte index of
        // the current character; use it plus its UTF-8 length as the slice
        // boundary for the remaining string.
        if expected == 't' {
            let end_of_prefix = idx + actual.len_utf8();
            let after_s3object = &after_from[end_of_prefix..];

            // If the very next character is '.' there is a sub-path.
            if let Some(rest) = after_s3object.strip_prefix('.') {
                let rest = rest.trim_start();
                if rest.is_empty() {
                    return None;
                }

                // Support quoted identifiers: s3object."my.path" or s3object.'my path'
                let mut chars = rest.chars();
                if let Some(first) = chars.next()
                    && (first == '"' || first == '\'')
                {
                    let quote = first;
                    let inner = &rest[first.len_utf8()..];
                    if let Some(end) = inner.find(quote) {
                        let path = &inner[..end];
                        if !path.trim().is_empty() {
                            return Some(path.to_string());
                        }
                    }
                    // Quoted but no closing quote or empty: treat as no sub-path.
                    return None;
                }

                // Unquoted identifier: collect characters until whitespace, '[', or ']'.
                let end = rest
                    .find(|c: char| c.is_whitespace() || c == '[' || c == ']')
                    .unwrap_or(rest.len());
                let path = rest[..end].trim();
                if !path.is_empty() {
                    return Some(path.to_string());
                }
            }
            return None;
        }
    }

    // We only reach here if the loop completed without hitting the 't'
    // branch above, which would be unexpected given S3OBJECT_LOWER.
    None
}

/// 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<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
    original_size: u64,
    json_sub_path: Option<String>,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>> {
    json_document_ndjson_stream_with_parser(
        stream,
        original_size,
        json_sub_path,
        memory_pool,
        query_tracker,
        |all_bytes, json_sub_path| parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()),
    )
}

fn json_document_ndjson_stream_with_parser<P>(
    stream: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
    original_size: u64,
    json_sub_path: Option<String>,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,
    parser: P,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>>
where
    P: FnOnce(Vec<u8>, Option<String>) -> std::io::Result<Vec<Bytes>> + Send + 'static,
{
    AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
        // Compact JSON can expand substantially into a serde_json DOM and
        // per-record output buffers, so reserve a conservative upper bound
        // before the source buffer is allocated.
        let buffer_capacity = usize::try_from(original_size).map_err(|_| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(DataFusionError::ResourcesExhausted(format!(
                "JSON DOCUMENT input size {original_size} does not fit in memory"
            ))),
        })?;
        let reservation_bytes = buffer_capacity
            .checked_mul(JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)
            .ok_or_else(|| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(DataFusionError::ResourcesExhausted(format!(
                    "JSON DOCUMENT memory reservation overflow for {original_size} input bytes"
                ))),
            })?;
        let reservation = MemoryConsumer::new("S3 Select JSON document").register(&memory_pool);
        reservation.try_resize(reservation_bytes).map_err(|err| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        })?;

        pin_mut!(stream);
        // ── 1. Read phase (lazy: only runs when the stream is polled) ────
        let mut all_bytes = Vec::with_capacity(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<Mutex<QueryExecutionGuardState>>,
}

enum QueryExecutionGuardState {
    Pending(Option<QueryExecutionTracker>),
    Started,
    Cancelled,
}

impl PendingQueryExecutionGuard {
    fn new(query_tracker: Option<QueryExecutionTracker>) -> Self {
        Self {
            state: Arc::new(Mutex::new(QueryExecutionGuardState::Pending(query_tracker))),
        }
    }

    fn task_state(&self) -> Arc<Mutex<QueryExecutionGuardState>> {
        Arc::clone(&self.state)
    }

    fn start(state: &Mutex<QueryExecutionGuardState>) -> std::io::Result<Option<QueryExecutionGuard>> {
        let mut state = state.lock();
        match std::mem::replace(&mut *state, QueryExecutionGuardState::Started) {
            QueryExecutionGuardState::Pending(None) => Ok(None),
            QueryExecutionGuardState::Pending(Some(query_tracker)) => query_tracker.query_guard().map(Some).ok_or_else(|| {
                std::io::Error::new(std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started")
            }),
            QueryExecutionGuardState::Cancelled => {
                *state = QueryExecutionGuardState::Cancelled;
                Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    "JSON DOCUMENT parse was cancelled before it started",
                ))
            }
            QueryExecutionGuardState::Started => {
                *state = QueryExecutionGuardState::Started;
                Err(std::io::Error::other("JSON DOCUMENT parse started more than once"))
            }
        }
    }
}

impl Drop for PendingQueryExecutionGuard {
    fn drop(&mut self) {
        let query_guard = {
            let mut state = self.state.lock();
            match std::mem::replace(&mut *state, QueryExecutionGuardState::Cancelled) {
                QueryExecutionGuardState::Pending(query_guard) => query_guard,
                QueryExecutionGuardState::Started => {
                    *state = QueryExecutionGuardState::Started;
                    None
                }
                QueryExecutionGuardState::Cancelled => None,
            }
        };
        drop(query_guard);
    }
}

/// Parse a JSON DOCUMENT (a single JSON value, possibly multi-line) into a
/// list of NDJSON lines – one [`Bytes`] per record.
///
/// `json_sub_path` – when the SQL expression contains `FROM s3object.<key>`,
/// 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<Vec<Bytes>> {
    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<Bytes> = 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<Bytes> {
    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<S>(stream: S, content_length: usize) -> impl Stream<Item = Result<Bytes>> + Send + 'static
where
    S: Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static,
{
    AsyncTryStream::<Bytes, o_Error, _>::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, SelectScanRange, bytes_stream,
        convert_field_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter, flatten_json_document_to_ndjson,
        http_range_spec_from_get_range, json_document_ndjson_stream, json_document_ndjson_stream_with_parser, replace_symbol,
        scan_range_from_bounds, scan_range_read_start, scan_range_stream, select_read_headers, validate_json_document_size,
        validated_object_size,
    };
    use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker};
    use crate::storage_api::SelectPutObjReader;
    use crate::storage_api::object_store::ObjectIO as _;
    use bytes::Bytes;
    use datafusion::{
        common::DataFusionError,
        execution::memory_pool::{GreedyMemoryPool, MemoryPool},
        object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path},
    };
    use futures::{StreamExt, TryStreamExt, stream};
    use s3s::dto::{
        CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput,
        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<SelectObjectContentInput>) -> s3s::S3Result<EcObjectStore> = EcObjectStore::new;
    }
    use tokio::sync::Semaphore;

    #[test]
    fn test_replace() {
        let result = replace_symbol(b"&&", b"dandan&&is&&best");
        assert_eq!(result, b"dandan,is,best");
    }

    #[test]
    fn test_validated_object_size_rejects_negative_metadata() {
        assert_eq!(validated_object_size(0).expect("zero object size should be valid"), 0);
        assert!(validated_object_size(-1).is_err());
    }

    #[tokio::test]
    async fn test_scan_range_stream_keeps_header_and_selected_record() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]);
        let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0, 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_rejects_early_eof() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n"))]);
        let mut output = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 7), false, 0, 8);

        assert_eq!(output.next().await.expect("first stream item").expect("first record"), b"1,a\n"[..]);
        let err = output
            .next()
            .await
            .expect("early EOF error")
            .expect_err("short ScanRange stream must fail");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
        assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
        assert!(source.to_string().contains("4 bytes remaining"));
        assert!(output.next().await.is_none());
    }

    #[test]
    fn test_scan_range_read_start_keeps_full_delimiter_boundary() {
        let range = SelectScanRange::new(10, 20);
        assert_eq!(scan_range_read_start(range, b"\n"), 9);
        assert_eq!(scan_range_read_start(range, b"\r\n"), 8);
        assert_eq!(scan_range_read_start(range, b"abcdef"), 4);
    }

    #[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, Some("&&".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, Some("&&".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, Some("&&".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_get_opts_validates_raw_length_before_delimiter_conversion() {
        let temp_root = tempfile::tempdir().expect("create s3select test temp root");
        let env = rustfs_test_utils::TestECStoreEnv::builder()
            .base_dir(temp_root.path())
            .build()
            .await;
        let bucket = "s3select-multi-byte-delimiter";
        let object = "input.csv";
        let input_bytes = b"a&&1\n";
        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("&&".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: "&&".to_string(),
            is_json_document: false,
            json_sub_path: None,
            memory_pool: Arc::new(GreedyMemoryPool::new(1024)),
            query_tracker: None,
            store: 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<Bytes> = stream.try_collect().await.expect("collect converted object stream");

        assert_eq!(chunks.concat(), b"a,1\n");

        let requested_range = 3..10;
        let ranges = store
            .get_ranges(&Path::from(object), std::slice::from_ref(&requested_range))
            .await
            .expect("bounded range past EOF should return the object remainder");
        assert_eq!(ranges, vec![Bytes::from_static(b"1\n")]);
    }

    #[tokio::test]
    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> = 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::<std::io::Error>().expect("I/O error source");
        assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
        assert!(source.to_string().contains("2 bytes remaining"));
        assert!(output.next().await.is_none());
    }

    #[tokio::test]
    async fn test_json_document_stream_respects_query_memory_pool() {
        let input = b"{}".to_vec();
        let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
        let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(required - 1));
        let mut output = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            None,
            memory_pool,
            None,
        );

        let err = output
            .next()
            .await
            .expect("memory error")
            .expect_err("reservation should exceed the pool");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        assert!(matches!(
            source.downcast_ref::<DataFusionError>(),
            Some(DataFusionError::ResourcesExhausted(_))
        ));
    }

    #[tokio::test]
    async fn test_json_document_stream_releases_memory_reservation() {
        let input = b"[1,2]".to_vec();
        let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
        let memory_pool = Arc::new(GreedyMemoryPool::new(required));
        let output: Vec<Bytes> = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            None,
            memory_pool.clone(),
            None,
        )
        .try_collect()
        .await
        .expect("JSON conversion should fit the pool");

        assert_eq!(output, vec![Bytes::from_static(b"1\n"), Bytes::from_static(b"2\n")]);
        assert_eq!(memory_pool.reserved(), 0);
    }

    #[tokio::test]
    async fn test_json_document_stream_rejects_early_eof() {
        let input = b"{}".to_vec();
        let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
        let mut output = json_document_ndjson_stream(Box::new(std::io::Cursor::new(input)), 4, None, memory_pool, None);

        let err = output
            .next()
            .await
            .expect("early EOF error")
            .expect_err("short JSON document must fail");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
        assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
        assert!(source.to_string().contains("2 bytes remaining"));
        assert!(output.next().await.is_none());
    }

    #[test]
    fn test_json_document_size_error_is_resource_exhausted() {
        assert!(validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES).is_ok());

        let err = validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES + 1).expect_err("oversized JSON document must fail");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        assert!(matches!(
            source.downcast_ref::<DataFusionError>(),
            Some(DataFusionError::ResourcesExhausted(_))
        ));
    }

    #[test]
    fn test_json_document_queued_parse_releases_query_guard_when_cancelled() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
            let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
            let blocker = tokio::task::spawn_blocking(move || {
                let _ = blocking_started_tx.send(());
                release_blocking_rx.recv().expect("release blocking worker");
            });
            blocking_started_rx.await.expect("blocking worker should start");

            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let query_guard: QueryExecutionGuard = Arc::new(permit);
            let query_tracker = QueryExecutionTracker::new(
                &QueryExecutionOwner::new(),
                query_guard,
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            let input = b"{}".to_vec();
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let mut output = json_document_ndjson_stream(
                Box::new(std::io::Cursor::new(input.clone())),
                input.len() as u64,
                None,
                memory_pool,
                Some(query_tracker),
            );

            {
                let next = output.next();
                futures::pin_mut!(next);
                assert!(futures::poll!(next.as_mut()).is_pending());
            }
            drop(output);

            let recovered_permit =
                tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
                    .await
                    .expect("queued JSON parse should be cancelled")
                    .expect("query admission should remain open");
            release_blocking_tx.send(()).expect("release blocking worker");
            blocker.await.expect("blocking worker should finish");
            drop(recovered_permit);
            assert_eq!(admission.available_permits(), 1);
        });
    }

    #[test]
    fn test_json_document_expired_queued_parse_does_not_start() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
            let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
            let blocker = tokio::task::spawn_blocking(move || {
                let _ = blocking_started_tx.send(());
                release_blocking_rx.recv().expect("release blocking worker");
            });
            blocking_started_rx.await.expect("blocking worker should start");

            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let owner = QueryExecutionOwner::new();
            let query_tracker = QueryExecutionTracker::new(
                &owner,
                Arc::new(permit),
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            let input = b"{}".to_vec();
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let parser_started = Arc::new(std::sync::atomic::AtomicBool::new(false));
            let parser_started_in_task = Arc::clone(&parser_started);
            let mut output = json_document_ndjson_stream_with_parser(
                Box::new(std::io::Cursor::new(input.clone())),
                input.len() as u64,
                None,
                memory_pool,
                Some(query_tracker.clone()),
                move |_, _| {
                    parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst);
                    Ok(vec![Bytes::from_static(b"{}\n")])
                },
            );

            {
                let next = output.next();
                futures::pin_mut!(next);
                assert!(futures::poll!(next.as_mut()).is_pending());
            }
            query_tracker.expire(&owner);
            assert_eq!(admission.available_permits(), 1);
            release_blocking_tx.send(()).expect("release blocking worker");
            blocker.await.expect("blocking worker should finish");

            let err = tokio::time::timeout(std::time::Duration::from_secs(5), output.next())
                .await
                .expect("queued parser should resume")
                .expect("queued parser should return an error")
                .expect_err("expired queued parser must not run");
            let object_store::Error::Generic { source, .. } = err else {
                panic!("expected generic object store error");
            };
            let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
            assert_eq!(source.kind(), std::io::ErrorKind::Interrupted);
            assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst));
        });
    }

    #[test]
    fn test_json_document_started_parse_retains_query_guard_when_cancelled() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let query_guard: QueryExecutionGuard = Arc::new(permit);
            let query_tracker = QueryExecutionTracker::new(
                &QueryExecutionOwner::new(),
                query_guard,
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            let input = b"{}".to_vec();
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let (parse_started_tx, parse_started_rx) = tokio::sync::oneshot::channel();
            let (release_parse_tx, release_parse_rx) = std::sync::mpsc::channel();
            let mut output = json_document_ndjson_stream_with_parser(
                Box::new(std::io::Cursor::new(input.clone())),
                input.len() as u64,
                None,
                memory_pool,
                Some(query_tracker),
                move |_, _| {
                    let _ = parse_started_tx.send(());
                    release_parse_rx.recv().expect("release JSON parser");
                    Ok(vec![Bytes::from_static(b"{}\n")])
                },
            );

            {
                let next = output.next();
                futures::pin_mut!(next);
                assert!(futures::poll!(next.as_mut()).is_pending());
            }
            parse_started_rx.await.expect("JSON parser should start");
            drop(output);

            assert!(Arc::clone(&admission).try_acquire_owned().is_err());
            release_parse_tx.send(()).expect("release JSON parser");
            let recovered_permit =
                tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
                    .await
                    .expect("started JSON parse should release the query guard")
                    .expect("query admission should remain open");
            drop(recovered_permit);
            assert_eq!(admission.available_permits(), 1);
        });
    }

    /// A JSON array is split into one NDJSON line per element.
    #[test]
    fn test_flatten_array_produces_one_line_per_element() {
        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::<serde_json::Value>(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_with_bracket() {
        // `FROM s3object.employees[*]` — bracket stops path collection.
        let sql = "SELECT e.name FROM s3object.employees[*] e";
        assert_eq!(extract_json_sub_path_from_expression(sql), Some("employees".to_string()));
    }
}