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