diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 910bb9797..536a1f3bd 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -91,6 +91,8 @@ 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 { @@ -111,6 +113,36 @@ pub struct EcObjectStore { 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, @@ -236,6 +268,11 @@ impl EcObjectStore { .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 @@ -261,13 +298,21 @@ impl EcObjectStore { .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) -> Result { + 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) @@ -293,18 +338,24 @@ impl EcObjectStore { } async fn read_raw_range(&self, range: Range) -> Result { - self.read_raw_range_with_opts(range, &self.object_options(&GetOptions::new())) + 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) -> Result { + 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).await?; + 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())); } @@ -314,6 +365,36 @@ impl EcObjectStore { 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 { @@ -388,8 +469,19 @@ fn http_range_spec_from_start(start: u64) -> HTTPRangeSpec { } } -fn scan_range_read_start(scan_range: SelectScanRange, delimiter: &[u8]) -> u64 { - scan_range.start().saturating_sub(delimiter.len() as u64) +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 { @@ -488,31 +580,45 @@ impl ObjectStore for EcObjectStore { 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 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)) + 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((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)); + 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 source_size { - Some(source_size) => source_size, + 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; @@ -525,13 +631,16 @@ impl ObjectStore for EcObjectStore { })?, 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) + 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. // @@ -554,12 +663,14 @@ impl ObjectStore for EcObjectStore { self.query_tracker.clone(), ); GetResultPayload::Stream(stream) - } else if let Some((_, scan_range)) = scan_context { + } 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 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?) + let header = if include_header && read_start > 0 { + Some( + self.read_header_record(original_size, &delimiter, &opts, &source_snapshot) + .await?, + ) } else { None }; @@ -577,13 +688,21 @@ impl ObjectStore for EcObjectStore { } else { stream }; - GetResultPayload::Stream(convert_field_delimiter_stream(stream, self.need_convert.then(|| self.delimiter.clone()))) + 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_field_delimiter_stream(stream, self.need_convert.then(|| self.delimiter.clone()))) + GetResultPayload::Stream(convert_csv_delimiter_stream( + stream, + record_delimiter, + self.need_convert.then(|| self.delimiter.clone()), + )) }; let meta = ObjectMeta { @@ -627,83 +746,92 @@ impl ObjectStore for EcObjectStore { } } -struct DelimiterConverter { - delimiter: Vec, +struct CsvDelimiterConverter { + record_delimiter: Option>, + field_delimiter: Option>, carry: Vec, } -impl DelimiterConverter { - fn new(delimiter: Vec) -> Self { +impl CsvDelimiterConverter { + fn new(record_delimiter: Option>, field_delimiter: Option>) -> Self { Self { - delimiter, + record_delimiter: record_delimiter.filter(|delimiter| !delimiter.is_empty()), + field_delimiter: field_delimiter.filter(|delimiter| !delimiter.is_empty()), carry: Vec::new(), } } - fn convert_chunk(&mut self, chunk: &[u8]) -> Vec { - if self.delimiter.is_empty() { - return chunk.to_vec(); - } + 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.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; - } - } + 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 { - if self.delimiter.is_empty() { - return std::mem::take(&mut self.carry); - } - let converted = replace_symbol(&self.delimiter, &self.carry); + let (converted, _) = self.convert_prefix(&self.carry, self.carry.len()); self.carry.clear(); converted } } -fn replace_symbol(delimiter: &[u8], slice: &[u8]) -> Vec { - 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(stream: S, delimiter: Option) -> BoxStream<'static, Result> +fn convert_delimiter_stream( + stream: S, + record_delimiter: Option>, + field_delimiter: Option>, +) -> BoxStream<'static, Result> where S: Stream> + Send + 'static, { - let Some(delimiter) = delimiter else { - return stream.boxed(); - }; AsyncTryStream::::new(|mut y| async move { - let mut converter = DelimiterConverter::new(delimiter.into_bytes()); + let mut converter = CsvDelimiterConverter::new(record_delimiter, field_delimiter); pin_mut!(stream); while let Some(result) = stream.next().await { let bytes = result?; @@ -721,6 +849,37 @@ where .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, @@ -1153,10 +1312,11 @@ fn incomplete_object_stream_error(remaining: impl std::fmt::Display) -> o_Error #[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, + 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}; @@ -1166,12 +1326,15 @@ mod test { 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, InputSerialization, OutputSerialization, SelectObjectContentInput, - SelectObjectContentRequest, + 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, @@ -1188,18 +1351,34 @@ mod test { } 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()); } + #[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"))]); @@ -1259,6 +1438,110 @@ mod test { 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"))]); @@ -1279,14 +1562,6 @@ mod test { 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)); @@ -1355,7 +1630,7 @@ mod 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 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()); @@ -1369,7 +1644,7 @@ mod test { 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 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()); @@ -1383,7 +1658,7 @@ mod test { 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 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()); @@ -1391,13 +1666,347 @@ mod test { 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&&1\n"; + 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 @@ -1417,7 +2026,8 @@ mod test { expression_type: ExpressionType::from_static(ExpressionType::SQL), input_serialization: InputSerialization { csv: Some(CSVInput { - field_delimiter: Some("&&".to_string()), + field_delimiter: Some("\r\n".to_string()), + record_delimiter: Some("^Y".to_string()), ..Default::default() }), ..Default::default() @@ -1433,7 +2043,7 @@ mod test { let store = super::EcObjectStore { input, need_convert: true, - delimiter: "&&".to_string(), + delimiter: "\r\n".to_string(), is_json_document: false, json_sub_path: None, memory_pool: Arc::new(GreedyMemoryPool::new(1024)), @@ -1450,14 +2060,14 @@ mod test { }; let chunks: Vec = stream.try_collect().await.expect("collect converted object stream"); - assert_eq!(chunks.concat(), b"a,1\n"); + 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\n")]); + assert_eq!(ranges, vec![Bytes::from_static(b"1^Y")]); } #[tokio::test] diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 7c0208206..81824a027 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -349,9 +349,17 @@ impl SessionCtxFactory { }; let rt = RuntimeEnvBuilder::new().with_memory_limit(memory_limit_bytes, 1.0).build()?; let config = SessionConfig::new().with_target_partitions(self.target_partitions); + let custom_two_byte_record_delimiter = context + .input + .request + .input_serialization + .csv + .as_ref() + .and_then(|csv| csv.record_delimiter.as_deref()) + .is_some_and(|delimiter| delimiter.len() == 2 && delimiter.as_bytes() != b"\r\n"); let scan_range_requires_single_file_scan = context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none(); - let config = if scan_range_requires_single_file_scan { + let config = if custom_two_byte_record_delimiter || scan_range_requires_single_file_scan { config.with_repartition_file_scans(false) } else { config @@ -586,6 +594,66 @@ mod tests { assert!(!session.inner().config().options().optimizer.repartition_file_scans); } + #[tokio::test] + async fn csv_scan_range_disables_file_repartitioning() { + let mut context = test_context(); + Arc::make_mut(&mut context.input).request.scan_range = Some(ScanRange { + start: Some(0), + end: Some(0), + }); + + let session = SessionCtxFactory::new(true) + .with_target_partitions(2) + .create_session_ctx(&context) + .await + .expect("CSV ScanRange session should be created"); + + assert!(!session.inner().config().options().optimizer.repartition_file_scans); + } + + #[tokio::test] + async fn two_byte_csv_record_delimiter_disables_file_scan_repartition() { + let mut context = test_context(); + Arc::get_mut(&mut context.input) + .expect("test context input should be uniquely owned") + .request + .input_serialization + .csv + .as_mut() + .expect("test context should use CSV") + .record_delimiter = Some("^Y".to_string()); + let session = SessionCtxFactory::new(true) + .with_target_partitions(4) + .create_session_ctx(&context) + .await + .expect("session should be created"); + + assert!(!session.inner().config().options().optimizer.repartition_file_scans); + } + + #[tokio::test] + async fn crlf_record_delimiter_retains_file_scan_repartition() { + let mut context = test_context(); + Arc::get_mut(&mut context.input) + .expect("test context input should be uniquely owned") + .request + .input_serialization + .csv + .as_mut() + .expect("test context should use CSV") + .record_delimiter = Some("\r\n".to_string()); + let session = SessionCtxFactory::new(true) + .with_target_partitions(4) + .create_session_ctx(&context) + .await + .expect("session should be created"); + + assert_eq!( + session.inner().config().options().optimizer.repartition_file_scans, + SessionConfig::new().options().optimizer.repartition_file_scans + ); + } + #[tokio::test] async fn session_factory_zero_target_partitions_uses_datafusion_default() { let factory = SessionCtxFactory::new(true); diff --git a/rustfs/src/app/select_object.rs b/rustfs/src/app/select_object.rs index 1e039faf3..e64c4bd17 100644 --- a/rustfs/src/app/select_object.rs +++ b/rustfs/src/app/select_object.rs @@ -250,7 +250,8 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> validate_single_byte(csv.comments.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; - validate_record_delimiter(csv.record_delimiter.as_deref())?; + validate_input_record_delimiter(csv.record_delimiter.as_deref())?; + validate_input_delimiter_pair(csv.field_delimiter.as_deref(), csv.record_delimiter.as_deref())?; } if let Some(json) = input.json.as_mut() { @@ -276,7 +277,7 @@ fn normalize_output_serialization(output: &mut OutputSerialization) -> S3Result< validate_single_byte(csv.field_delimiter.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; - validate_record_delimiter(csv.record_delimiter.as_deref())?; + validate_output_record_delimiter(csv.record_delimiter.as_deref())?; if let Some(quote_fields) = csv.quote_fields.as_ref() && !matches!(quote_fields.as_str(), QuoteFields::ALWAYS | QuoteFields::ASNEEDED) { @@ -351,7 +352,7 @@ fn validate_single_byte(value: Option<&str>, code: S3ErrorCode) -> S3Result<()> Ok(()) } -fn validate_record_delimiter(value: Option<&str>) -> S3Result<()> { +fn validate_output_record_delimiter(value: Option<&str>) -> S3Result<()> { if let Some(value) = value && value.len() != 1 && value != "\r\n" @@ -361,6 +362,30 @@ fn validate_record_delimiter(value: Option<&str>) -> S3Result<()> { Ok(()) } +fn validate_input_record_delimiter(value: Option<&str>) -> S3Result<()> { + if let Some(value) = value + && !(1..=2).contains(&value.len()) + { + return Err(S3Error::new(S3ErrorCode::InvalidRequestParameter)); + } + Ok(()) +} + +fn validate_input_delimiter_pair(field_delimiter: Option<&str>, record_delimiter: Option<&str>) -> S3Result<()> { + let field_delimiter = field_delimiter.unwrap_or(","); + let record_delimiter = record_delimiter.unwrap_or("\n"); + let normalized_field_delimiter = if field_delimiter.len() > 1 { "," } else { field_delimiter }; + let normalized_record_delimiter = if record_delimiter.len() == 2 { + "\r\n" + } else { + record_delimiter + }; + if record_delimiter.starts_with(field_delimiter) || normalized_record_delimiter.contains(normalized_field_delimiter) { + return Err(S3Error::new(S3ErrorCode::InvalidRequestParameter)); + } + Ok(()) +} + async fn preflight_select_object(headers: &http::HeaderMap, input: &SelectObjectContentInput) -> S3Result { let opts = get_opts(&input.bucket, &input.key, None, None, headers) .await @@ -1052,6 +1077,120 @@ mod tests { ); } + #[test] + fn validate_accepts_two_byte_csv_input_record_delimiter() { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV") + .record_delimiter = Some("^Y".to_string()); + + assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok()); + } + + #[test] + fn validate_accepts_supported_multibyte_csv_delimiter_pairs() { + let mut input = base_input(); + let csv = input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV"); + csv.field_delimiter = Some("\r\n".to_string()); + csv.record_delimiter = Some("^Y".to_string()); + + assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok()); + + let csv = input + .request + .input_serialization + .csv + .as_mut() + .expect("input should still use CSV"); + csv.field_delimiter = Some("\nX".to_string()); + csv.record_delimiter = None; + assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok()); + + let csv = input + .request + .input_serialization + .csv + .as_mut() + .expect("input should still use CSV"); + csv.field_delimiter = Some("aa".to_string()); + csv.record_delimiter = Some("a".to_string()); + assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok()); + } + + #[test] + fn validate_rejects_csv_input_record_delimiter_outside_one_to_two_bytes() { + for delimiter in ["", "^YZ"] { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV") + .record_delimiter = Some(delimiter.to_string()); + + let err = validate_select_request(&HeaderMap::new(), &mut input) + .expect_err("record delimiter outside the supported length must be rejected"); + + assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter); + } + } + + #[test] + fn validate_rejects_overlapping_csv_input_delimiters() { + for (field_delimiter, record_delimiter) in [ + (Some("^"), Some("^")), + (Some("\r"), Some("\r\n")), + (Some("a"), Some("aa")), + (None, Some(",")), + (Some("\n"), None), + (Some("||"), Some(",")), + (Some("\n"), Some("^Y")), + (Some("\r"), Some("^Y")), + ] { + let mut input = base_input(); + let csv = input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV"); + csv.field_delimiter = field_delimiter.map(str::to_string); + csv.record_delimiter = record_delimiter.map(str::to_string); + + let err = validate_select_request(&HeaderMap::new(), &mut input) + .expect_err("overlapping field and record delimiters must be rejected"); + + assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter); + } + } + + #[test] + fn validate_keeps_csv_output_record_delimiter_restriction() { + let mut input = base_input(); + input + .request + .output_serialization + .csv + .as_mut() + .expect("base output should use CSV") + .record_delimiter = Some("^Y".to_string()); + + let err = + validate_select_request(&HeaderMap::new(), &mut input).expect_err("multi-byte output delimiter must remain rejected"); + + assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter); + } + #[test] fn json_encoder_outputs_line_delimited_records() { let schema =