faucet_sink_bigquery/
config.rs

1//! BigQuery sink configuration.
2
3use faucet_core::DEFAULT_BATCH_SIZE;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7// Re-export the shared credentials type so end-user imports remain stable
8// (`use faucet_sink_bigquery::BigQueryCredentials;` keeps working).
9pub use faucet_common_bigquery::BigQueryCredentials;
10
11/// Configuration for the BigQuery streaming insert sink.
12#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
13pub struct BigQuerySinkConfig {
14    /// GCP project ID.
15    pub project_id: String,
16    /// BigQuery dataset ID.
17    pub dataset_id: String,
18    /// BigQuery table ID.
19    pub table_id: String,
20    /// Authentication credentials. YAML/JSON key is `auth` for consistency with
21    /// every other connector's auth block.
22    pub auth: BigQueryCredentials,
23    /// Maximum rows per `tabledata.insertAll` request. Defaults to
24    /// [`DEFAULT_BATCH_SIZE`].
25    ///
26    /// When the upstream `StreamPage` carries more records than `batch_size`,
27    /// the sink slices the page into `batch_size`-row chunks and issues one
28    /// `insertAll` HTTP call per chunk. When `batch_size = 0`, the page is
29    /// sent as a single request — useful when the source already chunks to
30    /// BigQuery's preferred size (e.g. ~500 rows for streaming inserts).
31    ///
32    /// `batch_size = 0` is the "no batching" sentinel: the entire upstream
33    /// page is forwarded in one `insertAll` call, subject to BigQuery's
34    /// natural per-request limits (~10MB body, ~500 rows recommended).
35    /// Larger pages may exceed those limits — keep the default unless the
36    /// upstream `StreamPage` size is already tuned for BigQuery.
37    #[serde(default = "default_batch_size")]
38    pub batch_size: usize,
39    /// Optional record field whose value is sent as the BigQuery streaming
40    /// `insertId` for each row. BigQuery uses `insertId` for best-effort
41    /// de-duplication over a short window, so a stable per-row key here makes
42    /// streaming inserts resilient to transport retries (which are otherwise
43    /// at-least-once and can produce duplicate rows) (#78/#31). When `None`
44    /// (the default) no `insertId` is sent. A row missing the field is
45    /// inserted without an `insertId` (no dedup for that row).
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub insert_id_field: Option<String>,
48    /// Write mode (append / upsert / delete) plus the `key` columns and optional
49    /// `delete_marker`. Flattened, so `write_mode` / `key` / `delete_marker`
50    /// appear at the config top level. Defaults to append (every existing
51    /// config keeps working). Upsert/delete merge by `key` in place via a
52    /// BigQuery `MERGE` over the page (no staging table); `key` must be real
53    /// column(s) of the target table.
54    #[serde(flatten)]
55    pub write: faucet_core::WriteSpec,
56    /// Scoped/windowed overwrite (#518): with `write_mode: overwrite`, replace
57    /// only the rows matching this scope (e.g. a partition/date window) instead
58    /// of the whole table — the out-of-scope rows are preserved.
59    #[serde(default, skip_serializing_if = "Option::is_none")]
60    pub scope: Option<faucet_core::OverwriteScope>,
61    /// Arrow columnar **load-job** mode (#380): buffer Arrow `RecordBatch`es to
62    /// Parquet, stage them on a GCS bucket, then run a BigQuery `PARQUET` load
63    /// job (`jobs.insert`) instead of the per-row `insertAll` path. Only
64    /// present in `arrow` builds; drives the columnar fast path the pipeline
65    /// negotiates when the source and sink are both columnar. Load jobs are
66    /// append/truncate only, so the sink advertises columnar support **only**
67    /// when the write mode is `append` — upsert/delete stay on the MERGE path.
68    #[cfg(feature = "arrow")]
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub bulk_load: Option<BigQueryLoadConfig>,
71}
72
73/// GCS-staged Parquet load-job configuration for the Arrow columnar path.
74#[cfg(feature = "arrow")]
75#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
76pub struct BigQueryLoadConfig {
77    /// GCS bucket used to stage Parquet files before the load job. Files are
78    /// written as `gs://<bucket>/<staging_prefix><uuid>.parquet`.
79    pub staging_bucket: String,
80    /// Object-key prefix within the bucket. Default `faucet-bq-load/`. A
81    /// trailing `/` is added if missing.
82    #[serde(default = "default_staging_prefix")]
83    pub staging_prefix: String,
84    /// Credentials for the GCS staging upload (independent of the BigQuery
85    /// `auth` used for the load job). Defaults to Application Default
86    /// Credentials.
87    #[serde(default)]
88    pub gcs_auth: faucet_common_gcs::GcsCredentials,
89    /// BigQuery load `writeDisposition` — `WRITE_APPEND` (default),
90    /// `WRITE_TRUNCATE`, or `WRITE_EMPTY`.
91    #[serde(default = "default_write_disposition")]
92    pub write_disposition: String,
93    /// Optional GCS storage endpoint override (e.g. a fake-gcs-server host for
94    /// tests). `None` uses the real Google endpoint.
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    pub storage_host: Option<String>,
97}
98
99#[cfg(feature = "arrow")]
100fn default_staging_prefix() -> String {
101    "faucet-bq-load/".to_string()
102}
103
104#[cfg(feature = "arrow")]
105fn default_write_disposition() -> String {
106    "WRITE_APPEND".to_string()
107}
108
109fn default_batch_size() -> usize {
110    DEFAULT_BATCH_SIZE
111}
112
113impl BigQuerySinkConfig {
114    /// Create a new config with the required fields and sensible defaults.
115    pub fn new(
116        project_id: impl Into<String>,
117        dataset_id: impl Into<String>,
118        table_id: impl Into<String>,
119        credentials: BigQueryCredentials,
120    ) -> Self {
121        Self {
122            project_id: project_id.into(),
123            dataset_id: dataset_id.into(),
124            table_id: table_id.into(),
125            auth: credentials,
126            batch_size: DEFAULT_BATCH_SIZE,
127            insert_id_field: None,
128            write: faucet_core::WriteSpec::default(),
129            scope: None,
130            #[cfg(feature = "arrow")]
131            bulk_load: None,
132        }
133    }
134
135    /// Enable Arrow columnar bulk-load via a GCS-staged Parquet load job (#380).
136    #[cfg(feature = "arrow")]
137    pub fn with_bulk_load(mut self, load: BigQueryLoadConfig) -> Self {
138        self.bulk_load = Some(load);
139        self
140    }
141
142    /// Set the record field used as the per-row BigQuery streaming `insertId`
143    /// for best-effort de-duplication on retry.
144    pub fn with_insert_id_field(mut self, field: impl Into<String>) -> Self {
145        self.insert_id_field = Some(field.into());
146        self
147    }
148
149    /// Set the per-request row count for `tabledata.insertAll`.
150    ///
151    /// Pass `0` to opt out of re-chunking — the sink forwards each upstream
152    /// [`StreamPage`](faucet_core::StreamPage) as a single `insertAll` call.
153    /// BigQuery's streaming-insert sweet spot is ~500 rows per request.
154    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
155        self.batch_size = batch_size;
156        self
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn batch_size_defaults_to_default_batch_size() {
166        let config = BigQuerySinkConfig::new(
167            "my-project",
168            "my_dataset",
169            "my_table",
170            BigQueryCredentials::ApplicationDefault,
171        );
172        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
173    }
174
175    #[test]
176    fn with_batch_size_overrides_default() {
177        let config =
178            BigQuerySinkConfig::new("proj", "ds", "tbl", BigQueryCredentials::ApplicationDefault)
179                .with_batch_size(500);
180        assert_eq!(config.batch_size, 500);
181    }
182
183    #[test]
184    fn config_stores_all_fields() {
185        let config = BigQuerySinkConfig::new(
186            "my-project",
187            "my_dataset",
188            "my_table",
189            BigQueryCredentials::ServiceAccountKeyPath {
190                path: "/path/to/key.json".into(),
191            },
192        );
193        assert_eq!(config.project_id, "my-project");
194        assert_eq!(config.dataset_id, "my_dataset");
195        assert_eq!(config.table_id, "my_table");
196        assert!(matches!(
197            config.auth,
198            BigQueryCredentials::ServiceAccountKeyPath { .. }
199        ));
200    }
201
202    #[test]
203    fn config_with_inline_key() {
204        let config = BigQuerySinkConfig::new(
205            "proj",
206            "ds",
207            "tbl",
208            BigQueryCredentials::ServiceAccountKey {
209                json: r#"{"type":"service_account"}"#.into(),
210            },
211        );
212        if let BigQueryCredentials::ServiceAccountKey { json } = &config.auth {
213            assert!(json.contains("service_account"));
214        } else {
215            panic!("expected ServiceAccountKey");
216        }
217    }
218
219    #[test]
220    fn config_builder_chaining() {
221        let config =
222            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
223                .with_batch_size(100)
224                .with_batch_size(250);
225        assert_eq!(config.batch_size, 250);
226    }
227
228    #[test]
229    fn config_clone() {
230        let config =
231            BigQuerySinkConfig::new("proj", "ds", "tbl", BigQueryCredentials::ApplicationDefault)
232                .with_batch_size(42);
233        let cloned = config.clone();
234        assert_eq!(cloned.project_id, "proj");
235        assert_eq!(cloned.batch_size, 42);
236    }
237
238    #[test]
239    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
240        let config =
241            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
242                .with_batch_size(0);
243        assert_eq!(config.batch_size, 0);
244        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
245    }
246
247    #[test]
248    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
249        let config =
250            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
251                .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
252        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
253    }
254
255    #[test]
256    fn insert_id_field_defaults_none_and_builder_sets_it() {
257        let config =
258            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault);
259        assert!(config.insert_id_field.is_none());
260        let config = config.with_insert_id_field("event_id");
261        assert_eq!(config.insert_id_field.as_deref(), Some("event_id"));
262    }
263
264    #[test]
265    fn insert_id_field_deserializes_from_json() {
266        let json = r#"{
267            "project_id": "p",
268            "dataset_id": "d",
269            "table_id": "t",
270            "auth": {"type": "application_default"},
271            "insert_id_field": "id"
272        }"#;
273        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
274        assert_eq!(config.insert_id_field.as_deref(), Some("id"));
275    }
276
277    #[test]
278    fn batch_size_deserializes_from_json() {
279        let json = r#"{
280            "project_id": "p",
281            "dataset_id": "d",
282            "table_id": "t",
283            "auth": {"type": "application_default"},
284            "batch_size": 250
285        }"#;
286        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
287        assert_eq!(config.batch_size, 250);
288    }
289
290    #[test]
291    fn batch_size_defaults_when_absent_in_json() {
292        let json = r#"{
293            "project_id": "p",
294            "dataset_id": "d",
295            "table_id": "t",
296            "auth": {"type": "application_default"}
297        }"#;
298        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
299        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
300    }
301
302    #[cfg(feature = "arrow")]
303    #[test]
304    fn bulk_load_builder_and_defaults() {
305        let cfg = BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
306            .with_bulk_load(BigQueryLoadConfig {
307                staging_bucket: "b".into(),
308                staging_prefix: default_staging_prefix(),
309                gcs_auth: Default::default(),
310                write_disposition: default_write_disposition(),
311                storage_host: None,
312            });
313        let load = cfg.bulk_load.expect("bulk_load set");
314        assert_eq!(load.staging_bucket, "b");
315        assert_eq!(load.staging_prefix, "faucet-bq-load/");
316        assert_eq!(load.write_disposition, "WRITE_APPEND");
317
318        // JSON: staging_prefix + write_disposition default when omitted.
319        let json = r#"{ "staging_bucket": "bk" }"#;
320        let l: BigQueryLoadConfig = serde_json::from_str(json).unwrap();
321        assert_eq!(l.staging_prefix, "faucet-bq-load/");
322        assert_eq!(l.write_disposition, "WRITE_APPEND");
323        assert!(l.storage_host.is_none());
324    }
325
326    #[test]
327    fn write_mode_defaults_to_append() {
328        let config =
329            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault);
330        assert_eq!(config.write.write_mode, faucet_core::WriteMode::Append);
331        assert!(config.write.key.is_empty());
332    }
333
334    #[test]
335    fn write_spec_deserializes_flattened() {
336        let json = r#"{
337            "project_id": "p",
338            "dataset_id": "d",
339            "table_id": "t",
340            "auth": {"type": "application_default"},
341            "write_mode": "upsert",
342            "key": ["id"],
343            "delete_marker": {"field": "__op", "values": ["d"]}
344        }"#;
345        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
346        assert_eq!(config.write.write_mode, faucet_core::WriteMode::Upsert);
347        assert_eq!(config.write.key, vec!["id".to_string()]);
348        let dm = config.write.delete_marker.expect("delete_marker");
349        assert_eq!(dm.field, "__op");
350        assert_eq!(dm.values, vec!["d".to_string()]);
351    }
352}
faucet_sink_bigquery/config.rs

faucet_sink_bigquery/
config.rs
