Skip to main content

faucet_source_s3/
stream.rs

1//! S3 source stream executor.
2
3use crate::config::{S3FileFormat, S3SourceConfig};
4use async_trait::async_trait;
5use aws_sdk_s3::Client;
6use faucet_core::shard::ShardSpec;
7use faucet_core::{FaucetError, Stream, StreamPage};
8use futures::stream::{self, StreamExt, TryStreamExt};
9use serde_json::Value;
10use std::pin::Pin;
11use std::sync::Mutex;
12use tokio::io::AsyncBufReadExt;
13
14/// An S3 source that lists and reads objects from a bucket.
15pub struct S3Source {
16    config: S3SourceConfig,
17    client: Client,
18    /// Shard applied by the cluster coordinator (Mode B): `(shards, index)`.
19    /// `None` (or `shards <= 1`) reads every listed object. Stored behind a
20    /// `Mutex` so `apply_shard(&self, …)` can record it before streaming.
21    applied_shard: Mutex<Option<(usize, usize)>>,
22}
23
24/// Stable FNV-1a hash of an object key, used to assign keys to shards.
25///
26/// Deterministic across processes and platforms (all cluster workers run the
27/// identical binary and this fixed algorithm), so every worker maps a given key
28/// to the same shard index — the partition is disjoint and complete.
29fn shard_hash(key: &str) -> u64 {
30    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
31    for b in key.as_bytes() {
32        h ^= *b as u64;
33        h = h.wrapping_mul(0x0000_0100_0000_01b3);
34    }
35    h
36}
37
38impl S3Source {
39    /// Create a new S3 source from the given configuration.
40    ///
41    /// Builds the S3 client eagerly so it is reused across calls.
42    pub async fn new(config: S3SourceConfig) -> Result<Self, FaucetError> {
43        let client = Self::build_client(&config).await?;
44        Ok(Self {
45            config,
46            client,
47            applied_shard: Mutex::new(None),
48        })
49    }
50
51    /// Retain only the keys belonging to the applied shard (hash-of-key modulo
52    /// `shards`). A no-op when no shard is applied or `shards <= 1`.
53    fn shard_filter(&self, keys: Vec<String>) -> Vec<String> {
54        match *self.applied_shard.lock().expect("shard mutex poisoned") {
55            Some((shards, index)) if shards > 1 => keys
56                .into_iter()
57                .filter(|k| (shard_hash(k) % shards as u64) == index as u64)
58                .collect(),
59            _ => keys,
60        }
61    }
62
63    /// Build an S3 client from the configuration.
64    async fn build_client(config: &S3SourceConfig) -> Result<Client, FaucetError> {
65        let mut config_loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
66
67        if let Some(ref region) = config.region {
68            config_loader = config_loader.region(aws_config::Region::new(region.clone()));
69        }
70
71        if let Some(ref endpoint) = config.endpoint_url {
72            config_loader = config_loader.endpoint_url(endpoint);
73        }
74
75        let sdk_config = config_loader.load().await;
76        let client = Client::new(&sdk_config);
77        Ok(client)
78    }
79
80    /// List object keys matching the configured bucket and prefix.
81    ///
82    /// When `prefix_override` is `Some`, it is used instead of `self.config.prefix`
83    /// (used for parent-context substitution).
84    async fn list_object_keys(
85        &self,
86        prefix_override: Option<&str>,
87    ) -> Result<Vec<String>, FaucetError> {
88        let mut keys = Vec::new();
89        let mut continuation_token: Option<String> = None;
90
91        let effective_prefix = prefix_override.or(self.config.prefix.as_deref());
92
93        loop {
94            let mut req = self.client.list_objects_v2().bucket(&self.config.bucket);
95
96            if let Some(prefix) = effective_prefix {
97                req = req.prefix(prefix);
98            }
99
100            if let Some(ref token) = continuation_token {
101                req = req.continuation_token(token);
102            }
103
104            let response = req.send().await.map_err(|e| {
105                FaucetError::Source(format!(
106                    "S3 list objects error for bucket '{}': {e}",
107                    self.config.bucket
108                ))
109            })?;
110
111            for object in response.contents() {
112                let key: &str = object.key().unwrap_or_default();
113                if key.is_empty() {
114                    continue;
115                }
116                keys.push(key.to_string());
117
118                if let Some(max) = self.config.max_objects
119                    && keys.len() >= max
120                {
121                    return Ok(self.shard_filter(keys));
122                }
123            }
124
125            if response.is_truncated() == Some(true) {
126                continuation_token = response.next_continuation_token().map(String::from);
127            } else {
128                break;
129            }
130        }
131
132        Ok(self.shard_filter(keys))
133    }
134
135    /// Read and parse a single S3 object into records.
136    async fn read_object(&self, key: &str) -> Result<Vec<Value>, FaucetError> {
137        let text = self.read_object_text(key).await?;
138        self.parse_content(key, &text)
139    }
140
141    /// Read the full body of a single S3 object into a UTF-8 `String`.
142    ///
143    /// Streams the (optionally decompressed) body straight into one `String`
144    /// via [`open_object_reader`](Self::open_object_reader) rather than
145    /// buffering the raw bytes AND the decompressed bytes AND the `String`
146    /// at once (#78/#25). The whole object is still one unit for
147    /// `JsonArray` / `RawText`, but peak memory is now ~1× the decoded size.
148    async fn read_object_text(&self, key: &str) -> Result<String, FaucetError> {
149        use tokio::io::AsyncReadExt as _;
150        let mut reader = self.open_object_reader(key).await?;
151        let mut text = String::new();
152        reader.read_to_string(&mut text).await.map_err(|e| {
153            FaucetError::Source(format!(
154                "S3 read/decode error for key '{key}' (not valid UTF-8?): {e}"
155            ))
156        })?;
157        Ok(text)
158    }
159
160    /// Open an S3 object as an [`AsyncBufRead`](tokio::io::AsyncBufRead) over
161    /// its body. Used by [`Source::stream_pages`](faucet_core::Source::stream_pages)
162    /// to decode `JsonLines` objects line-by-line without buffering the
163    /// whole file.
164    async fn open_object_reader(
165        &self,
166        key: &str,
167    ) -> Result<std::pin::Pin<Box<dyn tokio::io::AsyncBufRead + Send + Unpin>>, FaucetError> {
168        let mut request = self
169            .client
170            .get_object()
171            .bucket(&self.config.bucket)
172            .key(key);
173        // Ask S3 to return its stored checksum so we can verify the body (#161).
174        if self.config.verify_checksum {
175            request = request.checksum_mode(aws_sdk_s3::types::ChecksumMode::Enabled);
176        }
177        let response = request.send().await.map_err(|e| {
178            FaucetError::Source(format!("S3 get object error for key '{key}': {e}"))
179        })?;
180
181        // Read all metadata BEFORE consuming `body` (which partially moves
182        // `response`), so a cleanly-truncated/corrupted transfer is rejected
183        // rather than silently parsed as a complete object (#161).
184        let mut checks: Vec<Box<dyn faucet_core::IntegrityCheck>> = Vec::new();
185        match crate::verify::length_check(response.content_length(), self.config.verify_length) {
186            Some(check) => checks.push(check),
187            None if self.config.verify_length => tracing::debug!(
188                key = %key,
189                "S3 object reports no Content-Length; length verification skipped"
190            ),
191            None => {}
192        }
193        if self.config.verify_checksum {
194            let advertised = crate::verify::S3Checksums {
195                crc32: response.checksum_crc32().map(str::to_string),
196                crc32c: response.checksum_crc32_c().map(str::to_string),
197                crc64nvme: response.checksum_crc64_nvme().map(str::to_string),
198                sha256: response.checksum_sha256().map(str::to_string),
199                etag: response.e_tag().map(str::to_string),
200            };
201            match crate::verify::checksum_check(&advertised) {
202                Some(check) => checks.push(check),
203                None => tracing::warn!(
204                    key = %key,
205                    "verify_checksum is enabled but S3 advertised no verifiable checksum for \
206                     this object; relying on the length check only"
207                ),
208            }
209        }
210
211        // `ByteStream::into_async_read` returns `impl AsyncRead`. Wrap the RAW
212        // body in the verifier first so length/checksum cover the stored bytes
213        // (below any decompression), then `BufReader` so `.lines()` is usable
214        // and ownership is `Unpin`.
215        let verified = faucet_core::VerifyingReader::new(response.body.into_async_read(), checks);
216        let buffered = tokio::io::BufReader::new(verified);
217        #[cfg(feature = "compression")]
218        {
219            let codec = self.config.compression.resolve(key);
220            faucet_core::compression::warn_mismatch(key, codec);
221            Ok(faucet_core::compression::wrap_async_reader(buffered, codec))
222        }
223        #[cfg(not(feature = "compression"))]
224        {
225            Ok(Box::pin(buffered))
226        }
227    }
228
229    /// Parse file content into records based on the configured file format.
230    fn parse_content(&self, key: &str, text: &str) -> Result<Vec<Value>, FaucetError> {
231        match self.config.file_format {
232            S3FileFormat::JsonLines => {
233                let mut records = Vec::new();
234                for (line_num, line) in text.lines().enumerate() {
235                    let trimmed = line.trim();
236                    if trimmed.is_empty() {
237                        continue;
238                    }
239                    let value: Value = serde_json::from_str(trimmed).map_err(|e| {
240                        FaucetError::Source(format!(
241                            "S3 JSON parse error in '{key}' at line {}: {e}",
242                            line_num + 1
243                        ))
244                    })?;
245                    records.push(value);
246                }
247                Ok(records)
248            }
249            S3FileFormat::JsonArray => {
250                let value: Value = serde_json::from_str(text).map_err(|e| {
251                    FaucetError::Source(format!("S3 JSON parse error in '{key}': {e}"))
252                })?;
253                match value {
254                    Value::Array(arr) => Ok(arr),
255                    _ => Err(FaucetError::Source(format!(
256                        "S3 expected JSON array in '{key}', got {}",
257                        value_type_name(&value)
258                    ))),
259                }
260            }
261            S3FileFormat::RawText => {
262                let record = serde_json::json!({
263                    "key": key,
264                    "content": text,
265                });
266                Ok(vec![record])
267            }
268        }
269    }
270}
271
272#[async_trait]
273impl faucet_core::Source for S3Source {
274    async fn fetch_with_context(
275        &self,
276        context: &std::collections::HashMap<String, serde_json::Value>,
277    ) -> Result<Vec<Value>, FaucetError> {
278        // Substitute context into prefix when parent context is provided.
279        let substituted_prefix: Option<String> = if !context.is_empty() {
280            self.config
281                .prefix
282                .as_ref()
283                .map(|p| faucet_core::util::substitute_context(p, context))
284        } else {
285            None
286        };
287
288        let keys = self.list_object_keys(substituted_prefix.as_deref()).await?;
289
290        tracing::info!(
291            bucket = %self.config.bucket,
292            objects = keys.len(),
293            "Listed S3 objects"
294        );
295
296        let concurrency = self.config.concurrency.max(1);
297
298        let results: Vec<Vec<Value>> = stream::iter(keys)
299            .map(|key| async move {
300                let records = self.read_object(&key).await?;
301                tracing::debug!(key = %key, records = records.len(), "Read S3 object");
302                Ok::<Vec<Value>, FaucetError>(records)
303            })
304            .buffer_unordered(concurrency)
305            .try_collect()
306            .await?;
307
308        let all_records: Vec<Value> = results.into_iter().flatten().collect();
309
310        tracing::info!(total_records = all_records.len(), "S3 fetch complete");
311        Ok(all_records)
312    }
313
314    /// Stream records from listed S3 objects without buffering the full
315    /// scan. Each emitted [`StreamPage`] holds up to
316    /// [`S3SourceConfig::batch_size`] records.
317    ///
318    /// The trait-level `batch_size` argument is ignored in favour of the
319    /// config field — the config is the user-facing knob the README
320    /// documents, and routing the pipeline-supplied hint through it would
321    /// silently override an explicit config value.
322    ///
323    /// Behaviour by format:
324    ///
325    /// - `JsonLines` / `RawText`: the object body is decoded line-by-line
326    ///   via [`tokio::io::AsyncBufReadExt::lines`] so client-side memory is
327    ///   bounded at `O(batch_size)` per object. Multi-object scans are
328    ///   flattened — a single page may carry lines drawn from any object.
329    /// - `JsonArray`: each object is buffered fully (the JSON value can
330    ///   only be parsed once the array is complete) and then its records
331    ///   are chunked into pages of `batch_size`. See the README "Streaming
332    ///   and batching" section for the caveat.
333    ///
334    /// `batch_size = 0` is the "no batching" sentinel: one [`StreamPage`]
335    /// is emitted per S3 object (no within-object chunking and no
336    /// cross-object accumulation). The S3 source has no
337    /// incremental-replication mode today, so every emitted page carries
338    /// `bookmark: None`.
339    fn stream_pages<'a>(
340        &'a self,
341        context: &'a std::collections::HashMap<String, Value>,
342        _batch_size: usize,
343    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
344        let batch_size = self.config.batch_size;
345
346        Box::pin(async_stream::try_stream! {
347            // Substitute context into prefix when parent context is provided.
348            let substituted_prefix: Option<String> = if !context.is_empty() {
349                self.config
350                    .prefix
351                    .as_ref()
352                    .map(|p| faucet_core::util::substitute_context(p, context))
353            } else {
354                None
355            };
356
357            let keys = self.list_object_keys(substituted_prefix.as_deref()).await?;
358            tracing::info!(
359                bucket = %self.config.bucket,
360                objects = keys.len(),
361                "Listed S3 objects (stream)",
362            );
363
364            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
365            let initial_capacity = if batch_size == 0 { 1024 } else { batch_size };
366            let mut buffer: Vec<Value> = Vec::with_capacity(initial_capacity);
367            let mut total = 0usize;
368
369            for key in &keys {
370                match self.config.file_format {
371                    S3FileFormat::JsonLines => {
372                        let reader = self.open_object_reader(key).await?;
373                        let mut lines = reader.lines();
374                        let mut line_num: usize = 0;
375                        while let Some(line) = lines
376                            .next_line()
377                            .await
378                            .map_err(|e| FaucetError::Source(format!(
379                                "S3 read body error for key '{key}': {e}"
380                            )))?
381                        {
382                            line_num += 1;
383                            let trimmed = line.trim();
384                            if trimmed.is_empty() {
385                                continue;
386                            }
387                            let value: Value =
388                                serde_json::from_str(trimmed).map_err(|e| {
389                                    FaucetError::Source(format!(
390                                        "S3 JSON parse error in '{key}' at line {line_num}: {e}",
391                                    ))
392                                })?;
393                            buffer.push(value);
394                            if batch_size != 0 && buffer.len() >= chunk {
395                                let page = std::mem::replace(
396                                    &mut buffer,
397                                    Vec::with_capacity(initial_capacity),
398                                );
399                                total += page.len();
400                                yield StreamPage { records: page, bookmark: None };
401                            }
402                        }
403                        if batch_size == 0 && !buffer.is_empty() {
404                            let page = std::mem::take(&mut buffer);
405                            total += page.len();
406                            yield StreamPage { records: page, bookmark: None };
407                        }
408                    }
409                    S3FileFormat::RawText => {
410                        // RawText emits a single record per object; the
411                        // `key` + `content` shape is unchanged so we
412                        // continue to buffer the body fully. This still
413                        // streams *across* objects.
414                        let text = self.read_object_text(key).await?;
415                        let record = serde_json::json!({
416                            "key": key,
417                            "content": text,
418                        });
419                        buffer.push(record);
420                        if batch_size == 0 {
421                            let page = std::mem::take(&mut buffer);
422                            total += page.len();
423                            yield StreamPage { records: page, bookmark: None };
424                        } else if buffer.len() >= chunk {
425                            let page = std::mem::replace(
426                                &mut buffer,
427                                Vec::with_capacity(initial_capacity),
428                            );
429                            total += page.len();
430                            yield StreamPage { records: page, bookmark: None };
431                        }
432                    }
433                    S3FileFormat::JsonArray => {
434                        // JSON-array files cannot be parsed incrementally
435                        // (the closing `]` is required to validate the
436                        // structure), so each object is buffered fully and
437                        // then chunked. The caveat is documented in the
438                        // crate README.
439                        let text = self.read_object_text(key).await?;
440                        let value: Value = serde_json::from_str(&text).map_err(|e| {
441                            FaucetError::Source(format!("S3 JSON parse error in '{key}': {e}"))
442                        })?;
443                        let array = match value {
444                            Value::Array(arr) => arr,
445                            other => Err(FaucetError::Source(format!(
446                                "S3 expected JSON array in '{key}', got {}",
447                                value_type_name(&other)
448                            )))?,
449                        };
450                        if batch_size == 0 {
451                            // Flush any cross-object buffer first (none
452                            // here because each iteration completes its
453                            // own object — but keep symmetric with the
454                            // line-shaped branches).
455                            if !buffer.is_empty() {
456                                let page = std::mem::take(&mut buffer);
457                                total += page.len();
458                                yield StreamPage { records: page, bookmark: None };
459                            }
460                            total += array.len();
461                            yield StreamPage { records: array, bookmark: None };
462                        } else {
463                            for record in array {
464                                buffer.push(record);
465                                if buffer.len() >= chunk {
466                                    let page = std::mem::replace(
467                                        &mut buffer,
468                                        Vec::with_capacity(initial_capacity),
469                                    );
470                                    total += page.len();
471                                    yield StreamPage { records: page, bookmark: None };
472                                }
473                            }
474                        }
475                    }
476                }
477            }
478
479            if !buffer.is_empty() {
480                let page = std::mem::take(&mut buffer);
481                total += page.len();
482                yield StreamPage { records: page, bookmark: None };
483            }
484
485            tracing::info!(
486                total_records = total,
487                batch_size,
488                objects = keys.len(),
489                "S3 source stream complete",
490            );
491        })
492    }
493
494    fn config_schema(&self) -> serde_json::Value {
495        serde_json::to_value(faucet_core::schema_for!(S3SourceConfig))
496            .expect("schema serialization")
497    }
498
499    fn dataset_uri(&self) -> String {
500        match &self.config.prefix {
501            Some(p) => format!("s3://{}/{}", self.config.bucket, p),
502            None => format!("s3://{}", self.config.bucket),
503        }
504    }
505
506    /// The S3 source is always shardable: any object set can be split by
507    /// hash-of-key. Sharding only takes effect when the cluster coordinator
508    /// calls `apply_shard`; a plain `faucet run` reads
509    /// every object.
510    fn is_shardable(&self) -> bool {
511        true
512    }
513
514    /// Enumerate `target` hash-modulo shards. Each shard `i` will read the
515    /// objects whose key hashes to `i (mod target)`. No I/O: the partition is
516    /// defined by the hash function, so enumeration is cheap and stable as new
517    /// objects appear. `target <= 1` yields a single whole-dataset shard.
518    async fn enumerate_shards(&self, target: usize) -> Result<Vec<ShardSpec>, FaucetError> {
519        if target <= 1 {
520            return Ok(vec![ShardSpec::whole()]);
521        }
522        let shards = (0..target)
523            .map(|i| {
524                ShardSpec::new(
525                    i.to_string(),
526                    serde_json::json!({ "shards": target, "index": i }),
527                )
528            })
529            .collect();
530        Ok(shards)
531    }
532
533    /// Narrow this source to one hash-modulo shard. The whole-dataset shard
534    /// clears any filter (reads every object).
535    async fn apply_shard(&self, shard: &ShardSpec) -> Result<(), FaucetError> {
536        let parsed = if shard.is_whole() {
537            None
538        } else {
539            let shards = shard
540                .descriptor
541                .get("shards")
542                .and_then(Value::as_u64)
543                .ok_or_else(|| {
544                    FaucetError::Source(format!(
545                        "s3: invalid shard descriptor (missing 'shards'): {}",
546                        shard.descriptor
547                    ))
548                })?;
549            let index = shard
550                .descriptor
551                .get("index")
552                .and_then(Value::as_u64)
553                .ok_or_else(|| {
554                    FaucetError::Source(format!(
555                        "s3: invalid shard descriptor (missing 'index'): {}",
556                        shard.descriptor
557                    ))
558                })?;
559            Some((shards as usize, index as usize))
560        };
561        *self.applied_shard.lock().expect("shard mutex poisoned") = parsed;
562        Ok(())
563    }
564}
565
566/// Return a human-readable name for a JSON value type.
567fn value_type_name(v: &Value) -> &'static str {
568    match v {
569        Value::Null => "null",
570        Value::Bool(_) => "boolean",
571        Value::Number(_) => "number",
572        Value::String(_) => "string",
573        Value::Array(_) => "array",
574        Value::Object(_) => "object",
575    }
576}
577
578#[cfg(test)]
579mod tests {
580    use super::*;
581    use crate::config::S3SourceConfig;
582    use faucet_core::Source;
583    use serde_json::json;
584
585    /// Helper to build an S3Source synchronously for parse-only tests.
586    /// We construct it directly to avoid needing an async runtime for unit tests
587    /// that only exercise `parse_content`.
588    fn test_source(config: S3SourceConfig) -> S3Source {
589        // Build a dummy client — these tests never make network calls.
590        let sdk_config = aws_config::SdkConfig::builder()
591            .behavior_version(aws_config::BehaviorVersion::latest())
592            .build();
593        let client = Client::new(&sdk_config);
594        S3Source {
595            config,
596            client,
597            applied_shard: Mutex::new(None),
598        }
599    }
600
601    #[test]
602    fn parse_json_lines() {
603        let source = test_source(S3SourceConfig::new("test"));
604        let text = r#"{"id":1,"name":"Alice"}
605{"id":2,"name":"Bob"}
606"#;
607        let records = source.parse_content("test.jsonl", text).unwrap();
608        assert_eq!(records.len(), 2);
609        assert_eq!(records[0]["id"], 1);
610        assert_eq!(records[1]["name"], "Bob");
611    }
612
613    #[test]
614    fn parse_json_lines_skips_empty() {
615        let source = test_source(S3SourceConfig::new("test"));
616        let text = r#"{"id":1}
617
618{"id":2}
619
620"#;
621        let records = source.parse_content("test.jsonl", text).unwrap();
622        assert_eq!(records.len(), 2);
623    }
624
625    #[test]
626    fn parse_json_lines_invalid() {
627        let source = test_source(S3SourceConfig::new("test"));
628        let text = "not json\n";
629        let result = source.parse_content("test.jsonl", text);
630        assert!(result.is_err());
631        let err = result.unwrap_err().to_string();
632        assert!(err.contains("JSON parse error"));
633        assert!(err.contains("line 1"));
634    }
635
636    #[test]
637    fn parse_json_array() {
638        let source = test_source(S3SourceConfig::new("test").file_format(S3FileFormat::JsonArray));
639        let text = r#"[{"id":1},{"id":2}]"#;
640        let records = source.parse_content("test.json", text).unwrap();
641        assert_eq!(records.len(), 2);
642        assert_eq!(records[0]["id"], 1);
643    }
644
645    #[test]
646    fn parse_json_array_not_array() {
647        let source = test_source(S3SourceConfig::new("test").file_format(S3FileFormat::JsonArray));
648        let text = r#"{"id":1}"#;
649        let result = source.parse_content("test.json", text);
650        assert!(result.is_err());
651        let err = result.unwrap_err().to_string();
652        assert!(err.contains("expected JSON array"));
653    }
654
655    #[test]
656    fn parse_raw_text() {
657        let source = test_source(S3SourceConfig::new("test").file_format(S3FileFormat::RawText));
658        let text = "hello world\nline two";
659        let records = source.parse_content("data/file.txt", text).unwrap();
660        assert_eq!(records.len(), 1);
661        assert_eq!(
662            records[0],
663            json!({"key": "data/file.txt", "content": "hello world\nline two"})
664        );
665    }
666
667    #[cfg(feature = "compression")]
668    #[test]
669    fn compression_default_is_auto() {
670        let cfg = S3SourceConfig::new("bucket");
671        assert_eq!(cfg.compression, faucet_core::CompressionConfig::Auto);
672    }
673
674    // ── Hash-modulo sharding ────────────────────────────────────────────────
675
676    #[test]
677    fn shard_hash_is_deterministic() {
678        assert_eq!(
679            shard_hash("data/part-001.jsonl"),
680            shard_hash("data/part-001.jsonl")
681        );
682        assert_ne!(shard_hash("a"), shard_hash("b"));
683    }
684
685    #[tokio::test]
686    async fn enumerate_shards_returns_target_disjoint_shards() {
687        let source = test_source(S3SourceConfig::new("b"));
688        assert!(source.is_shardable());
689        let shards = source.enumerate_shards(3).await.unwrap();
690        assert_eq!(shards.len(), 3);
691        for (i, s) in shards.iter().enumerate() {
692            assert_eq!(s.descriptor["shards"], 3);
693            assert_eq!(s.descriptor["index"], i);
694        }
695    }
696
697    #[tokio::test]
698    async fn enumerate_shards_target_one_is_whole() {
699        let source = test_source(S3SourceConfig::new("b"));
700        let shards = source.enumerate_shards(1).await.unwrap();
701        assert_eq!(shards.len(), 1);
702        assert!(shards[0].is_whole());
703    }
704
705    // The union of every shard's filtered key set equals the full set, with no
706    // key in two shards — the core no-dup / no-loss guarantee.
707    #[tokio::test]
708    async fn shard_filter_partitions_keys_disjointly_and_completely() {
709        let keys: Vec<String> = (0..200).map(|i| format!("data/obj-{i}.jsonl")).collect();
710        let n = 4;
711        let mut union: Vec<String> = Vec::new();
712        for index in 0..n {
713            let source = test_source(S3SourceConfig::new("b"));
714            source
715                .apply_shard(&ShardSpec::new(
716                    index.to_string(),
717                    serde_json::json!({ "shards": n, "index": index }),
718                ))
719                .await
720                .unwrap();
721            let got = source.shard_filter(keys.clone());
722            union.extend(got);
723        }
724        union.sort();
725        let mut expected = keys.clone();
726        expected.sort();
727        assert_eq!(
728            union, expected,
729            "shards must union to the full key set, disjointly"
730        );
731    }
732
733    #[tokio::test]
734    async fn apply_whole_shard_reads_everything() {
735        let keys: Vec<String> = (0..20).map(|i| format!("k{i}")).collect();
736        let source = test_source(S3SourceConfig::new("b"));
737        source.apply_shard(&ShardSpec::whole()).await.unwrap();
738        assert_eq!(source.shard_filter(keys.clone()).len(), keys.len());
739    }
740
741    #[tokio::test]
742    async fn apply_shard_rejects_malformed_descriptor() {
743        let source = test_source(S3SourceConfig::new("b"));
744        let err = source
745            .apply_shard(&ShardSpec::new("0", serde_json::json!({ "index": 0 })))
746            .await
747            .unwrap_err();
748        assert!(matches!(err, FaucetError::Source(_)));
749    }
750
751    #[test]
752    fn dataset_uri_no_prefix() {
753        let source = test_source(S3SourceConfig::new("my-bucket"));
754        assert_eq!(source.dataset_uri(), "s3://my-bucket");
755    }
756
757    #[test]
758    fn dataset_uri_with_prefix() {
759        let source = test_source(S3SourceConfig::new("my-bucket").prefix("data/2026/"));
760        assert_eq!(source.dataset_uri(), "s3://my-bucket/data/2026/");
761    }
762}