Skip to main content

faucet_sink_bigquery/
sink.rs

1//! BigQuery streaming insert sink.
2
3use 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
24/// Max wall-clock spent polling an idempotent-write / token-read job to
25/// completion before giving up. Exactly-once pages are small, so this is a
26/// generous safety cap, not a steady-state wait.
27const IDEMPOTENT_JOB_TIMEOUT: Duration = Duration::from_secs(120);
28
29/// Server-side long-poll window per `getQueryResults` completion check —
30/// BigQuery holds the connection open up to this long, so we don't busy-wait.
31const JOB_POLL_LONG_POLL_MS: i32 = 10_000;
32
33/// `true` when a `tables.get` error is a 404 (table does not exist) — used by
34/// `current_schema` to report a not-yet-created target as `Ok(None)` rather
35/// than a hard error.
36fn is_table_not_found(err: &BQError) -> bool {
37    matches!(err, BQError::ResponseError { error } if error.error.code == 404)
38}
39
40/// Serialize planned delete key tuples into a JSON array of `{key_col: value}`
41/// objects for the `@deletes` parameter consumed by the semi-join `DELETE`.
42fn 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
56/// A sink that writes JSON records to a BigQuery table using the streaming
57/// insert API (`tabledata.insertAll`).
58pub struct BigQuerySink {
59    config: BigQuerySinkConfig,
60    client: Client,
61    /// Target table schema, fetched lazily on the first exactly-once / upsert
62    /// call and reused for every page in the run. `None` until first read, and
63    /// reset to `None` by [`evolve_schema`](faucet_core::Sink::evolve_schema) so
64    /// the next page diffs against the evolved table. Unused on the plain
65    /// streaming path.
66    schema_cache: RwLock<Option<Vec<idempotent::FieldSpec>>>,
67}
68
69impl BigQuerySink {
70    /// Create a new BigQuery sink from the given configuration.
71    ///
72    /// This initialises the BigQuery client and authenticates with GCP.
73    /// Returns a [`FaucetError::Auth`] if authentication fails.
74    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    /// Construct a sink from a pre-built BigQuery client.
86    ///
87    /// This is a low-level escape hatch for callers that build their own
88    /// [`gcp_bigquery_client::Client`] — for example to target the
89    /// [`bigquery-emulator`](https://github.com/goccy/bigquery-emulator) via
90    /// [`ClientBuilder::with_v2_base_url`](gcp_bigquery_client::client_builder::ClientBuilder::with_v2_base_url),
91    /// or to drive a wiremock-backed test fixture. Production code should
92    /// prefer [`BigQuerySink::new`], which handles credential loading.
93    #[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    /// Issue a single `tabledata.insertAll` call and return the raw response.
103    ///
104    /// Returns `Err` only on transport-level or HTTP-level failures. Per-row
105    /// `insertErrors` in the response body are surfaced to the caller as-is;
106    /// it is the caller's responsibility to inspect them.
107    ///
108    /// `skip_invalid_rows` maps to BigQuery's `skipInvalidRows` flag. When
109    /// `false` (the all-or-nothing [`write_batch`](Self::write_batch) path) a
110    /// single invalid row makes BigQuery commit *nothing* and return per-row
111    /// errors. When `true` (the [`write_batch_partial`] DLQ path) BigQuery
112    /// commits every valid row and reports `insertErrors` only for the rejected
113    /// ones — which is what makes the per-row `Ok`/`Err` mapping in
114    /// `write_batch_partial` truthful (without it, the "good" siblings are
115    /// reported `Ok` but were never actually committed → silent data loss).
116    ///
117    /// [`write_batch_partial`]: faucet_core::Sink::write_batch_partial
118    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            // When `insert_id_field` is configured, send that field's value as
129            // the streaming `insertId` so BigQuery can de-duplicate retries
130            // (#78/#31). A row lacking the field is inserted without one.
131            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    /// Insert a single chunk of rows in one `tabledata.insertAll` call,
154    /// collapsing any per-row errors into a single [`FaucetError::Sink`].
155    ///
156    /// Used by [`write_batch`](Self::write_batch). Callers that need per-row
157    /// error granularity should use
158    /// [`write_batch_partial`](faucet_core::Sink::write_batch_partial) instead,
159    /// which calls [`insert_chunk_raw`](Self::insert_chunk_raw) directly.
160    async fn insert_batch(&self, rows: &[Value]) -> Result<usize, FaucetError> {
161        if rows.is_empty() {
162            return Ok(0);
163        }
164
165        // All-or-nothing path: `skipInvalidRows=false` so BigQuery commits the
166        // whole chunk or nothing. Any `insertErrors` below becomes an outer
167        // `Err`, so the pipeline aborts before the bookmark advances — no
168        // partial commit to resume past.
169        let response = self.insert_chunk_raw(rows, false).await?;
170
171        // Check for per-row errors.
172        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    // -----------------------------------------------------------------------
192    // Exactly-once helpers
193    // -----------------------------------------------------------------------
194
195    /// Build a NAMED STRING query parameter.
196    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    /// Fetch the target table's schema fields directly via `tables.get`, with no
213    /// caching. Returns the raw [`idempotent::FieldSpec`]s (possibly empty for a
214    /// schemaless table); a missing table surfaces as the client's `BQError`.
215    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        // Table.schema is TableSchema (not Option); TableSchema.fields is Option<Vec<...>>.
227        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    /// Fetch (once) and cache the target table's schema as
240    /// [`idempotent::FieldSpec`]s, returning an owned clone. Used by the
241    /// exactly-once / upsert write paths, which require a table with a defined
242    /// schema — a missing table or empty schema is a hard error here.
243    ///
244    /// The cache is reset by [`evolve_schema`](faucet_core::Sink::evolve_schema)
245    /// so a later page re-fetches the evolved schema.
246    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        // Miss: fetch under the write lock so concurrent callers don't each
251        // issue a redundant tables.get. Re-check after acquiring in case a
252        // racing writer already filled it.
253        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    /// Backtick-quoted fully-qualified `` `project.dataset.table` `` reference.
273    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    /// Run one schema-evolution DDL statement through the same `jobs.query` +
282    /// authoritative job-status-verify path the data writes use, mapping any
283    /// failure to [`FaucetError::Sink`].
284    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    /// Create the commit-token watermark table if it does not exist.
297    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    /// Wait for a query/script job to finish, then authoritatively verify it
312    /// succeeded. Returns `Ok(())` only once the job reaches a terminal state
313    /// with no `errorResult`.
314    ///
315    /// Why `get_job` rather than the response `errors` field: the client maps
316    /// only non-2xx HTTP to `Err`, so a job that fails at *runtime* (a CAST
317    /// failure, a NULL into a REQUIRED column, …) comes back as `Ok` with the
318    /// failure recorded in the job body. `Job.status.error_result` is the
319    /// authoritative terminal-failure signal; the `errors` array can also carry
320    /// non-fatal warnings, so it must not be treated as failure on its own.
321    async fn await_query_complete(&self, initial: QueryResponse) -> Result<(), FaucetError> {
322        let (job_id, location) = Self::job_reference(&initial)?;
323
324        // Phase 1 — wait for completion via server-side long-poll (not a busy wait).
325        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                // The server long-poll normally blocks until completion, but if
352                // it returns early, back off so a still-running job can't turn
353                // this into a tight request-hammering loop.
354                tokio::time::sleep(Duration::from_millis(250)).await;
355            }
356        }
357
358        // Phase 2 — authoritative success check via the job's errorResult.
359        //
360        // We require an explicit terminal `DONE` state with no `errorResult`.
361        // A missing `status`, a non-`DONE` state, or a present `errorResult` all
362        // mean we cannot confirm the transaction durably committed — fail safe
363        // (returning `Ok` here would advance the bookmark over data that may
364        // never have landed, the silent-data-loss failure mode).
365        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    /// Extract `(job_id, location)` from a query response's job reference.
390    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    /// Run a planned upsert/delete page as one BigQuery multi-statement
402    /// transaction. When `token` is `Some((scope, tok))` the watermark `MERGE`
403    /// is appended inside the same transaction (exactly-once + upsert).
404    ///
405    /// The caller must have already validated `plan.failed` is empty (and, for
406    /// the exactly-once path, ensured the commit table exists). Returns the
407    /// number of rows applied (upserts + deletes).
408    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            // Nothing planned; for the exactly-once path the bookmark still
420            // needs its token, so emit the watermark MERGE alone.
421            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        // MERGE-by-key is idempotent, so a retried request is harmless; for the
473        // exactly-once path a deterministic request_id additionally suppresses
474        // duplicate jobs from a retried HTTP request within BigQuery's window.
475        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    /// Preflight check (`faucet doctor`).
511    ///
512    /// Runs a single read-only `tables.get` against the configured
513    /// `project_id.dataset_id.table_id` using the already-authenticated
514    /// client built in [`BigQuerySink::new`]. This mints/uses the access
515    /// token and confirms the credentials can read the target table's
516    /// metadata — without inserting any rows. Auth, missing-dataset,
517    /// missing-table, and permission errors all surface as a `Fail` probe
518    /// with a remediation hint. The access token is never included in the
519    /// reason or hint.
520    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    /// Write records to BigQuery.
569    ///
570    /// When `config.batch_size > 0` and the input slice is larger than
571    /// `batch_size`, the slice is split into chunks of `batch_size` rows and
572    /// each chunk is sent as a separate `tabledata.insertAll` call. When
573    /// `config.batch_size == 0`, the entire slice is sent in a single
574    /// `insertAll` request — useful when upstream `StreamPage`s are already
575    /// sized for BigQuery's per-request limits.
576    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            // Sentinel: pass the entire upstream page through in a single
594            // insertAll call. Subject to BigQuery's ~10MB request limit.
595            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    /// Write records to BigQuery, returning a per-row outcome vector.
617    ///
618    /// Unlike [`write_batch`](faucet_core::Sink::write_batch), which collapses all
619    /// `insertErrors` into a single `FaucetError`, this method maps each row
620    /// to `Ok(())` if BigQuery accepted it or `Err(FaucetError::Sink(...))` if
621    /// BigQuery reported a per-row error for it. This allows the pipeline's DLQ
622    /// router to quarantine only the rows that BigQuery actually rejected while
623    /// keeping already-committed siblings out of the dead-letter queue.
624    ///
625    /// Transport-level or HTTP-level failures (e.g. network errors, 4xx/5xx
626    /// responses) are still returned as an outer `Err` because no rows can be
627    /// considered committed in that case.
628    ///
629    /// Chunking follows the same `batch_size` semantics as `write_batch`:
630    /// `batch_size == 0` sends the entire slice in one call; `batch_size > 0`
631    /// splits the slice into chunks and concatenates the per-row outcomes in
632    /// input order.
633    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            // `skipInvalidRows=true`: BigQuery commits every valid row and
667            // returns `insertErrors` only for the rejected ones. This is what
668            // makes mapping the flagged indices to `Err` and the rest to
669            // `Ok(())` correct — the unflagged rows really were committed, so
670            // the DLQ router quarantines only the bad rows and the bookmark
671            // advances over genuinely-persisted data.
672            let response = self.insert_chunk_raw(chunk, true).await?;
673
674            // Build a set of failed row indices → first error message.
675            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    /// Read the last durably-committed token for `scope` from the watermark
721    /// table, so the pipeline can skip already-committed pages on resume.
722    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        // The watermark is a single tiny row, so `jobs.query` returns it inline.
740        // If BigQuery did not complete the read synchronously, fail safe: a
741        // wrong `None` here would re-run an already-committed page and produce
742        // duplicates, defeating exactly-once.
743        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        // `ResultSet` only yields rows when the response carries a schema; a
749        // completed `SELECT` always returns one. If it is somehow absent we
750        // cannot tell "no committed token" from "row present but unreadable",
751        // and a wrong `None` would replay committed pages — fail safe instead.
752        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    /// Atomically write `records` and record `token` for `scope` in one BigQuery
769    /// multi-statement transaction: a typed `INSERT … SELECT FROM
770    /// UNNEST(JSON_QUERY_ARRAY(@payload))` plus a watermark `MERGE`. Either both
771    /// the rows and the token commit, or neither does — so a crash/resume skips
772    /// the already-committed page (zero duplicates) and a failed page replays
773    /// cleanly.
774    ///
775    /// The entire page is one atomic unit (no `batch_size` re-chunking — core
776    /// issues exactly one token per page), so the page must serialize within
777    /// BigQuery's ~10 MB `jobs.query` request limit.
778    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    // -----------------------------------------------------------------------
842    // Schema drift (issue #194)
843    // -----------------------------------------------------------------------
844
845    fn supports_schema_evolution(&self) -> bool {
846        true
847    }
848
849    /// Read the live destination schema via a schema-only `tables.get`, mapped
850    /// to an `infer_schema`-shaped object so the drift policy can diff a page
851    /// against the real table.
852    ///
853    /// Returns `Ok(None)` when the target table does not exist yet (404) or
854    /// carries no field definitions — both mean "no schema to diff against",
855    /// so the drift pass treats every page column as new.
856    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    /// Apply an additive schema evolution to the target table via `ALTER TABLE`
868    /// DDL (issue #194):
869    ///
870    /// - additions → `ADD COLUMN IF NOT EXISTS <col> <type>`
871    /// - widenings → `ALTER COLUMN <col> SET DATA TYPE <type>`
872    /// - nullability relaxations → `ALTER COLUMN <col> DROP NOT NULL`
873    ///
874    /// Each statement runs as its own `jobs.query` job, verified to completion
875    /// via the authoritative job-status check. Every statement is idempotent so
876    /// concurrent runs converge. The cached schema is invalidated afterwards so
877    /// the next page re-fetches the evolved table.
878    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        // Invalidate the cached schema so the next exactly-once / upsert page
904        // (and the next drift diff) reads the evolved table.
905        *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    // dataset_uri test is skipped: BigQuerySink::new() requires GCP credentials
917    // (build_client fetches auth in new()), and from_parts() requires a
918    // gcp_bigquery_client::Client which cannot be constructed without auth.
919
920    #[test]
921    fn deletes_to_payload_preserves_number_type() {
922        // The delete payload must keep an integer key as a JSON number (not the
923        // string "2"), so the matching `CAST(JSON_VALUE(d, '$.id') AS INT64)`
924        // semi-join compares like-for-like.
925        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}