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    /// Arrow columnar **Storage Read API** mode (#380). When `true`, the source
102    /// reads [`read_table`](Self::read_table) directly via the BigQuery Storage
103    /// Read gRPC API as Arrow `RecordBatch`es (no `jobs.query`), driving the
104    /// columnar fast path when the sink is also columnar and decoding Arrow →
105    /// JSON on the row path otherwise. Requires a binary built with this
106    /// crate's `arrow` feature and a `read_table`; the `query` field is ignored
107    /// in this mode. Full extract only — no incremental bookmark.
108    #[serde(default)]
109    pub read_api: bool,
110    /// Table to read in `read_api` mode: `dataset.table` (billed to
111    /// `project_id`) or a fully-qualified `project.dataset.table`. Required
112    /// when `read_api` is set; ignored otherwise.
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub read_table: Option<String>,
115    /// Optional Storage Read API `row_restriction` — a SQL predicate (e.g.
116    /// `state = "CA"`) pushed to the read session so BigQuery filters rows
117    /// server-side. Only used in `read_api` mode.
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub row_restriction: Option<String>,
120    /// Optional column projection for `read_api` mode — the columns to read.
121    /// Empty means all columns.
122    #[serde(default, skip_serializing_if = "Vec::is_empty")]
123    pub selected_fields: Vec<String>,
124    /// Maximum number of Storage Read API streams to request (`read_api`
125    /// mode). Defaults to 1 (a single ordered stream). Higher values let
126    /// BigQuery shard large tables; the source reads the returned streams
127    /// sequentially.
128    #[serde(default = "default_max_streams")]
129    pub max_streams: i32,
130}
131
132fn default_max_streams() -> i32 {
133    1
134}
135
136impl std::fmt::Debug for BigQuerySourceConfig {
137    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
138        f.debug_struct("BigQuerySourceConfig")
139            .field("project_id", &self.project_id)
140            .field("auth", &self.auth)
141            .field("query", &self.query)
142            .field("use_legacy_sql", &self.use_legacy_sql)
143            .field("location", &self.location)
144            .field("max_results_per_page", &self.max_results_per_page)
145            .field("params", &self.params)
146            .field("statement_timeout", &self.statement_timeout)
147            .field("poll_timeout", &self.poll_timeout)
148            .field("batch_size", &self.batch_size)
149            .field("read_api", &self.read_api)
150            .field("read_table", &self.read_table)
151            .field("row_restriction", &self.row_restriction)
152            .field("selected_fields", &self.selected_fields)
153            .field("max_streams", &self.max_streams)
154            .finish()
155    }
156}
157
158impl BigQuerySourceConfig {
159    /// Create a new config with required fields and sensible defaults.
160    pub fn new(
161        project_id: impl Into<String>,
162        credentials: BigQueryCredentials,
163        query: impl Into<String>,
164    ) -> Self {
165        Self {
166            project_id: project_id.into(),
167            auth: credentials,
168            query: query.into(),
169            use_legacy_sql: default_use_legacy_sql(),
170            location: None,
171            max_results_per_page: default_max_results_per_page(),
172            params: Vec::new(),
173            statement_timeout: default_statement_timeout(),
174            poll_timeout: default_poll_timeout(),
175            batch_size: DEFAULT_BATCH_SIZE,
176            read_api: false,
177            read_table: None,
178            row_restriction: None,
179            selected_fields: Vec::new(),
180            max_streams: default_max_streams(),
181        }
182    }
183
184    /// Enable Arrow Storage Read API mode reading `table` (`dataset.table` or
185    /// `project.dataset.table`) instead of running a query (#380).
186    pub fn with_read_api(mut self, table: impl Into<String>) -> Self {
187        self.read_api = true;
188        self.read_table = Some(table.into());
189        self
190    }
191
192    /// Enable BigQuery's legacy SQL dialect.
193    pub fn with_use_legacy_sql(mut self, use_legacy: bool) -> Self {
194        self.use_legacy_sql = use_legacy;
195        self
196    }
197
198    /// Pin the job to a specific location (e.g. `"EU"`).
199    pub fn with_location(mut self, location: impl Into<String>) -> Self {
200        self.location = Some(location.into());
201        self
202    }
203
204    /// Set the maximum row count per `getQueryResults` page.
205    pub fn with_max_results_per_page(mut self, max_results: i32) -> Self {
206        self.max_results_per_page = max_results;
207        self
208    }
209
210    /// Set the positional bind parameters for the query.
211    pub fn with_params(mut self, params: Vec<Value>) -> Self {
212        self.params = params;
213        self
214    }
215
216    /// Set the per-statement server-side timeout.
217    pub fn with_statement_timeout(mut self, timeout: Duration) -> Self {
218        self.statement_timeout = timeout;
219        self
220    }
221
222    /// Set the maximum wall-clock time spent polling for job completion
223    /// before giving up. Pass `Duration::ZERO` to poll forever.
224    pub fn with_poll_timeout(mut self, timeout: Duration) -> Self {
225        self.poll_timeout = timeout;
226        self
227    }
228
229    /// Set the records-per-page hint for [`Source::stream_pages`](faucet_core::Source::stream_pages).
230    ///
231    /// Pass `0` to opt out of batching — the entire result set is emitted
232    /// in a single page.
233    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
234        self.batch_size = batch_size;
235        self
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use serde_json::json;
243
244    fn sample() -> BigQuerySourceConfig {
245        BigQuerySourceConfig::new(
246            "my-project",
247            BigQueryCredentials::ApplicationDefault,
248            "SELECT id FROM events",
249        )
250    }
251
252    #[test]
253    fn default_config() {
254        let c = sample();
255        assert_eq!(c.project_id, "my-project");
256        assert!(!c.use_legacy_sql);
257        assert!(c.location.is_none());
258        assert_eq!(c.max_results_per_page, 1000);
259        assert!(c.params.is_empty());
260        assert_eq!(c.statement_timeout, Duration::from_secs(60));
261        assert_eq!(c.poll_timeout, Duration::from_secs(300));
262        assert_eq!(c.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
263    }
264
265    #[test]
266    fn builder_chaining() {
267        let c = sample()
268            .with_use_legacy_sql(true)
269            .with_location("EU")
270            .with_max_results_per_page(500)
271            .with_params(vec![json!("us-east")])
272            .with_statement_timeout(Duration::from_secs(30))
273            .with_batch_size(250);
274        assert!(c.use_legacy_sql);
275        assert_eq!(c.location.as_deref(), Some("EU"));
276        assert_eq!(c.max_results_per_page, 500);
277        assert_eq!(c.params, vec![json!("us-east")]);
278        assert_eq!(c.statement_timeout, Duration::from_secs(30));
279        assert_eq!(c.batch_size, 250);
280    }
281
282    #[test]
283    fn deserializes_minimal_json() {
284        let json = r#"{
285            "project_id": "my-project",
286            "auth": {"type": "application_default"},
287            "query": "SELECT 1"
288        }"#;
289        let c: BigQuerySourceConfig = serde_json::from_str(json).unwrap();
290        assert!(!c.use_legacy_sql);
291        assert_eq!(c.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
292        assert_eq!(c.statement_timeout, Duration::from_secs(60));
293        assert_eq!(c.max_results_per_page, 1000);
294    }
295
296    #[test]
297    fn deserializes_all_fields() {
298        let json = r#"{
299            "project_id": "p",
300            "auth": {"type": "application_default"},
301            "query": "SELECT 1",
302            "use_legacy_sql": true,
303            "location": "EU",
304            "max_results_per_page": 500,
305            "params": ["us-east"],
306            "statement_timeout": 30,
307            "batch_size": 250
308        }"#;
309        let c: BigQuerySourceConfig = serde_json::from_str(json).unwrap();
310        assert!(c.use_legacy_sql);
311        assert_eq!(c.location.as_deref(), Some("EU"));
312        assert_eq!(c.max_results_per_page, 500);
313        assert_eq!(c.statement_timeout, Duration::from_secs(30));
314        assert_eq!(c.batch_size, 250);
315    }
316
317    #[test]
318    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
319        let c = sample().with_batch_size(0);
320        assert!(faucet_core::validate_batch_size(c.batch_size).is_ok());
321    }
322
323    #[test]
324    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
325        let c = sample().with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
326        assert!(faucet_core::validate_batch_size(c.batch_size).is_err());
327    }
328
329    #[test]
330    fn read_api_defaults_and_builder() {
331        let c = sample();
332        assert!(!c.read_api);
333        assert!(c.read_table.is_none());
334        assert_eq!(c.max_streams, 1);
335        assert!(c.selected_fields.is_empty());
336
337        let c = sample().with_read_api("ds.events");
338        assert!(c.read_api);
339        assert_eq!(c.read_table.as_deref(), Some("ds.events"));
340        // Debug renders the new fields without panicking.
341        assert!(format!("{c:?}").contains("read_api: true"));
342    }
343
344    #[test]
345    fn read_api_fields_deserialize() {
346        let json = r#"{
347            "project_id": "p",
348            "auth": {"type": "application_default"},
349            "query": "",
350            "read_api": true,
351            "read_table": "ds.t",
352            "row_restriction": "x = 1",
353            "selected_fields": ["a", "b"],
354            "max_streams": 3
355        }"#;
356        let c: BigQuerySourceConfig = serde_json::from_str(json).unwrap();
357        assert!(c.read_api);
358        assert_eq!(c.read_table.as_deref(), Some("ds.t"));
359        assert_eq!(c.row_restriction.as_deref(), Some("x = 1"));
360        assert_eq!(c.selected_fields, vec!["a".to_string(), "b".to_string()]);
361        assert_eq!(c.max_streams, 3);
362    }
363
364    #[test]
365    fn debug_masks_inline_credentials() {
366        let c = BigQuerySourceConfig::new(
367            "p",
368            BigQueryCredentials::ServiceAccountKey {
369                json: "secret".into(),
370            },
371            "SELECT 1",
372        );
373        let dbg = format!("{c:?}");
374        assert!(!dbg.contains("secret"));
375        assert!(dbg.contains("***"));
376    }
377}
faucet_source_bigquery/config.rs

faucet_source_bigquery/
config.rs
