1use crate::config::BigQuerySinkConfig;
4use crate::idempotent;
5use crate::merge;
6use async_trait::async_trait;
7use faucet_common_bigquery::build_client;
8use faucet_core::FaucetError;
9use faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL;
10use gcp_bigquery_client::Client;
11use gcp_bigquery_client::error::BQError;
12use gcp_bigquery_client::model::get_query_results_parameters::GetQueryResultsParameters;
13use gcp_bigquery_client::model::query_parameter::QueryParameter;
14use gcp_bigquery_client::model::query_parameter_type::QueryParameterType;
15use gcp_bigquery_client::model::query_parameter_value::QueryParameterValue;
16use gcp_bigquery_client::model::query_request::QueryRequest;
17use gcp_bigquery_client::model::query_response::{QueryResponse, ResultSet};
18use gcp_bigquery_client::model::table_data_insert_all_request::TableDataInsertAllRequest;
19use gcp_bigquery_client::model::table_data_insert_all_response::TableDataInsertAllResponse;
20use serde_json::Value;
21use std::time::Duration;
22use tokio::sync::RwLock;
23
24const IDEMPOTENT_JOB_TIMEOUT: Duration = Duration::from_secs(120);
28
29const JOB_POLL_LONG_POLL_MS: i32 = 10_000;
32
33fn is_table_not_found(err: &BQError) -> bool {
37 matches!(err, BQError::ResponseError { error } if error.error.code == 404)
38}
39
40fn deletes_to_payload(deletes: &[faucet_core::KeyTuple]) -> String {
43 let arr: Vec<Value> = deletes
44 .iter()
45 .map(|kt| {
46 let mut obj = serde_json::Map::new();
47 for (k, v) in &kt.0 {
48 obj.insert(k.clone(), v.clone());
49 }
50 Value::Object(obj)
51 })
52 .collect();
53 Value::Array(arr).to_string()
54}
55
56pub struct BigQuerySink {
59 config: BigQuerySinkConfig,
60 client: Client,
61 schema_cache: RwLock<Option<Vec<idempotent::FieldSpec>>>,
67}
68
69impl BigQuerySink {
70 pub async fn new(config: BigQuerySinkConfig) -> Result<Self, FaucetError> {
75 faucet_core::validate_batch_size(config.batch_size)?;
76 config.write.validate()?;
77 let client = build_client(&config.auth).await?;
78 Ok(Self {
79 config,
80 client,
81 schema_cache: RwLock::new(None),
82 })
83 }
84
85 #[doc(hidden)]
94 pub fn from_parts(config: BigQuerySinkConfig, client: Client) -> Self {
95 Self {
96 config,
97 client,
98 schema_cache: RwLock::new(None),
99 }
100 }
101
102 async fn insert_chunk_raw(
119 &self,
120 rows: &[Value],
121 skip_invalid_rows: bool,
122 ) -> Result<TableDataInsertAllResponse, FaucetError> {
123 let mut insert_request = TableDataInsertAllRequest::new();
124 if skip_invalid_rows {
125 insert_request.skip_invalid_rows();
126 }
127 for row in rows {
128 let insert_id = self.config.insert_id_field.as_ref().and_then(|field| {
132 row.get(field).map(|v| match v {
133 Value::String(s) => s.clone(),
134 other => other.to_string(),
135 })
136 });
137 insert_request.add_row(insert_id, row).map_err(|e| {
138 FaucetError::Sink(format!("failed to serialize row for BigQuery: {e}"))
139 })?;
140 }
141 self.client
142 .tabledata()
143 .insert_all(
144 &self.config.project_id,
145 &self.config.dataset_id,
146 &self.config.table_id,
147 insert_request,
148 )
149 .await
150 .map_err(|e| FaucetError::Sink(format!("BigQuery insertAll failed: {e}")))
151 }
152
153 async fn insert_batch(&self, rows: &[Value]) -> Result<usize, FaucetError> {
161 if rows.is_empty() {
162 return Ok(0);
163 }
164
165 let response = self.insert_chunk_raw(rows, false).await?;
170
171 if let Some(errors) = response.insert_errors
173 && !errors.is_empty()
174 {
175 let count = errors.len();
176 let first = &errors[0];
177 return Err(FaucetError::Sink(format!(
178 "BigQuery insertAll: {count} row(s) failed; first error on row {:?}: {:?}",
179 first.index,
180 first
181 .errors
182 .as_ref()
183 .and_then(|errs| errs.first())
184 .map(|e| &e.message),
185 )));
186 }
187
188 Ok(rows.len())
189 }
190
191 fn string_param(name: &str, value: &str) -> QueryParameter {
197 QueryParameter {
198 name: Some(name.to_string()),
199 parameter_type: Some(QueryParameterType {
200 r#type: "STRING".to_string(),
201 array_type: None,
202 struct_types: None,
203 }),
204 parameter_value: Some(QueryParameterValue {
205 value: Some(value.to_string()),
206 array_values: None,
207 struct_values: None,
208 }),
209 }
210 }
211
212 async fn fetch_schema_fields(&self) -> Result<Vec<idempotent::FieldSpec>, BQError> {
216 let table = self
217 .client
218 .table()
219 .get(
220 &self.config.project_id,
221 &self.config.dataset_id,
222 &self.config.table_id,
223 Some(vec!["schema"]),
224 )
225 .await?;
226 Ok(table
228 .schema
229 .fields
230 .as_ref()
231 .map(|fs| {
232 fs.iter()
233 .map(idempotent::FieldSpec::from_table_field)
234 .collect()
235 })
236 .unwrap_or_default())
237 }
238
239 async fn target_schema(&self) -> Result<Vec<idempotent::FieldSpec>, FaucetError> {
247 if let Some(fields) = self.schema_cache.read().await.as_ref() {
248 return Ok(fields.clone());
249 }
250 let mut guard = self.schema_cache.write().await;
254 if let Some(fields) = guard.as_ref() {
255 return Ok(fields.clone());
256 }
257 let fields = self
258 .fetch_schema_fields()
259 .await
260 .map_err(|e| FaucetError::Sink(format!("BigQuery tables.get (schema) failed: {e}")))?;
261 if fields.is_empty() {
262 return Err(FaucetError::Sink(format!(
263 "BigQuery target table {}.{}.{} has no schema fields; exactly-once \
264 delivery requires a table with a defined schema",
265 self.config.project_id, self.config.dataset_id, self.config.table_id
266 )));
267 }
268 *guard = Some(fields.clone());
269 Ok(fields)
270 }
271
272 fn table_ref(&self) -> String {
274 idempotent::table_ref(
275 &self.config.project_id,
276 &self.config.dataset_id,
277 &self.config.table_id,
278 )
279 }
280
281 async fn run_ddl(&self, sql: String) -> Result<(), FaucetError> {
285 let mut req = QueryRequest::new(sql);
286 req.use_legacy_sql = false;
287 let resp = self
288 .client
289 .job()
290 .query(&self.config.project_id, req)
291 .await
292 .map_err(|e| FaucetError::Sink(format!("BigQuery schema-evolution DDL failed: {e}")))?;
293 self.await_query_complete(resp).await
294 }
295
296 async fn ensure_commit_table(&self) -> Result<(), FaucetError> {
298 let sql =
299 idempotent::build_create_commit_table(&self.config.project_id, &self.config.dataset_id);
300 let mut req = QueryRequest::new(sql);
301 req.use_legacy_sql = false;
302 let resp = self
303 .client
304 .job()
305 .query(&self.config.project_id, req)
306 .await
307 .map_err(|e| FaucetError::Sink(format!("BigQuery commit-table create failed: {e}")))?;
308 self.await_query_complete(resp).await
309 }
310
311 async fn await_query_complete(&self, initial: QueryResponse) -> Result<(), FaucetError> {
322 let (job_id, location) = Self::job_reference(&initial)?;
323
324 if !initial.job_complete.unwrap_or(false) {
326 let started = std::time::Instant::now();
327 loop {
328 let params = GetQueryResultsParameters {
329 location: location.clone(),
330 timeout_ms: Some(JOB_POLL_LONG_POLL_MS),
331 max_results: Some(0),
332 ..Default::default()
333 };
334 let resp = self
335 .client
336 .job()
337 .get_query_results(&self.config.project_id, &job_id, params)
338 .await
339 .map_err(|e| {
340 FaucetError::Sink(format!("BigQuery jobs.getQueryResults failed: {e}"))
341 })?;
342 if resp.job_complete.unwrap_or(false) {
343 break;
344 }
345 if started.elapsed() >= IDEMPOTENT_JOB_TIMEOUT {
346 return Err(FaucetError::Sink(format!(
347 "BigQuery job '{job_id}' did not complete within {}s",
348 IDEMPOTENT_JOB_TIMEOUT.as_secs()
349 )));
350 }
351 tokio::time::sleep(Duration::from_millis(250)).await;
355 }
356 }
357
358 let job = self
366 .client
367 .job()
368 .get_job(&self.config.project_id, &job_id, location.as_deref())
369 .await
370 .map_err(|e| FaucetError::Sink(format!("BigQuery jobs.get failed: {e}")))?;
371 let status = job.status.ok_or_else(|| {
372 FaucetError::Sink(format!(
373 "BigQuery job '{job_id}' returned no status; cannot confirm durable commit"
374 ))
375 })?;
376 if let Some(err) = status.error_result {
377 return Err(FaucetError::Sink(format!(
378 "BigQuery query job '{job_id}' failed: {err}"
379 )));
380 }
381 match status.state.as_deref() {
382 Some("DONE") => Ok(()),
383 other => Err(FaucetError::Sink(format!(
384 "BigQuery job '{job_id}' is in state {other:?}, not DONE; cannot confirm durable commit"
385 ))),
386 }
387 }
388
389 fn job_reference(qr: &QueryResponse) -> Result<(String, Option<String>), FaucetError> {
391 let r = qr.job_reference.as_ref().ok_or_else(|| {
392 FaucetError::Sink("BigQuery query response missing jobReference".to_string())
393 })?;
394 let job_id = r
395 .job_id
396 .clone()
397 .ok_or_else(|| FaucetError::Sink("BigQuery jobReference missing jobId".to_string()))?;
398 Ok((job_id, r.location.clone()))
399 }
400
401 async fn run_upsert_script(
409 &self,
410 plan: &faucet_core::WritePlan,
411 token: Option<(&str, &str)>,
412 ) -> Result<usize, FaucetError> {
413 let columns = self.target_schema().await?;
414 merge::validate_keys_present(&columns, &self.config.write.key)?;
415
416 let has_upserts = !plan.upserts.is_empty();
417 let has_deletes = !plan.deletes.is_empty();
418 if !has_upserts && !has_deletes {
419 if token.is_none() {
422 return Ok(0);
423 }
424 }
425
426 let key = &self.config.write.key;
427 let (project, dataset, table) = (
428 &self.config.project_id,
429 &self.config.dataset_id,
430 &self.config.table_id,
431 );
432 let sql = match token {
433 Some(_) => merge::build_upsert_idempotent_sql(
434 &columns,
435 key,
436 has_upserts,
437 has_deletes,
438 project,
439 dataset,
440 table,
441 ),
442 None => merge::build_upsert_transaction_sql(
443 &columns,
444 key,
445 has_upserts,
446 has_deletes,
447 project,
448 dataset,
449 table,
450 ),
451 };
452
453 let mut params = Vec::new();
454 if has_upserts {
455 let payload = serde_json::to_string(&plan.upserts).map_err(|e| {
456 FaucetError::Sink(format!("bigquery upsert: serialize payload: {e}"))
457 })?;
458 params.push(Self::string_param("payload", &payload));
459 }
460 if has_deletes {
461 let deletes = deletes_to_payload(&plan.deletes);
462 params.push(Self::string_param("deletes", &deletes));
463 }
464 if let Some((scope, tok)) = token {
465 params.push(Self::string_param("scope", scope));
466 params.push(Self::string_param("token", tok));
467 }
468
469 let mut req = QueryRequest::new(sql);
470 req.use_legacy_sql = false;
471 req.parameter_mode = Some("NAMED".to_string());
472 if let Some((scope, tok)) = token {
476 req.request_id = Some(idempotent::build_request_id(scope, tok));
477 }
478 req.query_parameters = Some(params);
479
480 let resp = self
481 .client
482 .job()
483 .query(&self.config.project_id, req)
484 .await
485 .map_err(|e| FaucetError::Sink(format!("bigquery upsert write failed: {e}")))?;
486 self.await_query_complete(resp).await?;
487
488 Ok(plan.upserts.len() + plan.deletes.len())
489 }
490}
491
492#[async_trait]
493impl faucet_core::Sink for BigQuerySink {
494 fn connector_name(&self) -> &'static str {
495 "bigquery"
496 }
497
498 fn config_schema(&self) -> serde_json::Value {
499 serde_json::to_value(faucet_core::schema_for!(BigQuerySinkConfig))
500 .expect("schema serialization")
501 }
502
503 fn dataset_uri(&self) -> String {
504 format!(
505 "bigquery://{}.{}.{}",
506 self.config.project_id, self.config.dataset_id, self.config.table_id
507 )
508 }
509
510 async fn check(
521 &self,
522 ctx: &faucet_core::check::CheckContext,
523 ) -> Result<faucet_core::check::CheckReport, FaucetError> {
524 use faucet_core::check::{CheckReport, Probe};
525
526 let started = std::time::Instant::now();
527 let fqn = format!(
528 "{}.{}.{}",
529 self.config.project_id, self.config.dataset_id, self.config.table_id
530 );
531
532 let result = tokio::time::timeout(
533 ctx.timeout,
534 self.client.table().get(
535 &self.config.project_id,
536 &self.config.dataset_id,
537 &self.config.table_id,
538 Some(vec!["tableReference"]),
539 ),
540 )
541 .await;
542
543 let probe = match result {
544 Ok(Ok(_table)) => Probe::pass("auth", started.elapsed()),
545 Ok(Err(e)) => Probe::fail_hint(
546 "auth",
547 started.elapsed(),
548 format!("BigQuery tables.get on {fqn} failed: {e}"),
549 "Verify the service account has roles/bigquery.dataViewer (or \
550 read access) on the dataset and that the project, dataset, and \
551 table IDs are correct.",
552 ),
553 Err(_elapsed) => Probe::fail_hint(
554 "auth",
555 started.elapsed(),
556 format!(
557 "BigQuery tables.get on {fqn} timed out after {:?}",
558 ctx.timeout
559 ),
560 "Check network reachability to bigquery.googleapis.com and that \
561 credentials can be minted within the timeout.",
562 ),
563 };
564
565 Ok(CheckReport::single(probe))
566 }
567
568 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
577 if records.is_empty() {
578 return Ok(0);
579 }
580
581 if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
582 let plan = faucet_core::plan_writes(records, &self.config.write);
583 if let Some((idx, msg)) = plan.failed.first() {
584 return Err(FaucetError::Sink(format!(
585 "bigquery {}: row {idx}: {msg}",
586 self.config.write.write_mode.as_str()
587 )));
588 }
589 return self.run_upsert_script(&plan, None).await;
590 }
591
592 let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
593 vec![records]
596 } else {
597 records.chunks(self.config.batch_size).collect()
598 };
599
600 let mut total = 0;
601 for chunk in chunks {
602 total += self.insert_batch(chunk).await?;
603 }
604
605 tracing::info!(
606 table = %format!(
607 "{}.{}.{}",
608 self.config.project_id, self.config.dataset_id, self.config.table_id
609 ),
610 rows = total,
611 "BigQuery write complete"
612 );
613 Ok(total)
614 }
615
616 async fn write_batch_partial(
634 &self,
635 records: &[Value],
636 ) -> Result<Vec<faucet_core::RowOutcome>, FaucetError> {
637 use std::collections::HashMap;
638
639 if records.is_empty() {
640 return Ok(Vec::new());
641 }
642
643 if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
644 let plan = faucet_core::plan_writes(records, &self.config.write);
645 self.run_upsert_script(&plan, None).await?;
646 let mut outcomes: Vec<faucet_core::RowOutcome> =
647 records.iter().map(|_| Ok(())).collect();
648 for (idx, msg) in &plan.failed {
649 outcomes[*idx] = Err(FaucetError::Sink(format!(
650 "bigquery {}: {msg}",
651 self.config.write.write_mode.as_str()
652 )));
653 }
654 return Ok(outcomes);
655 }
656
657 let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
658 vec![records]
659 } else {
660 records.chunks(self.config.batch_size).collect()
661 };
662
663 let mut outcomes: Vec<faucet_core::RowOutcome> = Vec::with_capacity(records.len());
664
665 for chunk in chunks {
666 let response = self.insert_chunk_raw(chunk, true).await?;
673
674 let failed: HashMap<usize, String> = response
676 .insert_errors
677 .unwrap_or_default()
678 .into_iter()
679 .filter_map(|e| {
680 let idx = e.index? as usize;
681 let msg = e
682 .errors
683 .as_ref()
684 .and_then(|v| v.first())
685 .map(|er| er.message.clone().unwrap_or_default())
686 .unwrap_or_default();
687 Some((idx, msg))
688 })
689 .collect();
690
691 for i in 0..chunk.len() {
692 match failed.get(&i) {
693 Some(msg) => outcomes.push(Err(FaucetError::Sink(format!(
694 "BigQuery row rejected: {msg}"
695 )))),
696 None => outcomes.push(Ok(())),
697 }
698 }
699 }
700
701 Ok(outcomes)
702 }
703
704 fn supported_write_modes(&self) -> &'static [faucet_core::WriteMode] {
705 &[
706 faucet_core::WriteMode::Append,
707 faucet_core::WriteMode::Upsert,
708 faucet_core::WriteMode::Delete,
709 ]
710 }
711
712 fn dedups_by_key(&self) -> bool {
713 self.config.write.dedups_by_key()
714 }
715
716 fn supports_idempotent_writes(&self) -> bool {
717 true
718 }
719
720 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
723 self.ensure_commit_table().await?;
724 let mut req = QueryRequest::new(idempotent::build_select_token(
725 &self.config.project_id,
726 &self.config.dataset_id,
727 ));
728 req.use_legacy_sql = false;
729 req.parameter_mode = Some("NAMED".to_string());
730 req.query_parameters = Some(vec![Self::string_param("scope", scope)]);
731
732 let resp = self
733 .client
734 .job()
735 .query(&self.config.project_id, req)
736 .await
737 .map_err(|e| FaucetError::Sink(format!("BigQuery token read failed: {e}")))?;
738
739 if !resp.job_complete.unwrap_or(false) {
744 return Err(FaucetError::Sink(
745 "BigQuery watermark read did not complete synchronously".to_string(),
746 ));
747 }
748 if resp.schema.is_none() {
753 return Err(FaucetError::Sink(
754 "BigQuery watermark read returned no schema; cannot trust the token result"
755 .to_string(),
756 ));
757 }
758
759 let mut rs = ResultSet::new_from_query_response(resp);
760 if rs.next_row() {
761 rs.get_string_by_name(COMMIT_TOKEN_TOKEN_COL)
762 .map_err(|e| FaucetError::Sink(format!("BigQuery token decode failed: {e}")))
763 } else {
764 Ok(None)
765 }
766 }
767
768 async fn write_batch_idempotent(
779 &self,
780 records: &[Value],
781 scope: &str,
782 token: &str,
783 ) -> Result<usize, FaucetError> {
784 self.ensure_commit_table().await?;
785
786 if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
787 let plan = faucet_core::plan_writes(records, &self.config.write);
788 if let Some((idx, msg)) = plan.failed.first() {
789 return Err(FaucetError::Sink(format!(
790 "bigquery {}: row {idx}: {msg}",
791 self.config.write.write_mode.as_str()
792 )));
793 }
794 return self.run_upsert_script(&plan, Some((scope, token))).await;
795 }
796
797 let columns = self.target_schema().await?;
798
799 let payload = serde_json::to_string(records).map_err(|e| {
800 FaucetError::Sink(format!(
801 "BigQuery exactly-once: serialize page payload: {e}"
802 ))
803 })?;
804
805 let sql = idempotent::build_transaction_sql(
806 &columns,
807 &self.config.project_id,
808 &self.config.dataset_id,
809 &self.config.table_id,
810 );
811 let mut req = QueryRequest::new(sql);
812 req.use_legacy_sql = false;
813 req.parameter_mode = Some("NAMED".to_string());
814 req.request_id = Some(idempotent::build_request_id(scope, token));
815 req.query_parameters = Some(vec![
816 Self::string_param("payload", &payload),
817 Self::string_param("scope", scope),
818 Self::string_param("token", token),
819 ]);
820
821 let resp = self
822 .client
823 .job()
824 .query(&self.config.project_id, req)
825 .await
826 .map_err(|e| FaucetError::Sink(format!("BigQuery idempotent write failed: {e}")))?;
827 self.await_query_complete(resp).await?;
828
829 tracing::info!(
830 table = %format!(
831 "{}.{}.{}",
832 self.config.project_id, self.config.dataset_id, self.config.table_id
833 ),
834 rows = records.len(),
835 token = %token,
836 "BigQuery exactly-once page committed"
837 );
838 Ok(records.len())
839 }
840
841 fn supports_schema_evolution(&self) -> bool {
846 true
847 }
848
849 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
857 match self.fetch_schema_fields().await {
858 Ok(fields) if fields.is_empty() => Ok(None),
859 Ok(fields) => Ok(Some(idempotent::fieldspecs_to_json_schema(&fields))),
860 Err(e) if is_table_not_found(&e) => Ok(None),
861 Err(e) => Err(FaucetError::Sink(format!(
862 "BigQuery current_schema (tables.get) failed: {e}"
863 ))),
864 }
865 }
866
867 async fn evolve_schema(
879 &self,
880 evolution: &faucet_core::SchemaEvolution,
881 ) -> Result<(), FaucetError> {
882 let table_ref = self.table_ref();
883
884 for c in &evolution.additions {
885 let bq = idempotent::base_to_bq(
886 faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text),
887 );
888 self.run_ddl(idempotent::build_add_column_ddl(&table_ref, &c.name, bq))
889 .await?;
890 }
891 for c in &evolution.widenings {
892 let bq = idempotent::base_to_bq(
893 faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text),
894 );
895 self.run_ddl(idempotent::build_alter_type_ddl(&table_ref, &c.name, bq))
896 .await?;
897 }
898 for col in &evolution.relax_nullability {
899 self.run_ddl(idempotent::build_drop_not_null_ddl(&table_ref, col))
900 .await?;
901 }
902
903 *self.schema_cache.write().await = None;
906 Ok(())
907 }
908}
909
910#[cfg(test)]
911mod tests {
912 use super::deletes_to_payload;
913 use faucet_core::KeyTuple;
914 use serde_json::json;
915
916 #[test]
921 fn deletes_to_payload_preserves_number_type() {
922 let p = deletes_to_payload(&[KeyTuple(vec![("id".into(), json!(2))])]);
926 let v: serde_json::Value = serde_json::from_str(&p).expect("valid JSON");
927 assert_eq!(v, json!([{"id": 2}]));
928 assert!(v[0]["id"].is_number(), "id must serialize as a number: {p}");
929 }
930
931 #[test]
932 fn deletes_to_payload_composite_key_roundtrips() {
933 let p = deletes_to_payload(&[KeyTuple(vec![
934 ("tenant".into(), json!("acme")),
935 ("id".into(), json!(7)),
936 ])]);
937 let v: serde_json::Value = serde_json::from_str(&p).expect("valid JSON");
938 assert_eq!(v, json!([{"tenant": "acme", "id": 7}]));
939 }
940
941 #[test]
942 fn deletes_to_payload_multiple_rows() {
943 let p = deletes_to_payload(&[
944 KeyTuple(vec![("id".into(), json!(1))]),
945 KeyTuple(vec![("id".into(), json!(2))]),
946 ]);
947 let v: serde_json::Value = serde_json::from_str(&p).expect("valid JSON");
948 assert_eq!(v, json!([{"id": 1}, {"id": 2}]));
949 }
950}