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 supports_idempotent_writes(&self) -> bool {
713        true
714    }
715
716    /// Read the last durably-committed token for `scope` from the watermark
717    /// table, so the pipeline can skip already-committed pages on resume.
718    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        // The watermark is a single tiny row, so `jobs.query` returns it inline.
736        // If BigQuery did not complete the read synchronously, fail safe: a
737        // wrong `None` here would re-run an already-committed page and produce
738        // duplicates, defeating exactly-once.
739        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        // `ResultSet` only yields rows when the response carries a schema; a
745        // completed `SELECT` always returns one. If it is somehow absent we
746        // cannot tell "no committed token" from "row present but unreadable",
747        // and a wrong `None` would replay committed pages — fail safe instead.
748        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    /// Atomically write `records` and record `token` for `scope` in one BigQuery
765    /// multi-statement transaction: a typed `INSERT … SELECT FROM
766    /// UNNEST(JSON_QUERY_ARRAY(@payload))` plus a watermark `MERGE`. Either both
767    /// the rows and the token commit, or neither does — so a crash/resume skips
768    /// the already-committed page (zero duplicates) and a failed page replays
769    /// cleanly.
770    ///
771    /// The entire page is one atomic unit (no `batch_size` re-chunking — core
772    /// issues exactly one token per page), so the page must serialize within
773    /// BigQuery's ~10 MB `jobs.query` request limit.
774    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    // -----------------------------------------------------------------------
838    // Schema drift (issue #194)
839    // -----------------------------------------------------------------------
840
841    fn supports_schema_evolution(&self) -> bool {
842        true
843    }
844
845    /// Read the live destination schema via a schema-only `tables.get`, mapped
846    /// to an `infer_schema`-shaped object so the drift policy can diff a page
847    /// against the real table.
848    ///
849    /// Returns `Ok(None)` when the target table does not exist yet (404) or
850    /// carries no field definitions — both mean "no schema to diff against",
851    /// so the drift pass treats every page column as new.
852    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    /// Apply an additive schema evolution to the target table via `ALTER TABLE`
864    /// DDL (issue #194):
865    ///
866    /// - additions → `ADD COLUMN IF NOT EXISTS <col> <type>`
867    /// - widenings → `ALTER COLUMN <col> SET DATA TYPE <type>`
868    /// - nullability relaxations → `ALTER COLUMN <col> DROP NOT NULL`
869    ///
870    /// Each statement runs as its own `jobs.query` job, verified to completion
871    /// via the authoritative job-status check. Every statement is idempotent so
872    /// concurrent runs converge. The cached schema is invalidated afterwards so
873    /// the next page re-fetches the evolved table.
874    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        // Invalidate the cached schema so the next exactly-once / upsert page
900        // (and the next drift diff) reads the evolved table.
901        *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    // dataset_uri test is skipped: BigQuerySink::new() requires GCP credentials
913    // (build_client fetches auth in new()), and from_parts() requires a
914    // gcp_bigquery_client::Client which cannot be constructed without auth.
915
916    #[test]
917    fn deletes_to_payload_preserves_number_type() {
918        // The delete payload must keep an integer key as a JSON number (not the
919        // string "2"), so the matching `CAST(JSON_VALUE(d, '$.id') AS INT64)`
920        // semi-join compares like-for-like.
921        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}