Skip to main content

faucet_source_gcs/
config.rs

1//! GCS source configuration.
2
3use faucet_common_gcs::GcsCredentials;
4use faucet_core::DEFAULT_BATCH_SIZE;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8/// Format of files stored in GCS.
9#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
10#[serde(rename_all = "snake_case")]
11pub enum GcsFileFormat {
12    /// Each line in the file is a separate JSON record.
13    #[default]
14    JsonLines,
15    /// The entire file is a JSON array of records.
16    JsonArray,
17    /// Each file becomes a single record with `"key"` and `"content"` fields.
18    RawText,
19}
20
21/// Configuration for the GCS source connector.
22#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
23pub struct GcsSourceConfig {
24    /// GCS bucket name.
25    pub bucket: String,
26    /// Object name prefix filter. Ignored when `object_keys` is set.
27    pub prefix: Option<String>,
28    /// Explicit object names. When set, listing is skipped and `prefix`
29    /// is ignored.
30    pub object_keys: Option<Vec<String>>,
31    /// Credential source.
32    #[serde(default)]
33    pub auth: GcsCredentials,
34    /// File format.
35    #[serde(default)]
36    pub file_format: GcsFileFormat,
37    /// Hard cap on the number of objects read (after listing).
38    pub max_objects: Option<usize>,
39    /// Maximum concurrent object reads (default: 10).
40    #[serde(default = "default_concurrency")]
41    pub concurrency: usize,
42    /// Records per emitted `StreamPage`. See "Streaming and batching"
43    /// in the README. `batch_size = 0` is the "no batching" sentinel and
44    /// emits one page per object.
45    #[serde(default = "default_batch_size")]
46    pub batch_size: usize,
47    /// Verify each object's byte length against the `size` GCS reports for it,
48    /// failing the read with [`FaucetError::Source`](faucet_core::FaucetError::Source)
49    /// on a short (truncated) or over-long transfer (#161). Cheap (a byte
50    /// counter over the body that is read anyway) and defaults to `true`.
51    /// The check is automatically skipped for an object served with a
52    /// non-empty `Content-Encoding` (GCS may decompressively transcode it on
53    /// read, so the received byte count would not match the stored `size`).
54    #[serde(default = "default_true")]
55    pub verify_length: bool,
56    /// Verify each object's body against the CRC32C (or MD5) checksum GCS
57    /// reports for it (#161). Stronger than the length check but costs a hash
58    /// over the full body, so it defaults to `false`. Skipped for an object
59    /// with no usable checksum or one served with a non-empty
60    /// `Content-Encoding` (the stored checksum covers the stored bytes, which
61    /// transcoding would not return).
62    #[serde(default)]
63    pub verify_checksum: bool,
64    /// Optional storage-host override (e.g. `http://localhost:4443` for
65    /// fake-gcs-server). Production users should leave this unset.
66    pub storage_host: Option<String>,
67    /// Compression codec applied to each downloaded object. Defaults to
68    /// [`CompressionConfig::Auto`](faucet_core::CompressionConfig::Auto) —
69    /// the codec is resolved per-object-key, so a single source can read a
70    /// mix of compressed and uncompressed objects. Requires the
71    /// crate-local `compression` feature.
72    #[cfg(feature = "compression")]
73    #[serde(default)]
74    pub compression: faucet_core::CompressionConfig,
75}
76
77fn default_batch_size() -> usize {
78    DEFAULT_BATCH_SIZE
79}
80fn default_concurrency() -> usize {
81    10
82}
83fn default_true() -> bool {
84    true
85}
86
87impl GcsSourceConfig {
88    /// Create a new config with the required bucket name and sensible defaults.
89    pub fn new(bucket: impl Into<String>) -> Self {
90        Self {
91            bucket: bucket.into(),
92            prefix: None,
93            object_keys: None,
94            auth: GcsCredentials::default(),
95            file_format: GcsFileFormat::default(),
96            max_objects: None,
97            concurrency: default_concurrency(),
98            batch_size: default_batch_size(),
99            verify_length: true,
100            verify_checksum: false,
101            storage_host: None,
102            #[cfg(feature = "compression")]
103            compression: faucet_core::CompressionConfig::default(),
104        }
105    }
106
107    pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
108        self.prefix = Some(prefix.into());
109        self
110    }
111
112    pub fn object_keys(mut self, keys: Vec<String>) -> Self {
113        self.object_keys = Some(keys);
114        self
115    }
116
117    pub fn auth(mut self, creds: GcsCredentials) -> Self {
118        self.auth = creds;
119        self
120    }
121
122    pub fn file_format(mut self, format: GcsFileFormat) -> Self {
123        self.file_format = format;
124        self
125    }
126
127    pub fn max_objects(mut self, max: usize) -> Self {
128        self.max_objects = Some(max);
129        self
130    }
131
132    pub fn concurrency(mut self, concurrency: usize) -> Self {
133        self.concurrency = concurrency;
134        self
135    }
136
137    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
138        self.batch_size = batch_size;
139        self
140    }
141
142    pub fn storage_host(mut self, host: impl Into<String>) -> Self {
143        self.storage_host = Some(host.into());
144        self
145    }
146
147    /// Enable or disable the per-object length verification (default `true`).
148    /// See [`verify_length`](Self::verify_length).
149    pub fn verify_length(mut self, verify: bool) -> Self {
150        self.verify_length = verify;
151        self
152    }
153
154    /// Enable or disable per-object checksum verification (default `false`).
155    /// See [`verify_checksum`](Self::verify_checksum).
156    pub fn verify_checksum(mut self, verify: bool) -> Self {
157        self.verify_checksum = verify;
158        self
159    }
160
161    /// Set the compression codec. Available only with the `compression` feature.
162    #[cfg(feature = "compression")]
163    pub fn compression(mut self, c: faucet_core::CompressionConfig) -> Self {
164        self.compression = c;
165        self
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn default_config() {
175        let config = GcsSourceConfig::new("my-bucket");
176        assert_eq!(config.bucket, "my-bucket");
177        assert!(config.prefix.is_none());
178        assert!(config.object_keys.is_none());
179        assert!(matches!(config.auth, GcsCredentials::ApplicationDefault));
180        assert!(matches!(config.file_format, GcsFileFormat::JsonLines));
181        assert!(config.max_objects.is_none());
182        assert_eq!(config.concurrency, 10);
183        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
184        assert!(config.storage_host.is_none());
185    }
186
187    #[test]
188    fn builder_methods() {
189        let config = GcsSourceConfig::new("my-bucket")
190            .prefix("data/")
191            .file_format(GcsFileFormat::JsonArray)
192            .max_objects(5)
193            .concurrency(20)
194            .with_batch_size(250)
195            .storage_host("http://localhost:4443");
196
197        assert_eq!(config.bucket, "my-bucket");
198        assert_eq!(config.prefix.as_deref(), Some("data/"));
199        assert!(matches!(config.file_format, GcsFileFormat::JsonArray));
200        assert_eq!(config.max_objects, Some(5));
201        assert_eq!(config.concurrency, 20);
202        assert_eq!(config.batch_size, 250);
203        assert_eq!(
204            config.storage_host.as_deref(),
205            Some("http://localhost:4443")
206        );
207    }
208
209    #[test]
210    fn file_format_default_is_json_lines() {
211        assert!(matches!(GcsFileFormat::default(), GcsFileFormat::JsonLines));
212    }
213
214    #[test]
215    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
216        let config = GcsSourceConfig::new("b").with_batch_size(0);
217        assert_eq!(config.batch_size, 0);
218        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
219    }
220
221    #[test]
222    fn batch_size_above_max_is_rejected() {
223        let config = GcsSourceConfig::new("b").with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
224        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
225    }
226
227    #[cfg(feature = "compression")]
228    #[test]
229    fn compression_default_is_auto() {
230        let cfg = GcsSourceConfig::new("bucket");
231        assert_eq!(cfg.compression, faucet_core::CompressionConfig::Auto);
232    }
233
234    #[test]
235    fn verify_defaults_length_on_checksum_off() {
236        let cfg = GcsSourceConfig::new("b");
237        assert!(cfg.verify_length);
238        assert!(!cfg.verify_checksum);
239    }
240
241    #[test]
242    fn verify_builders_override() {
243        let cfg = GcsSourceConfig::new("b")
244            .verify_length(false)
245            .verify_checksum(true);
246        assert!(!cfg.verify_length);
247        assert!(cfg.verify_checksum);
248    }
249
250    #[test]
251    fn verify_fields_default_when_absent_from_json() {
252        let json = r#"{
253            "bucket": "my-bucket",
254            "prefix": null,
255            "object_keys": null,
256            "file_format": "json_lines",
257            "max_objects": null,
258            "concurrency": 10,
259            "storage_host": null
260        }"#;
261        let config: GcsSourceConfig = serde_json::from_str(json).unwrap();
262        assert!(config.verify_length);
263        assert!(!config.verify_checksum);
264    }
265
266    #[test]
267    fn batch_size_defaults_when_omitted_from_json() {
268        let json = r#"{
269            "bucket": "my-bucket",
270            "prefix": null,
271            "object_keys": null,
272            "file_format": "json_lines",
273            "max_objects": null,
274            "concurrency": 10,
275            "storage_host": null
276        }"#;
277        let config: GcsSourceConfig = serde_json::from_str(json).unwrap();
278        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
279    }
280}