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    /// Arrow columnar **load-job** mode (#380): buffer Arrow `RecordBatch`es to
57    /// Parquet, stage them on a GCS bucket, then run a BigQuery `PARQUET` load
58    /// job (`jobs.insert`) instead of the per-row `insertAll` path. Only
59    /// present in `arrow` builds; drives the columnar fast path the pipeline
60    /// negotiates when the source and sink are both columnar. Load jobs are
61    /// append/truncate only, so the sink advertises columnar support **only**
62    /// when the write mode is `append` — upsert/delete stay on the MERGE path.
63    #[cfg(feature = "arrow")]
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub bulk_load: Option<BigQueryLoadConfig>,
66}
67
68/// GCS-staged Parquet load-job configuration for the Arrow columnar path.
69#[cfg(feature = "arrow")]
70#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
71pub struct BigQueryLoadConfig {
72    /// GCS bucket used to stage Parquet files before the load job. Files are
73    /// written as `gs://<bucket>/<staging_prefix><uuid>.parquet`.
74    pub staging_bucket: String,
75    /// Object-key prefix within the bucket. Default `faucet-bq-load/`. A
76    /// trailing `/` is added if missing.
77    #[serde(default = "default_staging_prefix")]
78    pub staging_prefix: String,
79    /// Credentials for the GCS staging upload (independent of the BigQuery
80    /// `auth` used for the load job). Defaults to Application Default
81    /// Credentials.
82    #[serde(default)]
83    pub gcs_auth: faucet_common_gcs::GcsCredentials,
84    /// BigQuery load `writeDisposition` — `WRITE_APPEND` (default),
85    /// `WRITE_TRUNCATE`, or `WRITE_EMPTY`.
86    #[serde(default = "default_write_disposition")]
87    pub write_disposition: String,
88    /// Optional GCS storage endpoint override (e.g. a fake-gcs-server host for
89    /// tests). `None` uses the real Google endpoint.
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub storage_host: Option<String>,
92}
93
94#[cfg(feature = "arrow")]
95fn default_staging_prefix() -> String {
96    "faucet-bq-load/".to_string()
97}
98
99#[cfg(feature = "arrow")]
100fn default_write_disposition() -> String {
101    "WRITE_APPEND".to_string()
102}
103
104fn default_batch_size() -> usize {
105    DEFAULT_BATCH_SIZE
106}
107
108impl BigQuerySinkConfig {
109    /// Create a new config with the required fields and sensible defaults.
110    pub fn new(
111        project_id: impl Into<String>,
112        dataset_id: impl Into<String>,
113        table_id: impl Into<String>,
114        credentials: BigQueryCredentials,
115    ) -> Self {
116        Self {
117            project_id: project_id.into(),
118            dataset_id: dataset_id.into(),
119            table_id: table_id.into(),
120            auth: credentials,
121            batch_size: DEFAULT_BATCH_SIZE,
122            insert_id_field: None,
123            write: faucet_core::WriteSpec::default(),
124            #[cfg(feature = "arrow")]
125            bulk_load: None,
126        }
127    }
128
129    /// Enable Arrow columnar bulk-load via a GCS-staged Parquet load job (#380).
130    #[cfg(feature = "arrow")]
131    pub fn with_bulk_load(mut self, load: BigQueryLoadConfig) -> Self {
132        self.bulk_load = Some(load);
133        self
134    }
135
136    /// Set the record field used as the per-row BigQuery streaming `insertId`
137    /// for best-effort de-duplication on retry.
138    pub fn with_insert_id_field(mut self, field: impl Into<String>) -> Self {
139        self.insert_id_field = Some(field.into());
140        self
141    }
142
143    /// Set the per-request row count for `tabledata.insertAll`.
144    ///
145    /// Pass `0` to opt out of re-chunking — the sink forwards each upstream
146    /// [`StreamPage`](faucet_core::StreamPage) as a single `insertAll` call.
147    /// BigQuery's streaming-insert sweet spot is ~500 rows per request.
148    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
149        self.batch_size = batch_size;
150        self
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn batch_size_defaults_to_default_batch_size() {
160        let config = BigQuerySinkConfig::new(
161            "my-project",
162            "my_dataset",
163            "my_table",
164            BigQueryCredentials::ApplicationDefault,
165        );
166        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
167    }
168
169    #[test]
170    fn with_batch_size_overrides_default() {
171        let config =
172            BigQuerySinkConfig::new("proj", "ds", "tbl", BigQueryCredentials::ApplicationDefault)
173                .with_batch_size(500);
174        assert_eq!(config.batch_size, 500);
175    }
176
177    #[test]
178    fn config_stores_all_fields() {
179        let config = BigQuerySinkConfig::new(
180            "my-project",
181            "my_dataset",
182            "my_table",
183            BigQueryCredentials::ServiceAccountKeyPath {
184                path: "/path/to/key.json".into(),
185            },
186        );
187        assert_eq!(config.project_id, "my-project");
188        assert_eq!(config.dataset_id, "my_dataset");
189        assert_eq!(config.table_id, "my_table");
190        assert!(matches!(
191            config.auth,
192            BigQueryCredentials::ServiceAccountKeyPath { .. }
193        ));
194    }
195
196    #[test]
197    fn config_with_inline_key() {
198        let config = BigQuerySinkConfig::new(
199            "proj",
200            "ds",
201            "tbl",
202            BigQueryCredentials::ServiceAccountKey {
203                json: r#"{"type":"service_account"}"#.into(),
204            },
205        );
206        if let BigQueryCredentials::ServiceAccountKey { json } = &config.auth {
207            assert!(json.contains("service_account"));
208        } else {
209            panic!("expected ServiceAccountKey");
210        }
211    }
212
213    #[test]
214    fn config_builder_chaining() {
215        let config =
216            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
217                .with_batch_size(100)
218                .with_batch_size(250);
219        assert_eq!(config.batch_size, 250);
220    }
221
222    #[test]
223    fn config_clone() {
224        let config =
225            BigQuerySinkConfig::new("proj", "ds", "tbl", BigQueryCredentials::ApplicationDefault)
226                .with_batch_size(42);
227        let cloned = config.clone();
228        assert_eq!(cloned.project_id, "proj");
229        assert_eq!(cloned.batch_size, 42);
230    }
231
232    #[test]
233    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
234        let config =
235            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
236                .with_batch_size(0);
237        assert_eq!(config.batch_size, 0);
238        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
239    }
240
241    #[test]
242    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
243        let config =
244            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
245                .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
246        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
247    }
248
249    #[test]
250    fn insert_id_field_defaults_none_and_builder_sets_it() {
251        let config =
252            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault);
253        assert!(config.insert_id_field.is_none());
254        let config = config.with_insert_id_field("event_id");
255        assert_eq!(config.insert_id_field.as_deref(), Some("event_id"));
256    }
257
258    #[test]
259    fn insert_id_field_deserializes_from_json() {
260        let json = r#"{
261            "project_id": "p",
262            "dataset_id": "d",
263            "table_id": "t",
264            "auth": {"type": "application_default"},
265            "insert_id_field": "id"
266        }"#;
267        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
268        assert_eq!(config.insert_id_field.as_deref(), Some("id"));
269    }
270
271    #[test]
272    fn batch_size_deserializes_from_json() {
273        let json = r#"{
274            "project_id": "p",
275            "dataset_id": "d",
276            "table_id": "t",
277            "auth": {"type": "application_default"},
278            "batch_size": 250
279        }"#;
280        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
281        assert_eq!(config.batch_size, 250);
282    }
283
284    #[test]
285    fn batch_size_defaults_when_absent_in_json() {
286        let json = r#"{
287            "project_id": "p",
288            "dataset_id": "d",
289            "table_id": "t",
290            "auth": {"type": "application_default"}
291        }"#;
292        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
293        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
294    }
295
296    #[cfg(feature = "arrow")]
297    #[test]
298    fn bulk_load_builder_and_defaults() {
299        let cfg = BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
300            .with_bulk_load(BigQueryLoadConfig {
301                staging_bucket: "b".into(),
302                staging_prefix: default_staging_prefix(),
303                gcs_auth: Default::default(),
304                write_disposition: default_write_disposition(),
305                storage_host: None,
306            });
307        let load = cfg.bulk_load.expect("bulk_load set");
308        assert_eq!(load.staging_bucket, "b");
309        assert_eq!(load.staging_prefix, "faucet-bq-load/");
310        assert_eq!(load.write_disposition, "WRITE_APPEND");
311
312        // JSON: staging_prefix + write_disposition default when omitted.
313        let json = r#"{ "staging_bucket": "bk" }"#;
314        let l: BigQueryLoadConfig = serde_json::from_str(json).unwrap();
315        assert_eq!(l.staging_prefix, "faucet-bq-load/");
316        assert_eq!(l.write_disposition, "WRITE_APPEND");
317        assert!(l.storage_host.is_none());
318    }
319
320    #[test]
321    fn write_mode_defaults_to_append() {
322        let config =
323            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault);
324        assert_eq!(config.write.write_mode, faucet_core::WriteMode::Append);
325        assert!(config.write.key.is_empty());
326    }
327
328    #[test]
329    fn write_spec_deserializes_flattened() {
330        let json = r#"{
331            "project_id": "p",
332            "dataset_id": "d",
333            "table_id": "t",
334            "auth": {"type": "application_default"},
335            "write_mode": "upsert",
336            "key": ["id"],
337            "delete_marker": {"field": "__op", "values": ["d"]}
338        }"#;
339        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
340        assert_eq!(config.write.write_mode, faucet_core::WriteMode::Upsert);
341        assert_eq!(config.write.key, vec!["id".to_string()]);
342        let dm = config.write.delete_marker.expect("delete_marker");
343        assert_eq!(dm.field, "__op");
344        assert_eq!(dm.values, vec!["d".to_string()]);
345    }
346}
faucet_sink_bigquery/config.rs

faucet_sink_bigquery/
config.rs
