Skip to main content

faucet_source_s3/
config.rs

1//! S3 source configuration.
2
3use faucet_core::DEFAULT_BATCH_SIZE;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// Format of files stored in S3.
8#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
9#[serde(rename_all = "snake_case")]
10pub enum S3FileFormat {
11    /// Each line in the file is a separate JSON record.
12    #[default]
13    JsonLines,
14    /// The entire file is a JSON array of records.
15    JsonArray,
16    /// Each file becomes a single record with `"key"` and `"content"` fields.
17    RawText,
18    /// Apache Parquet objects. Each object is decoded via the Arrow Parquet
19    /// reader; its `RecordBatch`es feed both the row path (converted to JSON
20    /// records) and the **columnar** fast path
21    /// ([`Source::stream_batches`](faucet_core::Source::stream_batches)) so an
22    /// `s3(parquet) → parquet`/`delta` chain never materializes
23    /// `serde_json::Value`. Requires the crate-local `arrow` feature
24    /// (RFC 0002 / #375).
25    #[cfg(feature = "arrow")]
26    Parquet,
27}
28
29/// Configuration for the S3 source connector.
30#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
31pub struct S3SourceConfig {
32    /// S3 bucket name.
33    pub bucket: String,
34    /// Object key prefix filter.
35    pub prefix: Option<String>,
36    /// AWS region. `None` uses the SDK default.
37    pub region: Option<String>,
38    /// Custom endpoint URL for S3-compatible services (e.g. MinIO).
39    pub endpoint_url: Option<String>,
40    /// Format of the files to read.
41    pub file_format: S3FileFormat,
42    /// Maximum number of objects to read.
43    pub max_objects: Option<usize>,
44    /// Maximum number of concurrent object reads (default: 10).
45    pub concurrency: usize,
46    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). For
47    /// `JsonLines` and `RawText` formats, the object body is decoded
48    /// line-by-line via [`tokio::io::AsyncBufReadExt`] and a page is yielded
49    /// whenever the buffer reaches this size; multi-object scans flatten so
50    /// a single page may contain lines from any object. For `JsonArray`,
51    /// each object is buffered fully before its records are chunked into
52    /// pages of this size (see the README "Streaming and batching" section
53    /// for the caveat). Defaults to [`DEFAULT_BATCH_SIZE`].
54    ///
55    /// `batch_size = 0` is the "no batching" sentinel: every page is one
56    /// complete object — no within-object chunking. Useful for small
57    /// lookup files, or for sinks (e.g. SQL `COPY`, BigQuery load jobs)
58    /// that prefer one large request per file to many small ones.
59    #[serde(default = "default_batch_size")]
60    pub batch_size: usize,
61    /// Verify each object's byte length against the `Content-Length` the
62    /// store advertises, failing the read with [`FaucetError::Source`](faucet_core::FaucetError::Source)
63    /// on a short (truncated) or over-long transfer (#161). The check is
64    /// cheap (a byte counter over the body that is read anyway) and defaults
65    /// to `true`. Disable it only for an S3-compatible store that does not
66    /// return a reliable `Content-Length`. When the store reports no length,
67    /// the check is skipped (a debug log notes it) rather than failing.
68    #[serde(default = "default_true")]
69    pub verify_length: bool,
70    /// Verify each object's body against the checksum the store advertises —
71    /// an `x-amz-checksum-{crc32,crc32c,sha1,sha256}` header when present, or
72    /// the ETag as MD5 for a non-multipart upload (#161). Stronger than the
73    /// length check but costs a hash over the full body, so it defaults to
74    /// `false`. Enabling it sets `ChecksumMode::Enabled` on each `GetObject`
75    /// so the store returns its stored checksum. When the store advertises no
76    /// usable checksum for an object, verification is skipped for that object
77    /// (a debug log notes it); the length check still applies.
78    #[serde(default)]
79    pub verify_checksum: bool,
80    /// Compression codec applied to each downloaded object. Defaults to
81    /// [`CompressionConfig::Auto`](faucet_core::CompressionConfig::Auto) —
82    /// the codec is resolved per-object-key, so a single source can read a
83    /// mix of compressed and uncompressed objects. Requires the
84    /// crate-local `compression` feature.
85    #[cfg(feature = "compression")]
86    #[serde(default)]
87    pub compression: faucet_core::CompressionConfig,
88}
89
90fn default_batch_size() -> usize {
91    DEFAULT_BATCH_SIZE
92}
93
94fn default_true() -> bool {
95    true
96}
97
98impl S3SourceConfig {
99    /// Create a new config with the required bucket name and sensible defaults.
100    pub fn new(bucket: impl Into<String>) -> Self {
101        Self {
102            bucket: bucket.into(),
103            prefix: None,
104            region: None,
105            endpoint_url: None,
106            file_format: S3FileFormat::default(),
107            max_objects: None,
108            concurrency: 10,
109            batch_size: DEFAULT_BATCH_SIZE,
110            verify_length: true,
111            verify_checksum: false,
112            #[cfg(feature = "compression")]
113            compression: faucet_core::CompressionConfig::default(),
114        }
115    }
116
117    /// Set the object key prefix filter.
118    pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
119        self.prefix = Some(prefix.into());
120        self
121    }
122
123    /// Set the AWS region.
124    pub fn region(mut self, region: impl Into<String>) -> Self {
125        self.region = Some(region.into());
126        self
127    }
128
129    /// Set a custom endpoint URL for S3-compatible services.
130    pub fn endpoint_url(mut self, url: impl Into<String>) -> Self {
131        self.endpoint_url = Some(url.into());
132        self
133    }
134
135    /// Set the file format.
136    pub fn file_format(mut self, format: S3FileFormat) -> Self {
137        self.file_format = format;
138        self
139    }
140
141    /// Set the maximum number of objects to read.
142    pub fn max_objects(mut self, max: usize) -> Self {
143        self.max_objects = Some(max);
144        self
145    }
146
147    /// Set the maximum number of concurrent object reads.
148    pub fn concurrency(mut self, concurrency: usize) -> Self {
149        self.concurrency = concurrency;
150        self
151    }
152
153    /// Set the per-page record count for [`Source::stream_pages`](faucet_core::Source::stream_pages).
154    ///
155    /// Pass `0` to opt out of within-object chunking — every emitted
156    /// [`StreamPage`](faucet_core::StreamPage) corresponds to exactly one
157    /// S3 object.
158    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
159        self.batch_size = batch_size;
160        self
161    }
162
163    /// Enable or disable the per-object `Content-Length` verification
164    /// (default `true`). See [`verify_length`](Self::verify_length).
165    pub fn verify_length(mut self, verify: bool) -> Self {
166        self.verify_length = verify;
167        self
168    }
169
170    /// Enable or disable per-object checksum verification (default `false`).
171    /// See [`verify_checksum`](Self::verify_checksum).
172    pub fn verify_checksum(mut self, verify: bool) -> Self {
173        self.verify_checksum = verify;
174        self
175    }
176
177    /// Set the compression codec. Available only with the `compression` feature.
178    #[cfg(feature = "compression")]
179    pub fn compression(mut self, c: faucet_core::CompressionConfig) -> Self {
180        self.compression = c;
181        self
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn default_config() {
191        let config = S3SourceConfig::new("my-bucket");
192        assert_eq!(config.bucket, "my-bucket");
193        assert!(config.prefix.is_none());
194        assert!(config.region.is_none());
195        assert!(config.endpoint_url.is_none());
196        assert!(matches!(config.file_format, S3FileFormat::JsonLines));
197        assert!(config.max_objects.is_none());
198    }
199
200    #[test]
201    fn builder_methods() {
202        let config = S3SourceConfig::new("my-bucket")
203            .prefix("data/")
204            .region("us-west-2")
205            .endpoint_url("http://localhost:9000")
206            .file_format(S3FileFormat::JsonArray)
207            .max_objects(10);
208
209        assert_eq!(config.bucket, "my-bucket");
210        assert_eq!(config.prefix.as_deref(), Some("data/"));
211        assert_eq!(config.region.as_deref(), Some("us-west-2"));
212        assert_eq!(
213            config.endpoint_url.as_deref(),
214            Some("http://localhost:9000")
215        );
216        assert!(matches!(config.file_format, S3FileFormat::JsonArray));
217        assert_eq!(config.max_objects, Some(10));
218    }
219
220    #[test]
221    fn file_format_default_is_json_lines() {
222        let format = S3FileFormat::default();
223        assert!(matches!(format, S3FileFormat::JsonLines));
224    }
225
226    #[test]
227    fn batch_size_defaults_to_default_batch_size() {
228        let config = S3SourceConfig::new("my-bucket");
229        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
230    }
231
232    #[test]
233    fn with_batch_size_overrides_default() {
234        let config = S3SourceConfig::new("my-bucket").with_batch_size(500);
235        assert_eq!(config.batch_size, 500);
236    }
237
238    #[test]
239    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
240        let config = S3SourceConfig::new("my-bucket").with_batch_size(0);
241        assert_eq!(config.batch_size, 0);
242        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
243    }
244
245    #[test]
246    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
247        let config =
248            S3SourceConfig::new("my-bucket").with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
249        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
250    }
251
252    #[test]
253    fn batch_size_deserializes_from_json() {
254        let json = r#"{
255            "bucket": "my-bucket",
256            "prefix": null,
257            "region": null,
258            "endpoint_url": null,
259            "file_format": "json_lines",
260            "max_objects": null,
261            "concurrency": 10,
262            "batch_size": 250
263        }"#;
264        let config: S3SourceConfig = serde_json::from_str(json).unwrap();
265        assert_eq!(config.batch_size, 250);
266    }
267
268    #[test]
269    fn verify_defaults_length_on_checksum_off() {
270        let cfg = S3SourceConfig::new("b");
271        assert!(cfg.verify_length, "length verification defaults on");
272        assert!(!cfg.verify_checksum, "checksum verification defaults off");
273    }
274
275    #[test]
276    fn verify_fields_default_when_absent_from_json() {
277        // An existing config that predates these fields must still parse, with
278        // length verification on and checksum off.
279        let json = r#"{
280            "bucket": "my-bucket",
281            "prefix": null,
282            "region": null,
283            "endpoint_url": null,
284            "file_format": "json_lines",
285            "max_objects": null,
286            "concurrency": 10,
287            "batch_size": 250
288        }"#;
289        let config: S3SourceConfig = serde_json::from_str(json).unwrap();
290        assert!(config.verify_length);
291        assert!(!config.verify_checksum);
292    }
293
294    #[test]
295    fn verify_builders_override() {
296        let cfg = S3SourceConfig::new("b")
297            .verify_length(false)
298            .verify_checksum(true);
299        assert!(!cfg.verify_length);
300        assert!(cfg.verify_checksum);
301    }
302
303    #[cfg(feature = "compression")]
304    #[test]
305    fn compression_default_is_auto() {
306        let cfg = S3SourceConfig::new("bucket");
307        assert_eq!(cfg.compression, faucet_core::CompressionConfig::Auto);
308    }
309}