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}
57
58fn default_batch_size() -> usize {
59    DEFAULT_BATCH_SIZE
60}
61
62impl BigQuerySinkConfig {
63    /// Create a new config with the required fields and sensible defaults.
64    pub fn new(
65        project_id: impl Into<String>,
66        dataset_id: impl Into<String>,
67        table_id: impl Into<String>,
68        credentials: BigQueryCredentials,
69    ) -> Self {
70        Self {
71            project_id: project_id.into(),
72            dataset_id: dataset_id.into(),
73            table_id: table_id.into(),
74            auth: credentials,
75            batch_size: DEFAULT_BATCH_SIZE,
76            insert_id_field: None,
77            write: faucet_core::WriteSpec::default(),
78        }
79    }
80
81    /// Set the record field used as the per-row BigQuery streaming `insertId`
82    /// for best-effort de-duplication on retry.
83    pub fn with_insert_id_field(mut self, field: impl Into<String>) -> Self {
84        self.insert_id_field = Some(field.into());
85        self
86    }
87
88    /// Set the per-request row count for `tabledata.insertAll`.
89    ///
90    /// Pass `0` to opt out of re-chunking — the sink forwards each upstream
91    /// [`StreamPage`](faucet_core::StreamPage) as a single `insertAll` call.
92    /// BigQuery's streaming-insert sweet spot is ~500 rows per request.
93    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
94        self.batch_size = batch_size;
95        self
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn batch_size_defaults_to_default_batch_size() {
105        let config = BigQuerySinkConfig::new(
106            "my-project",
107            "my_dataset",
108            "my_table",
109            BigQueryCredentials::ApplicationDefault,
110        );
111        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
112    }
113
114    #[test]
115    fn with_batch_size_overrides_default() {
116        let config =
117            BigQuerySinkConfig::new("proj", "ds", "tbl", BigQueryCredentials::ApplicationDefault)
118                .with_batch_size(500);
119        assert_eq!(config.batch_size, 500);
120    }
121
122    #[test]
123    fn config_stores_all_fields() {
124        let config = BigQuerySinkConfig::new(
125            "my-project",
126            "my_dataset",
127            "my_table",
128            BigQueryCredentials::ServiceAccountKeyPath {
129                path: "/path/to/key.json".into(),
130            },
131        );
132        assert_eq!(config.project_id, "my-project");
133        assert_eq!(config.dataset_id, "my_dataset");
134        assert_eq!(config.table_id, "my_table");
135        assert!(matches!(
136            config.auth,
137            BigQueryCredentials::ServiceAccountKeyPath { .. }
138        ));
139    }
140
141    #[test]
142    fn config_with_inline_key() {
143        let config = BigQuerySinkConfig::new(
144            "proj",
145            "ds",
146            "tbl",
147            BigQueryCredentials::ServiceAccountKey {
148                json: r#"{"type":"service_account"}"#.into(),
149            },
150        );
151        if let BigQueryCredentials::ServiceAccountKey { json } = &config.auth {
152            assert!(json.contains("service_account"));
153        } else {
154            panic!("expected ServiceAccountKey");
155        }
156    }
157
158    #[test]
159    fn config_builder_chaining() {
160        let config =
161            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
162                .with_batch_size(100)
163                .with_batch_size(250);
164        assert_eq!(config.batch_size, 250);
165    }
166
167    #[test]
168    fn config_clone() {
169        let config =
170            BigQuerySinkConfig::new("proj", "ds", "tbl", BigQueryCredentials::ApplicationDefault)
171                .with_batch_size(42);
172        let cloned = config.clone();
173        assert_eq!(cloned.project_id, "proj");
174        assert_eq!(cloned.batch_size, 42);
175    }
176
177    #[test]
178    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
179        let config =
180            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
181                .with_batch_size(0);
182        assert_eq!(config.batch_size, 0);
183        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
184    }
185
186    #[test]
187    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
188        let config =
189            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault)
190                .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
191        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
192    }
193
194    #[test]
195    fn insert_id_field_defaults_none_and_builder_sets_it() {
196        let config =
197            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault);
198        assert!(config.insert_id_field.is_none());
199        let config = config.with_insert_id_field("event_id");
200        assert_eq!(config.insert_id_field.as_deref(), Some("event_id"));
201    }
202
203    #[test]
204    fn insert_id_field_deserializes_from_json() {
205        let json = r#"{
206            "project_id": "p",
207            "dataset_id": "d",
208            "table_id": "t",
209            "auth": {"type": "application_default"},
210            "insert_id_field": "id"
211        }"#;
212        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
213        assert_eq!(config.insert_id_field.as_deref(), Some("id"));
214    }
215
216    #[test]
217    fn batch_size_deserializes_from_json() {
218        let json = r#"{
219            "project_id": "p",
220            "dataset_id": "d",
221            "table_id": "t",
222            "auth": {"type": "application_default"},
223            "batch_size": 250
224        }"#;
225        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
226        assert_eq!(config.batch_size, 250);
227    }
228
229    #[test]
230    fn batch_size_defaults_when_absent_in_json() {
231        let json = r#"{
232            "project_id": "p",
233            "dataset_id": "d",
234            "table_id": "t",
235            "auth": {"type": "application_default"}
236        }"#;
237        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
238        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
239    }
240
241    #[test]
242    fn write_mode_defaults_to_append() {
243        let config =
244            BigQuerySinkConfig::new("p", "d", "t", BigQueryCredentials::ApplicationDefault);
245        assert_eq!(config.write.write_mode, faucet_core::WriteMode::Append);
246        assert!(config.write.key.is_empty());
247    }
248
249    #[test]
250    fn write_spec_deserializes_flattened() {
251        let json = r#"{
252            "project_id": "p",
253            "dataset_id": "d",
254            "table_id": "t",
255            "auth": {"type": "application_default"},
256            "write_mode": "upsert",
257            "key": ["id"],
258            "delete_marker": {"field": "__op", "values": ["d"]}
259        }"#;
260        let config: BigQuerySinkConfig = serde_json::from_str(json).unwrap();
261        assert_eq!(config.write.write_mode, faucet_core::WriteMode::Upsert);
262        assert_eq!(config.write.key, vec!["id".to_string()]);
263        let dm = config.write.delete_marker.expect("delete_marker");
264        assert_eq!(dm.field, "__op");
265        assert_eq!(dm.values, vec!["d".to_string()]);
266    }
267}
faucet_sink_bigquery/config.rs

faucet_sink_bigquery/
config.rs
