faucet_source_gcs/
config.rs1use faucet_common_gcs::GcsCredentials;
4use faucet_core::DEFAULT_BATCH_SIZE;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
10#[serde(rename_all = "snake_case")]
11pub enum GcsFileFormat {
12 #[default]
14 JsonLines,
15 JsonArray,
17 RawText,
19 #[cfg(feature = "arrow")]
27 Parquet,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
32pub struct GcsSourceConfig {
33 pub bucket: String,
35 pub prefix: Option<String>,
37 pub object_keys: Option<Vec<String>>,
40 #[serde(default)]
42 pub auth: GcsCredentials,
43 #[serde(default)]
45 pub file_format: GcsFileFormat,
46 pub max_objects: Option<usize>,
48 #[serde(default = "default_concurrency")]
50 pub concurrency: usize,
51 #[serde(default = "default_batch_size")]
55 pub batch_size: usize,
56 #[serde(default = "default_true")]
64 pub verify_length: bool,
65 #[serde(default)]
72 pub verify_checksum: bool,
73 pub storage_host: Option<String>,
76 #[cfg(feature = "compression")]
82 #[serde(default)]
83 pub compression: faucet_core::CompressionConfig,
84}
85
86fn default_batch_size() -> usize {
87 DEFAULT_BATCH_SIZE
88}
89fn default_concurrency() -> usize {
90 10
91}
92fn default_true() -> bool {
93 true
94}
95
96impl GcsSourceConfig {
97 pub fn new(bucket: impl Into<String>) -> Self {
99 Self {
100 bucket: bucket.into(),
101 prefix: None,
102 object_keys: None,
103 auth: GcsCredentials::default(),
104 file_format: GcsFileFormat::default(),
105 max_objects: None,
106 concurrency: default_concurrency(),
107 batch_size: default_batch_size(),
108 verify_length: true,
109 verify_checksum: false,
110 storage_host: None,
111 #[cfg(feature = "compression")]
112 compression: faucet_core::CompressionConfig::default(),
113 }
114 }
115
116 pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
117 self.prefix = Some(prefix.into());
118 self
119 }
120
121 pub fn object_keys(mut self, keys: Vec<String>) -> Self {
122 self.object_keys = Some(keys);
123 self
124 }
125
126 pub fn auth(mut self, creds: GcsCredentials) -> Self {
127 self.auth = creds;
128 self
129 }
130
131 pub fn file_format(mut self, format: GcsFileFormat) -> Self {
132 self.file_format = format;
133 self
134 }
135
136 pub fn max_objects(mut self, max: usize) -> Self {
137 self.max_objects = Some(max);
138 self
139 }
140
141 pub fn concurrency(mut self, concurrency: usize) -> Self {
142 self.concurrency = concurrency;
143 self
144 }
145
146 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
147 self.batch_size = batch_size;
148 self
149 }
150
151 pub fn storage_host(mut self, host: impl Into<String>) -> Self {
152 self.storage_host = Some(host.into());
153 self
154 }
155
156 pub fn verify_length(mut self, verify: bool) -> Self {
159 self.verify_length = verify;
160 self
161 }
162
163 pub fn verify_checksum(mut self, verify: bool) -> Self {
166 self.verify_checksum = verify;
167 self
168 }
169
170 #[cfg(feature = "compression")]
172 pub fn compression(mut self, c: faucet_core::CompressionConfig) -> Self {
173 self.compression = c;
174 self
175 }
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181
182 #[test]
183 fn default_config() {
184 let config = GcsSourceConfig::new("my-bucket");
185 assert_eq!(config.bucket, "my-bucket");
186 assert!(config.prefix.is_none());
187 assert!(config.object_keys.is_none());
188 assert!(matches!(config.auth, GcsCredentials::ApplicationDefault));
189 assert!(matches!(config.file_format, GcsFileFormat::JsonLines));
190 assert!(config.max_objects.is_none());
191 assert_eq!(config.concurrency, 10);
192 assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
193 assert!(config.storage_host.is_none());
194 }
195
196 #[test]
197 fn builder_methods() {
198 let config = GcsSourceConfig::new("my-bucket")
199 .prefix("data/")
200 .file_format(GcsFileFormat::JsonArray)
201 .max_objects(5)
202 .concurrency(20)
203 .with_batch_size(250)
204 .storage_host("http://localhost:4443");
205
206 assert_eq!(config.bucket, "my-bucket");
207 assert_eq!(config.prefix.as_deref(), Some("data/"));
208 assert!(matches!(config.file_format, GcsFileFormat::JsonArray));
209 assert_eq!(config.max_objects, Some(5));
210 assert_eq!(config.concurrency, 20);
211 assert_eq!(config.batch_size, 250);
212 assert_eq!(
213 config.storage_host.as_deref(),
214 Some("http://localhost:4443")
215 );
216 }
217
218 #[test]
219 fn file_format_default_is_json_lines() {
220 assert!(matches!(GcsFileFormat::default(), GcsFileFormat::JsonLines));
221 }
222
223 #[test]
224 fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
225 let config = GcsSourceConfig::new("b").with_batch_size(0);
226 assert_eq!(config.batch_size, 0);
227 assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
228 }
229
230 #[test]
231 fn batch_size_above_max_is_rejected() {
232 let config = GcsSourceConfig::new("b").with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
233 assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
234 }
235
236 #[cfg(feature = "compression")]
237 #[test]
238 fn compression_default_is_auto() {
239 let cfg = GcsSourceConfig::new("bucket");
240 assert_eq!(cfg.compression, faucet_core::CompressionConfig::Auto);
241 }
242
243 #[test]
244 fn verify_defaults_length_on_checksum_off() {
245 let cfg = GcsSourceConfig::new("b");
246 assert!(cfg.verify_length);
247 assert!(!cfg.verify_checksum);
248 }
249
250 #[test]
251 fn verify_builders_override() {
252 let cfg = GcsSourceConfig::new("b")
253 .verify_length(false)
254 .verify_checksum(true);
255 assert!(!cfg.verify_length);
256 assert!(cfg.verify_checksum);
257 }
258
259 #[test]
260 fn verify_fields_default_when_absent_from_json() {
261 let json = r#"{
262 "bucket": "my-bucket",
263 "prefix": null,
264 "object_keys": null,
265 "file_format": "json_lines",
266 "max_objects": null,
267 "concurrency": 10,
268 "storage_host": null
269 }"#;
270 let config: GcsSourceConfig = serde_json::from_str(json).unwrap();
271 assert!(config.verify_length);
272 assert!(!config.verify_checksum);
273 }
274
275 #[test]
276 fn batch_size_defaults_when_omitted_from_json() {
277 let json = r#"{
278 "bucket": "my-bucket",
279 "prefix": null,
280 "object_keys": null,
281 "file_format": "json_lines",
282 "max_objects": null,
283 "concurrency": 10,
284 "storage_host": null
285 }"#;
286 let config: GcsSourceConfig = serde_json::from_str(json).unwrap();
287 assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
288 }
289}