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 supports_idempotent_writes(&self) -> bool {
713 true
714 }
715
716 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
719 self.ensure_commit_table().await?;
720 let mut req = QueryRequest::new(idempotent::build_select_token(
721 &self.config.project_id,
722 &self.config.dataset_id,
723 ));
724 req.use_legacy_sql = false;
725 req.parameter_mode = Some("NAMED".to_string());
726 req.query_parameters = Some(vec![Self::string_param("scope", scope)]);
727
728 let resp = self
729 .client
730 .job()
731 .query(&self.config.project_id, req)
732 .await
733 .map_err(|e| FaucetError::Sink(format!("BigQuery token read failed: {e}")))?;
734
735 if !resp.job_complete.unwrap_or(false) {
740 return Err(FaucetError::Sink(
741 "BigQuery watermark read did not complete synchronously".to_string(),
742 ));
743 }
744 if resp.schema.is_none() {
749 return Err(FaucetError::Sink(
750 "BigQuery watermark read returned no schema; cannot trust the token result"
751 .to_string(),
752 ));
753 }
754
755 let mut rs = ResultSet::new_from_query_response(resp);
756 if rs.next_row() {
757 rs.get_string_by_name(COMMIT_TOKEN_TOKEN_COL)
758 .map_err(|e| FaucetError::Sink(format!("BigQuery token decode failed: {e}")))
759 } else {
760 Ok(None)
761 }
762 }
763
764 async fn write_batch_idempotent(
775 &self,
776 records: &[Value],
777 scope: &str,
778 token: &str,
779 ) -> Result<usize, FaucetError> {
780 self.ensure_commit_table().await?;
781
782 if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
783 let plan = faucet_core::plan_writes(records, &self.config.write);
784 if let Some((idx, msg)) = plan.failed.first() {
785 return Err(FaucetError::Sink(format!(
786 "bigquery {}: row {idx}: {msg}",
787 self.config.write.write_mode.as_str()
788 )));
789 }
790 return self.run_upsert_script(&plan, Some((scope, token))).await;
791 }
792
793 let columns = self.target_schema().await?;
794
795 let payload = serde_json::to_string(records).map_err(|e| {
796 FaucetError::Sink(format!(
797 "BigQuery exactly-once: serialize page payload: {e}"
798 ))
799 })?;
800
801 let sql = idempotent::build_transaction_sql(
802 &columns,
803 &self.config.project_id,
804 &self.config.dataset_id,
805 &self.config.table_id,
806 );
807 let mut req = QueryRequest::new(sql);
808 req.use_legacy_sql = false;
809 req.parameter_mode = Some("NAMED".to_string());
810 req.request_id = Some(idempotent::build_request_id(scope, token));
811 req.query_parameters = Some(vec![
812 Self::string_param("payload", &payload),
813 Self::string_param("scope", scope),
814 Self::string_param("token", token),
815 ]);
816
817 let resp = self
818 .client
819 .job()
820 .query(&self.config.project_id, req)
821 .await
822 .map_err(|e| FaucetError::Sink(format!("BigQuery idempotent write failed: {e}")))?;
823 self.await_query_complete(resp).await?;
824
825 tracing::info!(
826 table = %format!(
827 "{}.{}.{}",
828 self.config.project_id, self.config.dataset_id, self.config.table_id
829 ),
830 rows = records.len(),
831 token = %token,
832 "BigQuery exactly-once page committed"
833 );
834 Ok(records.len())
835 }
836
837 fn supports_schema_evolution(&self) -> bool {
842 true
843 }
844
845 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
853 match self.fetch_schema_fields().await {
854 Ok(fields) if fields.is_empty() => Ok(None),
855 Ok(fields) => Ok(Some(idempotent::fieldspecs_to_json_schema(&fields))),
856 Err(e) if is_table_not_found(&e) => Ok(None),
857 Err(e) => Err(FaucetError::Sink(format!(
858 "BigQuery current_schema (tables.get) failed: {e}"
859 ))),
860 }
861 }
862
863 async fn evolve_schema(
875 &self,
876 evolution: &faucet_core::SchemaEvolution,
877 ) -> Result<(), FaucetError> {
878 let table_ref = self.table_ref();
879
880 for c in &evolution.additions {
881 let bq = idempotent::base_to_bq(
882 faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text),
883 );
884 self.run_ddl(idempotent::build_add_column_ddl(&table_ref, &c.name, bq))
885 .await?;
886 }
887 for c in &evolution.widenings {
888 let bq = idempotent::base_to_bq(
889 faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text),
890 );
891 self.run_ddl(idempotent::build_alter_type_ddl(&table_ref, &c.name, bq))
892 .await?;
893 }
894 for col in &evolution.relax_nullability {
895 self.run_ddl(idempotent::build_drop_not_null_ddl(&table_ref, col))
896 .await?;
897 }
898
899 *self.schema_cache.write().await = None;
902 Ok(())
903 }
904}
905
906#[cfg(test)]
907mod tests {
908 use super::deletes_to_payload;
909 use faucet_core::KeyTuple;
910 use serde_json::json;
911
912 #[test]
917 fn deletes_to_payload_preserves_number_type() {
918 let p = deletes_to_payload(&[KeyTuple(vec![("id".into(), json!(2))])]);
922 let v: serde_json::Value = serde_json::from_str(&p).expect("valid JSON");
923 assert_eq!(v, json!([{"id": 2}]));
924 assert!(v[0]["id"].is_number(), "id must serialize as a number: {p}");
925 }
926
927 #[test]
928 fn deletes_to_payload_composite_key_roundtrips() {
929 let p = deletes_to_payload(&[KeyTuple(vec![
930 ("tenant".into(), json!("acme")),
931 ("id".into(), json!(7)),
932 ])]);
933 let v: serde_json::Value = serde_json::from_str(&p).expect("valid JSON");
934 assert_eq!(v, json!([{"tenant": "acme", "id": 7}]));
935 }
936
937 #[test]
938 fn deletes_to_payload_multiple_rows() {
939 let p = deletes_to_payload(&[
940 KeyTuple(vec![("id".into(), json!(1))]),
941 KeyTuple(vec![("id".into(), json!(2))]),
942 ]);
943 let v: serde_json::Value = serde_json::from_str(&p).expect("valid JSON");
944 assert_eq!(v, json!([{"id": 1}, {"id": 2}]));
945 }
946}