1use 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
14pub struct S3Source {
16 config: S3SourceConfig,
17 client: Client,
18 applied_shard: Mutex<Option<(usize, usize)>>,
22}
23
24fn 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 fn is_shardable(&self) -> bool {
511 true
512 }
513
514 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 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
566fn 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 fn test_source(config: S3SourceConfig) -> S3Source {
589 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 #[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 #[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}