faucet_source_bigquery/
config.rs

1//! BigQuery source configuration.
2
3use faucet_core::DEFAULT_BATCH_SIZE;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::time::Duration;
8
9// Re-export the shared credentials type so end-user imports remain stable.
10pub use faucet_common_bigquery::BigQueryCredentials;
11
12fn default_use_legacy_sql() -> bool {
13    false
14}
15
16fn default_max_results_per_page() -> i32 {
17    1000
18}
19
20fn default_statement_timeout() -> Duration {
21    Duration::from_secs(60)
22}
23
24fn default_poll_timeout() -> Duration {
25    Duration::from_secs(300)
26}
27
28fn default_batch_size() -> usize {
29    DEFAULT_BATCH_SIZE
30}
31
32/// Configuration for the BigQuery query source.
33#[derive(Clone, Serialize, Deserialize, JsonSchema)]
34pub struct BigQuerySourceConfig {
35    /// GCP project ID against which the query is billed and run.
36    pub project_id: String,
37    /// Authentication — the `auth` field, consistent with every other connector.
38    pub auth: BigQueryCredentials,
39    /// SQL statement to execute. May contain `${field.path}` placeholders that
40    /// are resolved against the parent-record context at runtime as
41    /// positional `?` markers; matched values are appended to
42    /// [`params`](Self::params) when the query is sent.
43    pub query: String,
44    /// Whether to use BigQuery's legacy SQL dialect. Defaults to `false`
45    /// (Standard SQL). Set to `true` only for tables that use legacy
46    /// `[project:dataset.table]` references.
47    #[serde(default = "default_use_legacy_sql")]
48    pub use_legacy_sql: bool,
49    /// Optional location override for non-`US` jobs (`"EU"`, `"asia-east1"`).
50    /// When `None`, BigQuery uses the default location for the queried tables.
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub location: Option<String>,
53    /// Maximum rows per page when calling `jobs.getQueryResults`. Smaller
54    /// values trade more HTTP round-trips for lower memory; larger values
55    /// trade memory for fewer requests. Defaults to 1000.
56    #[serde(default = "default_max_results_per_page")]
57    pub max_results_per_page: i32,
58    /// Positional bind parameters for the query, sent as
59    /// [`POSITIONAL`](https://cloud.google.com/bigquery/docs/parameterized-queries)
60    /// query parameters in declaration order before any context-derived
61    /// values. Each value is shipped as a STRING parameter; BigQuery casts
62    /// as needed at execution time.
63    #[serde(default)]
64    pub params: Vec<Value>,
65    /// Per-statement server-side timeout. Forwarded to the
66    /// `timeoutMs` field on `jobs.query`. Defaults to 60 seconds. If
67    /// BigQuery does not finish the query within this window it responds
68    /// with `jobComplete=false`; the source then polls
69    /// `jobs.getQueryResults` until the job completes.
70    #[serde(
71        default = "default_statement_timeout",
72        with = "faucet_core::config::duration_secs"
73    )]
74    #[schemars(with = "u64")]
75    pub statement_timeout: Duration,
76    /// Maximum wall-clock time the source will spend polling
77    /// `jobs.getQueryResults` for a job that keeps reporting
78    /// `jobComplete=false`, before giving up with
79    /// [`FaucetError::Source`](faucet_core::FaucetError::Source).
80    /// Without this cap a job that never completes would loop forever.
81    /// Defaults to 300 seconds. Set to `0` to disable the cap and poll
82    /// indefinitely. Only the *completion* wait is bounded; once the job is
83    /// complete, ordinary `pageToken` paging is unaffected.
84    #[serde(
85        default = "default_poll_timeout",
86        with = "faucet_core::config::duration_secs"
87    )]
88    #[schemars(with = "u64")]
89    pub poll_timeout: Duration,
90    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). Rows
91    /// returned by BigQuery are re-framed into pages of this size — every
92    /// time the buffer reaches `batch_size`, a page is yielded. Defaults to
93    /// [`DEFAULT_BATCH_SIZE`].
94    ///
95    /// `batch_size = 0` is the **"no batching" sentinel**: the entire
96    /// result set is buffered and emitted in a single page. Useful for
97    /// small lookup tables, or for sinks that prefer one large request to
98    /// many small ones.
99    #[serde(default = "default_batch_size")]
100    pub batch_size: usize,
101}
102
103impl std::fmt::Debug for BigQuerySourceConfig {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        f.debug_struct("BigQuerySourceConfig")
106            .field("project_id", &self.project_id)
107            .field("auth", &self.auth)
108            .field("query", &self.query)
109            .field("use_legacy_sql", &self.use_legacy_sql)
110            .field("location", &self.location)
111            .field("max_results_per_page", &self.max_results_per_page)
112            .field("params", &self.params)
113            .field("statement_timeout", &self.statement_timeout)
114            .field("poll_timeout", &self.poll_timeout)
115            .field("batch_size", &self.batch_size)
116            .finish()
117    }
118}
119
120impl BigQuerySourceConfig {
121    /// Create a new config with required fields and sensible defaults.
122    pub fn new(
123        project_id: impl Into<String>,
124        credentials: BigQueryCredentials,
125        query: impl Into<String>,
126    ) -> Self {
127        Self {
128            project_id: project_id.into(),
129            auth: credentials,
130            query: query.into(),
131            use_legacy_sql: default_use_legacy_sql(),
132            location: None,
133            max_results_per_page: default_max_results_per_page(),
134            params: Vec::new(),
135            statement_timeout: default_statement_timeout(),
136            poll_timeout: default_poll_timeout(),
137            batch_size: DEFAULT_BATCH_SIZE,
138        }
139    }
140
141    /// Enable BigQuery's legacy SQL dialect.
142    pub fn with_use_legacy_sql(mut self, use_legacy: bool) -> Self {
143        self.use_legacy_sql = use_legacy;
144        self
145    }
146
147    /// Pin the job to a specific location (e.g. `"EU"`).
148    pub fn with_location(mut self, location: impl Into<String>) -> Self {
149        self.location = Some(location.into());
150        self
151    }
152
153    /// Set the maximum row count per `getQueryResults` page.
154    pub fn with_max_results_per_page(mut self, max_results: i32) -> Self {
155        self.max_results_per_page = max_results;
156        self
157    }
158
159    /// Set the positional bind parameters for the query.
160    pub fn with_params(mut self, params: Vec<Value>) -> Self {
161        self.params = params;
162        self
163    }
164
165    /// Set the per-statement server-side timeout.
166    pub fn with_statement_timeout(mut self, timeout: Duration) -> Self {
167        self.statement_timeout = timeout;
168        self
169    }
170
171    /// Set the maximum wall-clock time spent polling for job completion
172    /// before giving up. Pass `Duration::ZERO` to poll forever.
173    pub fn with_poll_timeout(mut self, timeout: Duration) -> Self {
174        self.poll_timeout = timeout;
175        self
176    }
177
178    /// Set the records-per-page hint for [`Source::stream_pages`](faucet_core::Source::stream_pages).
179    ///
180    /// Pass `0` to opt out of batching — the entire result set is emitted
181    /// in a single page.
182    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
183        self.batch_size = batch_size;
184        self
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use serde_json::json;
192
193    fn sample() -> BigQuerySourceConfig {
194        BigQuerySourceConfig::new(
195            "my-project",
196            BigQueryCredentials::ApplicationDefault,
197            "SELECT id FROM events",
198        )
199    }
200
201    #[test]
202    fn default_config() {
203        let c = sample();
204        assert_eq!(c.project_id, "my-project");
205        assert!(!c.use_legacy_sql);
206        assert!(c.location.is_none());
207        assert_eq!(c.max_results_per_page, 1000);
208        assert!(c.params.is_empty());
209        assert_eq!(c.statement_timeout, Duration::from_secs(60));
210        assert_eq!(c.poll_timeout, Duration::from_secs(300));
211        assert_eq!(c.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
212    }
213
214    #[test]
215    fn builder_chaining() {
216        let c = sample()
217            .with_use_legacy_sql(true)
218            .with_location("EU")
219            .with_max_results_per_page(500)
220            .with_params(vec![json!("us-east")])
221            .with_statement_timeout(Duration::from_secs(30))
222            .with_batch_size(250);
223        assert!(c.use_legacy_sql);
224        assert_eq!(c.location.as_deref(), Some("EU"));
225        assert_eq!(c.max_results_per_page, 500);
226        assert_eq!(c.params, vec![json!("us-east")]);
227        assert_eq!(c.statement_timeout, Duration::from_secs(30));
228        assert_eq!(c.batch_size, 250);
229    }
230
231    #[test]
232    fn deserializes_minimal_json() {
233        let json = r#"{
234            "project_id": "my-project",
235            "auth": {"type": "application_default"},
236            "query": "SELECT 1"
237        }"#;
238        let c: BigQuerySourceConfig = serde_json::from_str(json).unwrap();
239        assert!(!c.use_legacy_sql);
240        assert_eq!(c.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
241        assert_eq!(c.statement_timeout, Duration::from_secs(60));
242        assert_eq!(c.max_results_per_page, 1000);
243    }
244
245    #[test]
246    fn deserializes_all_fields() {
247        let json = r#"{
248            "project_id": "p",
249            "auth": {"type": "application_default"},
250            "query": "SELECT 1",
251            "use_legacy_sql": true,
252            "location": "EU",
253            "max_results_per_page": 500,
254            "params": ["us-east"],
255            "statement_timeout": 30,
256            "batch_size": 250
257        }"#;
258        let c: BigQuerySourceConfig = serde_json::from_str(json).unwrap();
259        assert!(c.use_legacy_sql);
260        assert_eq!(c.location.as_deref(), Some("EU"));
261        assert_eq!(c.max_results_per_page, 500);
262        assert_eq!(c.statement_timeout, Duration::from_secs(30));
263        assert_eq!(c.batch_size, 250);
264    }
265
266    #[test]
267    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
268        let c = sample().with_batch_size(0);
269        assert!(faucet_core::validate_batch_size(c.batch_size).is_ok());
270    }
271
272    #[test]
273    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
274        let c = sample().with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
275        assert!(faucet_core::validate_batch_size(c.batch_size).is_err());
276    }
277
278    #[test]
279    fn debug_masks_inline_credentials() {
280        let c = BigQuerySourceConfig::new(
281            "p",
282            BigQueryCredentials::ServiceAccountKey {
283                json: "secret".into(),
284            },
285            "SELECT 1",
286        );
287        let dbg = format!("{c:?}");
288        assert!(!dbg.contains("secret"));
289        assert!(dbg.contains("***"));
290    }
291}
faucet_source_bigquery/config.rs

faucet_source_bigquery/
config.rs
