Skip to main content

risingwave_connector/sink/
big_query.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use core::pin::Pin;
16use core::time::Duration;
17use std::collections::{BTreeMap, HashMap, VecDeque};
18use std::time::{SystemTime, UNIX_EPOCH};
19
20use anyhow::{Context, anyhow};
21use async_trait::async_trait;
22use base64::Engine;
23use base64::prelude::BASE64_STANDARD;
24use futures::future::pending;
25use futures::prelude::Future;
26use futures::{Stream, StreamExt};
27use futures_async_stream::try_stream;
28use gcp_bigquery_client::Client;
29use gcp_bigquery_client::error::BQError;
30use gcp_bigquery_client::model::query_request::QueryRequest;
31use gcp_bigquery_client::model::query_response::ResultSet;
32use gcp_bigquery_client::model::table::Table;
33use gcp_bigquery_client::model::table_field_schema::TableFieldSchema;
34use gcp_bigquery_client::model::table_schema::TableSchema;
35use google_cloud_bigquery::grpc::apiv1::conn_pool::ConnectionManager;
36use google_cloud_gax::conn::{ConnectionOptions, Environment};
37use google_cloud_gax::grpc::{Request, Response, Status};
38use google_cloud_googleapis::cloud::bigquery::storage::v1::append_rows_request::{
39    MissingValueInterpretation, ProtoData, Rows as AppendRowsRequestRows,
40};
41use google_cloud_googleapis::cloud::bigquery::storage::v1::{
42    AppendRowsRequest, AppendRowsResponse, ProtoRows, ProtoSchema,
43};
44use google_cloud_pubsub::client::google_cloud_auth;
45use google_cloud_pubsub::client::google_cloud_auth::credentials::CredentialsFile;
46use parking_lot::Mutex;
47use phf::{Set, phf_set};
48use prost_reflect::{FieldDescriptor, MessageDescriptor};
49use prost_types::{
50    DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet,
51    field_descriptor_proto,
52};
53use risingwave_common::array::{Op, StreamChunk};
54use risingwave_common::catalog::{Field, Schema};
55use risingwave_common::types::DataType;
56use risingwave_common::util::iter_util::ZipEqDebug;
57use serde::Deserialize;
58use serde_with::{DisplayFromStr, serde_as};
59use simd_json::prelude::ArrayTrait;
60use tokio::sync::mpsc;
61use url::Url;
62use uuid::Uuid;
63use with_options::WithOptions;
64use yup_oauth2::ServiceAccountKey;
65
66use super::encoder::{ProtoEncoder, ProtoHeader, RowEncoder, SerTo};
67use super::log_store::{LogStoreReadItem, TruncateOffset};
68use super::{
69    LogSinker, SINK_TYPE_APPEND_ONLY, SINK_TYPE_OPTION, SINK_TYPE_UPSERT, SinkError, SinkLogReader,
70};
71use crate::aws_utils::load_file_descriptor_from_s3;
72use crate::connector_common::AwsAuthProps;
73use crate::enforce_secret::EnforceSecret;
74use crate::sink::{Result, Sink, SinkParam, SinkWriterParam};
75
76pub const BIGQUERY_SINK: &str = "bigquery";
77pub const CHANGE_TYPE: &str = "_CHANGE_TYPE";
78pub const CHANGE_SEQUENCE_NUMBER: &str = "_CHANGE_SEQUENCE_NUMBER";
79const DEFAULT_GRPC_CHANNEL_NUMS: usize = 4;
80const CONNECT_TIMEOUT: Option<Duration> = Some(Duration::from_secs(30));
81const CONNECTION_TIMEOUT: Option<Duration> = None;
82const BIGQUERY_SEND_FUTURE_BUFFER_MAX_SIZE: usize = 65536;
83// < 10MB, we set 8MB
84const MAX_ROW_SIZE: usize = 8 * 1024 * 1024;
85
86#[serde_as]
87#[derive(Deserialize, Debug, Clone, WithOptions)]
88pub struct BigQueryCommon {
89    #[serde(rename = "bigquery.local.path")]
90    pub local_path: Option<String>,
91    #[serde(rename = "bigquery.s3.path")]
92    pub s3_path: Option<String>,
93    #[serde(rename = "bigquery.project")]
94    pub project: String,
95    #[serde(rename = "bigquery.dataset")]
96    pub dataset: String,
97    #[serde(rename = "bigquery.table")]
98    pub table: String,
99    #[serde(default, alias = "create_table_if_not_exists")] // default false
100    #[serde_as(as = "DisplayFromStr")]
101    pub auto_create: bool,
102    #[serde(rename = "bigquery.credentials")]
103    pub credentials: Option<String>,
104}
105
106impl EnforceSecret for BigQueryCommon {
107    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
108        "bigquery.credentials",
109    };
110}
111
112struct BigQueryFutureManager {
113    // `offset_queue` holds the Some corresponding to each future.
114    // When TruncateOffset is barrier, the num is 0, we don't need to wait for the return of `resp_stream`.
115    // When TruncateOffset is chunk:
116    // 1. chunk has no rows. we didn't send, the num is 0, we don't need to wait for the return of `resp_stream`.
117    // 2. chunk is less than `MAX_ROW_SIZE`, we only sent once, the num is 1 and we only have to wait once for `resp_stream`.
118    // 3. chunk is less than `MAX_ROW_SIZE`, we only sent n, the num is n and we need to wait n times for r.
119    offset_queue: VecDeque<(TruncateOffset, usize)>,
120    resp_stream: Pin<Box<dyn Stream<Item = Result<()>> + Send>>,
121}
122impl BigQueryFutureManager {
123    pub fn new(
124        max_future_num: usize,
125        resp_stream: impl Stream<Item = Result<()>> + Send + 'static,
126    ) -> Self {
127        let offset_queue = VecDeque::with_capacity(max_future_num);
128        Self {
129            offset_queue,
130            resp_stream: Box::pin(resp_stream),
131        }
132    }
133
134    pub fn add_offset(&mut self, offset: TruncateOffset, resp_num: usize) {
135        self.offset_queue.push_back((offset, resp_num));
136    }
137
138    pub async fn next_offset(&mut self) -> Result<TruncateOffset> {
139        if let Some((_offset, remaining_resp_num)) = self.offset_queue.front_mut() {
140            if *remaining_resp_num == 0 {
141                return Ok(self.offset_queue.pop_front().unwrap().0);
142            }
143            while *remaining_resp_num > 0 {
144                self.resp_stream
145                    .next()
146                    .await
147                    .ok_or_else(|| SinkError::BigQuery(anyhow::anyhow!("end of stream")))??;
148                *remaining_resp_num -= 1;
149            }
150            Ok(self.offset_queue.pop_front().unwrap().0)
151        } else {
152            pending().await
153        }
154    }
155}
156pub struct BigQueryLogSinker {
157    writer: BigQuerySinkWriter,
158    bigquery_future_manager: BigQueryFutureManager,
159    future_num: usize,
160}
161impl BigQueryLogSinker {
162    pub fn new(
163        writer: BigQuerySinkWriter,
164        resp_stream: impl Stream<Item = Result<()>> + Send + 'static,
165        future_num: usize,
166    ) -> Self {
167        Self {
168            writer,
169            bigquery_future_manager: BigQueryFutureManager::new(future_num, resp_stream),
170            future_num,
171        }
172    }
173}
174
175#[async_trait]
176impl LogSinker for BigQueryLogSinker {
177    async fn consume_log_and_sink(mut self, mut log_reader: impl SinkLogReader) -> Result<!> {
178        log_reader.start_from(None).await?;
179        loop {
180            tokio::select!(
181                offset = self.bigquery_future_manager.next_offset() => {
182                        log_reader.truncate(offset?)?;
183                }
184                item_result = log_reader.next_item(), if self.bigquery_future_manager.offset_queue.len() <= self.future_num => {
185                    let (epoch, item) = item_result?;
186                    match item {
187                        LogStoreReadItem::StreamChunk { chunk_id, chunk } => {
188                            let resp_num = self.writer.write_chunk(chunk)?;
189                            self.bigquery_future_manager
190                                .add_offset(TruncateOffset::Chunk { epoch, chunk_id },resp_num);
191                        }
192                        LogStoreReadItem::Barrier { .. } => {
193                            self.bigquery_future_manager
194                                .add_offset(TruncateOffset::Barrier { epoch },0);
195                        }
196                    }
197                }
198            )
199        }
200    }
201}
202
203impl BigQueryCommon {
204    async fn build_client(&self, aws_auth_props: &AwsAuthProps) -> Result<Client> {
205        let auth_json = self.get_auth_json_from_path(aws_auth_props).await?;
206
207        let service_account =
208            if let Ok(auth_json_from_base64) = BASE64_STANDARD.decode(auth_json.clone()) {
209                serde_json::from_slice::<ServiceAccountKey>(&auth_json_from_base64)
210            } else {
211                serde_json::from_str::<ServiceAccountKey>(&auth_json)
212            }
213            .map_err(|e| SinkError::BigQuery(e.into()))?;
214
215        let client: Client = Client::from_service_account_key(service_account, false)
216            .await
217            .map_err(|err| SinkError::BigQuery(anyhow::anyhow!(err)))?;
218        Ok(client)
219    }
220
221    async fn build_writer_client(
222        &self,
223        aws_auth_props: &AwsAuthProps,
224    ) -> Result<(StorageWriterClient, impl Stream<Item = Result<()>> + use<>)> {
225        let auth_json = self.get_auth_json_from_path(aws_auth_props).await?;
226
227        let credentials_file =
228            if let Ok(auth_json_from_base64) = BASE64_STANDARD.decode(auth_json.clone()) {
229                serde_json::from_slice::<CredentialsFile>(&auth_json_from_base64)
230            } else {
231                serde_json::from_str::<CredentialsFile>(&auth_json)
232            }
233            .map_err(|e| SinkError::BigQuery(e.into()))?;
234
235        StorageWriterClient::new(credentials_file).await
236    }
237
238    async fn get_auth_json_from_path(&self, aws_auth_props: &AwsAuthProps) -> Result<String> {
239        if let Some(credentials) = &self.credentials {
240            Ok(credentials.clone())
241        } else if let Some(local_path) = &self.local_path {
242            std::fs::read_to_string(local_path)
243                .map_err(|err| SinkError::BigQuery(anyhow::anyhow!(err)))
244        } else if let Some(s3_path) = &self.s3_path {
245            let url =
246                Url::parse(s3_path).map_err(|err| SinkError::BigQuery(anyhow::anyhow!(err)))?;
247            let auth_vec = load_file_descriptor_from_s3(&url, aws_auth_props)
248                .await
249                .map_err(|err| SinkError::BigQuery(anyhow::anyhow!(err)))?;
250            Ok(String::from_utf8(auth_vec).map_err(|e| SinkError::BigQuery(e.into()))?)
251        } else {
252            Err(SinkError::BigQuery(anyhow::anyhow!(
253                "`bigquery.local.path` and `bigquery.s3.path` set at least one, configure as needed."
254            )))
255        }
256    }
257}
258
259#[serde_as]
260#[derive(Clone, Debug, Deserialize, WithOptions)]
261pub struct BigQueryConfig {
262    #[serde(flatten)]
263    pub common: BigQueryCommon,
264    #[serde(flatten)]
265    pub aws_auth_props: AwsAuthProps,
266    pub r#type: String, // accept "append-only" or "upsert"
267
268    #[serde(flatten)]
269    pub unknown_fields: std::collections::HashMap<String, String>,
270}
271
272crate::impl_sink_unknown_fields!(BigQueryConfig);
273
274impl EnforceSecret for BigQueryConfig {
275    fn enforce_one(prop: &str) -> crate::error::ConnectorResult<()> {
276        BigQueryCommon::enforce_one(prop)?;
277        AwsAuthProps::enforce_one(prop)?;
278        Ok(())
279    }
280}
281
282impl BigQueryConfig {
283    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
284        let config =
285            serde_json::from_value::<BigQueryConfig>(serde_json::to_value(properties).unwrap())
286                .map_err(|e| SinkError::Config(anyhow!(e)))?;
287        if config.r#type != SINK_TYPE_APPEND_ONLY && config.r#type != SINK_TYPE_UPSERT {
288            return Err(SinkError::Config(anyhow!(
289                "`{}` must be {}, or {}",
290                SINK_TYPE_OPTION,
291                SINK_TYPE_APPEND_ONLY,
292                SINK_TYPE_UPSERT
293            )));
294        }
295        Ok(config)
296    }
297}
298
299#[derive(Debug)]
300pub struct BigQuerySink {
301    pub config: BigQueryConfig,
302    schema: Schema,
303    pk_indices: Vec<usize>,
304    is_append_only: bool,
305}
306
307impl EnforceSecret for BigQuerySink {
308    fn enforce_secret<'a>(
309        prop_iter: impl Iterator<Item = &'a str>,
310    ) -> crate::error::ConnectorResult<()> {
311        for prop in prop_iter {
312            BigQueryConfig::enforce_one(prop)?;
313        }
314        Ok(())
315    }
316}
317
318impl BigQuerySink {
319    pub fn new(
320        config: BigQueryConfig,
321        schema: Schema,
322        pk_indices: Vec<usize>,
323        is_append_only: bool,
324    ) -> Result<Self> {
325        Ok(Self {
326            config,
327            schema,
328            pk_indices,
329            is_append_only,
330        })
331    }
332}
333
334impl BigQuerySink {
335    fn is_decimal_type_compatible(bigquery_type: &str) -> bool {
336        // BigQuery INFORMATION_SCHEMA reports parameterized decimal columns as
337        // `NUMERIC(p, s)` or `BIGNUMERIC(p, s)`. RisingWave `Decimal` does not
338        // carry typmod, so schema validation accepts the BigQuery decimal type
339        // family instead of requiring exact string equality.
340        let normalized = bigquery_type.trim().to_ascii_uppercase();
341        matches!(
342            normalized
343                .split_once('(')
344                .map_or(normalized.as_str(), |(prefix, _)| prefix),
345            "NUMERIC" | "BIGNUMERIC"
346        )
347    }
348
349    fn is_data_type_compatible(rw_data_type: &DataType, bigquery_type: &str) -> Result<bool> {
350        if matches!(rw_data_type, DataType::Decimal) {
351            return Ok(Self::is_decimal_type_compatible(bigquery_type));
352        }
353
354        Ok(Self::get_string_and_check_support_from_datatype(rw_data_type)? == bigquery_type)
355    }
356
357    fn check_column_name_and_type(
358        &self,
359        big_query_columns_desc: HashMap<String, String>,
360    ) -> Result<()> {
361        let rw_fields_name = self.schema.fields();
362        if big_query_columns_desc.is_empty() {
363            return Err(SinkError::BigQuery(anyhow::anyhow!(
364                "Cannot find table in bigquery"
365            )));
366        }
367        if rw_fields_name.len().ne(&big_query_columns_desc.len()) {
368            return Err(SinkError::BigQuery(anyhow::anyhow!(
369                "The length of the RisingWave column {} must be equal to the length of the bigquery column {}",
370                rw_fields_name.len(),
371                big_query_columns_desc.len()
372            )));
373        }
374
375        for i in rw_fields_name {
376            let value = big_query_columns_desc.get(&i.name).ok_or_else(|| {
377                SinkError::BigQuery(anyhow::anyhow!(
378                    "Column `{:?}` on RisingWave side is not found on BigQuery side.",
379                    i.name
380                ))
381            })?;
382            let data_type_string = Self::get_string_and_check_support_from_datatype(&i.data_type)?;
383            if !Self::is_data_type_compatible(&i.data_type, value)? {
384                return Err(SinkError::BigQuery(anyhow::anyhow!(
385                    "Data type mismatch for column `{:?}`. BigQuery side: `{:?}`, RisingWave side: `{:?}`. ",
386                    i.name,
387                    value,
388                    data_type_string
389                )));
390            };
391        }
392        Ok(())
393    }
394
395    fn get_string_and_check_support_from_datatype(rw_data_type: &DataType) -> Result<String> {
396        match rw_data_type {
397            DataType::Boolean => Ok("BOOL".to_owned()),
398            DataType::Int16 => Ok("INT64".to_owned()),
399            DataType::Int32 => Ok("INT64".to_owned()),
400            DataType::Int64 => Ok("INT64".to_owned()),
401            DataType::Float32 => Err(SinkError::BigQuery(anyhow::anyhow!(
402                "REAL is not supported for BigQuery sink. Please convert to FLOAT64 or other supported types."
403            ))),
404            DataType::Float64 => Ok("FLOAT64".to_owned()),
405            DataType::Decimal => Ok("NUMERIC".to_owned()),
406            DataType::Date => Ok("DATE".to_owned()),
407            DataType::Varchar => Ok("STRING".to_owned()),
408            DataType::Time => Ok("TIME".to_owned()),
409            DataType::Timestamp => Ok("DATETIME".to_owned()),
410            DataType::Timestamptz => Ok("TIMESTAMP".to_owned()),
411            DataType::Interval => Ok("INTERVAL".to_owned()),
412            DataType::Struct(structs) => {
413                let mut elements_vec = vec![];
414                for (name, datatype) in structs.iter() {
415                    let element_string =
416                        Self::get_string_and_check_support_from_datatype(datatype)?;
417                    elements_vec.push(format!("{} {}", name, element_string));
418                }
419                Ok(format!("STRUCT<{}>", elements_vec.join(", ")))
420            }
421            DataType::List(l) => {
422                let element_string = Self::get_string_and_check_support_from_datatype(l.elem())?;
423                Ok(format!("ARRAY<{}>", element_string))
424            }
425            DataType::Bytea => Ok("BYTES".to_owned()),
426            DataType::Jsonb => Ok("JSON".to_owned()),
427            DataType::Variant => Err(SinkError::BigQuery(anyhow::anyhow!(
428                "VARIANT is not supported for BigQuery sink."
429            ))),
430            DataType::Serial => Ok("INT64".to_owned()),
431            DataType::Int256 => Err(SinkError::BigQuery(anyhow::anyhow!(
432                "INT256 is not supported for BigQuery sink."
433            ))),
434            DataType::Map(_) => Err(SinkError::BigQuery(anyhow::anyhow!(
435                "MAP is not supported for BigQuery sink."
436            ))),
437            DataType::Vector(_) => Err(SinkError::BigQuery(anyhow::anyhow!(
438                "VECTOR is not supported for BigQuery sink."
439            ))),
440        }
441    }
442
443    fn map_field(rw_field: &Field) -> Result<TableFieldSchema> {
444        let tfs = match &rw_field.data_type {
445            DataType::Boolean => TableFieldSchema::bool(&rw_field.name),
446            DataType::Int16 | DataType::Int32 | DataType::Int64 | DataType::Serial => {
447                TableFieldSchema::integer(&rw_field.name)
448            }
449            DataType::Float32 => {
450                return Err(SinkError::BigQuery(anyhow::anyhow!(
451                    "REAL is not supported for BigQuery sink. Please convert to FLOAT64 or other supported types."
452                )));
453            }
454            DataType::Float64 => TableFieldSchema::float(&rw_field.name),
455            DataType::Decimal => TableFieldSchema::numeric(&rw_field.name),
456            DataType::Date => TableFieldSchema::date(&rw_field.name),
457            DataType::Varchar => TableFieldSchema::string(&rw_field.name),
458            DataType::Time => TableFieldSchema::time(&rw_field.name),
459            DataType::Timestamp => TableFieldSchema::date_time(&rw_field.name),
460            DataType::Timestamptz => TableFieldSchema::timestamp(&rw_field.name),
461            DataType::Interval => {
462                return Err(SinkError::BigQuery(anyhow::anyhow!(
463                    "INTERVAL is not supported for BigQuery sink. Please convert to VARCHAR or other supported types."
464                )));
465            }
466            DataType::Struct(st) => {
467                let mut sub_fields = Vec::with_capacity(st.len());
468                for (name, dt) in st.iter() {
469                    let rw_field = Field::with_name(dt.clone(), name);
470                    let field = Self::map_field(&rw_field)?;
471                    sub_fields.push(field);
472                }
473                TableFieldSchema::record(&rw_field.name, sub_fields)
474            }
475            DataType::List(lt) => {
476                let inner_field =
477                    Self::map_field(&Field::with_name(lt.elem().clone(), &rw_field.name))?;
478                TableFieldSchema {
479                    mode: Some("REPEATED".to_owned()),
480                    ..inner_field
481                }
482            }
483
484            DataType::Bytea => TableFieldSchema::bytes(&rw_field.name),
485            DataType::Jsonb => TableFieldSchema::json(&rw_field.name),
486            DataType::Variant => {
487                return Err(SinkError::BigQuery(anyhow::anyhow!(
488                    "VARIANT is not supported for BigQuery sink."
489                )));
490            }
491            DataType::Int256 => {
492                return Err(SinkError::BigQuery(anyhow::anyhow!(
493                    "INT256 is not supported for BigQuery sink."
494                )));
495            }
496            DataType::Map(_) => {
497                return Err(SinkError::BigQuery(anyhow::anyhow!(
498                    "MAP is not supported for BigQuery sink."
499                )));
500            }
501            DataType::Vector(_) => {
502                return Err(SinkError::BigQuery(anyhow::anyhow!(
503                    "VECTOR is not supported for BigQuery sink."
504                )));
505            }
506        };
507        Ok(tfs)
508    }
509
510    async fn create_table(
511        &self,
512        client: &Client,
513        project_id: &str,
514        dataset_id: &str,
515        table_id: &str,
516        fields: &Vec<Field>,
517    ) -> Result<Table> {
518        let dataset = client
519            .dataset()
520            .get(project_id, dataset_id)
521            .await
522            .map_err(|e| SinkError::BigQuery(e.into()))?;
523        let fields: Vec<_> = fields.iter().map(Self::map_field).collect::<Result<_>>()?;
524        let table = Table::from_dataset(&dataset, table_id, TableSchema::new(fields));
525
526        client
527            .table()
528            .create(table)
529            .await
530            .map_err(|e| SinkError::BigQuery(e.into()))
531    }
532}
533
534impl Sink for BigQuerySink {
535    type LogSinker = BigQueryLogSinker;
536
537    const SINK_NAME: &'static str = BIGQUERY_SINK;
538
539    crate::impl_validate_sink_unknown_fields!();
540
541    async fn new_log_sinker(&self, _writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
542        let (writer, resp_stream) = BigQuerySinkWriter::new(
543            self.config.clone(),
544            self.schema.clone(),
545            self.pk_indices.clone(),
546            self.is_append_only,
547        )
548        .await?;
549        Ok(BigQueryLogSinker::new(
550            writer,
551            resp_stream,
552            BIGQUERY_SEND_FUTURE_BUFFER_MAX_SIZE,
553        ))
554    }
555
556    async fn validate(&self) -> Result<()> {
557        risingwave_common::license::Feature::BigQuerySink
558            .check_available()
559            .map_err(|e| anyhow::anyhow!(e))?;
560        if !self.is_append_only && self.pk_indices.is_empty() {
561            return Err(SinkError::Config(anyhow!(
562                "Primary key not defined for upsert bigquery sink (please define in `primary_key` field)"
563            )));
564        }
565        let client = self
566            .config
567            .common
568            .build_client(&self.config.aws_auth_props)
569            .await?;
570        let BigQueryCommon {
571            project: project_id,
572            dataset: dataset_id,
573            table: table_id,
574            ..
575        } = &self.config.common;
576
577        if self.config.common.auto_create {
578            match client
579                .table()
580                .get(project_id, dataset_id, table_id, None)
581                .await
582            {
583                Err(BQError::ResponseError { error }) if error.error.code == 404 => {
584                    // early return: no need to query schema to check column and type
585                    return self
586                        .create_table(
587                            &client,
588                            project_id,
589                            dataset_id,
590                            table_id,
591                            &self.schema.fields,
592                        )
593                        .await
594                        .map(|_| ());
595                }
596                Err(e) => return Err(SinkError::BigQuery(e.into())),
597                _ => {}
598            }
599        }
600
601        let rs = client
602            .job()
603            .query(
604                &self.config.common.project,
605                QueryRequest::new(format!(
606                    "SELECT column_name, data_type FROM `{}.{}.INFORMATION_SCHEMA.COLUMNS` WHERE table_name = '{}'",
607                    project_id, dataset_id, table_id,
608                )),
609            ).await.map_err(|e| SinkError::BigQuery(e.into()))?;
610        let mut rs = ResultSet::new_from_query_response(rs);
611
612        let mut big_query_schema = HashMap::default();
613        while rs.next_row() {
614            big_query_schema.insert(
615                rs.get_string_by_name("column_name")
616                    .map_err(|e| SinkError::BigQuery(e.into()))?
617                    .ok_or_else(|| {
618                        SinkError::BigQuery(anyhow::anyhow!("Cannot find column_name"))
619                    })?,
620                rs.get_string_by_name("data_type")
621                    .map_err(|e| SinkError::BigQuery(e.into()))?
622                    .ok_or_else(|| {
623                        SinkError::BigQuery(anyhow::anyhow!("Cannot find column_name"))
624                    })?,
625            );
626        }
627
628        self.check_column_name_and_type(big_query_schema)?;
629        Ok(())
630    }
631}
632
633pub struct BigQuerySinkWriter {
634    pub config: BigQueryConfig,
635    #[expect(dead_code)]
636    schema: Schema,
637    #[expect(dead_code)]
638    pk_indices: Vec<usize>,
639    client: StorageWriterClient,
640    is_append_only: bool,
641    row_encoder: ProtoEncoder,
642    writer_pb_schema: ProtoSchema,
643    #[expect(dead_code)]
644    message_descriptor: MessageDescriptor,
645    write_stream: String,
646    proto_field: Option<FieldDescriptor>,
647    change_sequence_number_field: Option<FieldDescriptor>,
648}
649
650impl TryFrom<SinkParam> for BigQuerySink {
651    type Error = SinkError;
652
653    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
654        let schema = param.schema();
655        let pk_indices = param.downstream_pk_or_empty();
656        let config = BigQueryConfig::from_btreemap(param.properties)?;
657        BigQuerySink::new(config, schema, pk_indices, param.sink_type.is_append_only())
658    }
659}
660
661impl BigQuerySinkWriter {
662    pub async fn new(
663        config: BigQueryConfig,
664        schema: Schema,
665        pk_indices: Vec<usize>,
666        is_append_only: bool,
667    ) -> Result<(Self, impl Stream<Item = Result<()>>)> {
668        let (client, resp_stream) = config
669            .common
670            .build_writer_client(&config.aws_auth_props)
671            .await?;
672        let writer = Self::with_client(config, schema, pk_indices, is_append_only, client)?;
673        Ok((writer, resp_stream))
674    }
675
676    fn with_client(
677        config: BigQueryConfig,
678        schema: Schema,
679        pk_indices: Vec<usize>,
680        is_append_only: bool,
681        client: StorageWriterClient,
682    ) -> Result<Self> {
683        let mut descriptor_proto = build_protobuf_schema(
684            schema
685                .fields()
686                .iter()
687                .map(|f| (f.name.as_str(), &f.data_type)),
688            config.common.table.clone(),
689        )?;
690
691        if !is_append_only {
692            let field = FieldDescriptorProto {
693                name: Some(CHANGE_TYPE.to_owned()),
694                number: Some((schema.len() + 1) as i32),
695                r#type: Some(field_descriptor_proto::Type::String.into()),
696                ..Default::default()
697            };
698            descriptor_proto.field.push(field);
699            let field = FieldDescriptorProto {
700                name: Some(CHANGE_SEQUENCE_NUMBER.to_owned()),
701                number: Some((schema.len() + 2) as i32),
702                r#type: Some(field_descriptor_proto::Type::String.into()),
703                ..Default::default()
704            };
705            descriptor_proto.field.push(field);
706        }
707
708        let descriptor_pool = build_protobuf_descriptor_pool(&descriptor_proto)?;
709        let message_descriptor = descriptor_pool
710            .get_message_by_name(&config.common.table)
711            .ok_or_else(|| {
712                SinkError::BigQuery(anyhow::anyhow!(
713                    "Can't find message proto {}",
714                    config.common.table
715                ))
716            })?;
717        let proto_field = if !is_append_only {
718            let proto_field = message_descriptor
719                .get_field_by_name(CHANGE_TYPE)
720                .ok_or_else(|| {
721                    SinkError::BigQuery(anyhow::anyhow!("Can't find {}", CHANGE_TYPE))
722                })?;
723            Some(proto_field)
724        } else {
725            None
726        };
727        let change_sequence_number_field = if !is_append_only {
728            let field = message_descriptor
729                .get_field_by_name(CHANGE_SEQUENCE_NUMBER)
730                .ok_or_else(|| {
731                    SinkError::BigQuery(anyhow::anyhow!("Can't find {}", CHANGE_SEQUENCE_NUMBER))
732                })?;
733            Some(field)
734        } else {
735            None
736        };
737        let row_encoder = ProtoEncoder::new(
738            schema.clone(),
739            None,
740            message_descriptor.clone(),
741            ProtoHeader::None,
742        )?;
743        Ok(Self {
744            write_stream: format!(
745                "projects/{}/datasets/{}/tables/{}/streams/_default",
746                config.common.project, config.common.dataset, config.common.table
747            ),
748            config,
749            schema,
750            pk_indices,
751            client,
752            is_append_only,
753            row_encoder,
754            message_descriptor,
755            proto_field,
756            change_sequence_number_field,
757            writer_pb_schema: ProtoSchema {
758                proto_descriptor: Some(descriptor_proto.clone()),
759            },
760        })
761    }
762
763    fn append_only(&mut self, chunk: StreamChunk) -> Result<Vec<Vec<u8>>> {
764        let mut serialized_rows: Vec<Vec<u8>> = Vec::with_capacity(chunk.capacity());
765        for (op, row) in chunk.rows() {
766            if op != Op::Insert {
767                continue;
768            }
769            serialized_rows.push(self.row_encoder.encode(row)?.ser_to()?)
770        }
771        Ok(serialized_rows)
772    }
773
774    fn upsert(&mut self, chunk: StreamChunk) -> Result<Vec<Vec<u8>>> {
775        let sequence_numbers = CHANGE_SEQUENCE_CLOCK.reserve(chunk.cardinality());
776        let mut serialized_rows: Vec<Vec<u8>> = Vec::with_capacity(chunk.capacity());
777        for ((op, row), sequence_number) in chunk.rows().zip_eq_debug(sequence_numbers) {
778            if op == Op::UpdateDelete {
779                continue;
780            }
781            let mut pb_row = self.row_encoder.encode(row)?;
782            match op {
783                Op::Insert => pb_row
784                    .message
785                    .try_set_field(
786                        self.proto_field.as_ref().unwrap(),
787                        prost_reflect::Value::String("UPSERT".to_owned()),
788                    )
789                    .map_err(|e| SinkError::BigQuery(e.into()))?,
790                Op::Delete => pb_row
791                    .message
792                    .try_set_field(
793                        self.proto_field.as_ref().unwrap(),
794                        prost_reflect::Value::String("DELETE".to_owned()),
795                    )
796                    .map_err(|e| SinkError::BigQuery(e.into()))?,
797                Op::UpdateDelete => continue,
798                Op::UpdateInsert => pb_row
799                    .message
800                    .try_set_field(
801                        self.proto_field.as_ref().unwrap(),
802                        prost_reflect::Value::String("UPSERT".to_owned()),
803                    )
804                    .map_err(|e| SinkError::BigQuery(e.into()))?,
805            };
806            pb_row
807                .message
808                .try_set_field(
809                    self.change_sequence_number_field.as_ref().unwrap(),
810                    prost_reflect::Value::String(sequence_number),
811                )
812                .map_err(|e| SinkError::BigQuery(e.into()))?;
813
814            serialized_rows.push(pb_row.ser_to()?)
815        }
816        Ok(serialized_rows)
817    }
818
819    fn write_chunk(&mut self, chunk: StreamChunk) -> Result<usize> {
820        let serialized_rows = if self.is_append_only {
821            self.append_only(chunk)?
822        } else {
823            self.upsert(chunk)?
824        };
825        if serialized_rows.is_empty() {
826            return Ok(0);
827        }
828        let mut result = Vec::new();
829        let mut result_inner = Vec::new();
830        let mut size_count = 0;
831        for i in serialized_rows {
832            size_count += i.len();
833            if size_count > MAX_ROW_SIZE {
834                result.push(result_inner);
835                result_inner = Vec::new();
836                size_count = i.len();
837            }
838            result_inner.push(i);
839        }
840        if !result_inner.is_empty() {
841            result.push(result_inner);
842        }
843        let len = result.len();
844        for serialized_rows in result {
845            let rows = AppendRowsRequestRows::ProtoRows(ProtoData {
846                writer_schema: Some(self.writer_pb_schema.clone()),
847                rows: Some(ProtoRows { serialized_rows }),
848            });
849            self.client.append_rows(rows, self.write_stream.clone())?;
850        }
851        Ok(len)
852    }
853}
854
855/// Hands out `_CHANGE_SEQUENCE_NUMBER`s as `{physical:X}/{logical:X}`, with the wall clock in
856/// microseconds as `physical`. Each number is greater than any reserved before from the same clock,
857/// even if the wall clock moves back.
858///
859/// The change with the largest number wins for each key, so numbering rows in send order is
860/// enough: the sinker sends the changes of a key in log order, and every replay (rewind, recovery,
861/// rescale) resends all changes after the truncated offset, so the last change sent for a key is
862/// its latest one.
863struct ChangeSequenceClock {
864    /// The physical time of the last reserved numbers and the next logical number after them.
865    next: Mutex<(u64, u64)>,
866}
867
868impl ChangeSequenceClock {
869    const fn new() -> Self {
870        Self {
871            next: parking_lot::const_mutex((0, 0)),
872        }
873    }
874
875    /// Reserves `count` numbers at the current wall clock.
876    fn reserve(&self, count: usize) -> impl Iterator<Item = String> + use<> {
877        let now = SystemTime::now()
878            .duration_since(UNIX_EPOCH)
879            .map_or(0, |since_epoch| since_epoch.as_micros() as u64);
880        let count = count as u64;
881        let (physical, start) = self.reserve_at(now, count);
882        (start..start + count).map(move |logical| format!("{physical:X}/{logical:X}"))
883    }
884
885    /// Reserves `count` logical numbers at physical time `now`, or after the last reserved ones if
886    /// `now` has not moved past their physical time. Returns the physical time and the first
887    /// logical number.
888    fn reserve_at(&self, now: u64, count: u64) -> (u64, u64) {
889        let mut next = self.next.lock();
890        let (physical, logical) = if now > next.0 { (now, 0) } else { *next };
891        *next = (physical, logical + count);
892        (physical, logical)
893    }
894}
895
896/// Shared by all writers in the process, so that the numbers keep increasing when a writer is
897/// rebuilt or a key moves to another actor in the process.
898static CHANGE_SEQUENCE_CLOCK: ChangeSequenceClock = ChangeSequenceClock::new();
899
900#[try_stream(ok = (), error = SinkError)]
901pub async fn resp_to_stream(
902    resp_stream: impl Future<
903        Output = std::result::Result<
904            Response<google_cloud_gax::grpc::Streaming<AppendRowsResponse>>,
905            Status,
906        >,
907    >
908    + 'static
909    + Send,
910) {
911    let mut resp_stream = resp_stream
912        .await
913        .map_err(|e| SinkError::BigQuery(e.into()))?
914        .into_inner();
915    loop {
916        match resp_stream
917            .message()
918            .await
919            .map_err(|e| SinkError::BigQuery(e.into()))?
920        {
921            Some(append_rows_response) => {
922                if !append_rows_response.row_errors.is_empty() {
923                    return Err(SinkError::BigQuery(anyhow::anyhow!(
924                        "bigquery insert error {:?}",
925                        append_rows_response.row_errors
926                    )));
927                }
928                if let Some(google_cloud_googleapis::cloud::bigquery::storage::v1::append_rows_response::Response::Error(status)) = append_rows_response.response{
929                            return Err(SinkError::BigQuery(anyhow::anyhow!(
930                                "bigquery insert error {:?}",
931                                status
932                            )));
933                        }
934                yield ();
935            }
936            None => {
937                return Err(SinkError::BigQuery(anyhow::anyhow!(
938                    "bigquery insert error: end of resp stream",
939                )));
940            }
941        }
942    }
943}
944
945struct StorageWriterClient {
946    #[expect(dead_code)]
947    environment: Environment,
948    request_sender: mpsc::UnboundedSender<AppendRowsRequest>,
949}
950impl StorageWriterClient {
951    pub async fn new(
952        credentials: CredentialsFile,
953    ) -> Result<(Self, impl Stream<Item = Result<()>>)> {
954        let ts_grpc = google_cloud_auth::token::DefaultTokenSourceProvider::new_with_credentials(
955            Self::bigquery_grpc_auth_config(),
956            Box::new(credentials),
957        )
958        .await
959        .map_err(|e| SinkError::BigQuery(e.into()))?;
960        let conn_options = ConnectionOptions {
961            connect_timeout: CONNECT_TIMEOUT,
962            timeout: CONNECTION_TIMEOUT,
963            ..Default::default()
964        };
965        let environment = Environment::GoogleCloud(Box::new(ts_grpc));
966        let conn = ConnectionManager::new(DEFAULT_GRPC_CHANNEL_NUMS, &environment, &conn_options)
967            .await
968            .map_err(|e| SinkError::BigQuery(e.into()))?;
969        let mut client = conn.writer();
970
971        let (tx, rx) = mpsc::unbounded_channel();
972        let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx);
973
974        let resp = async move { client.append_rows(Request::new(stream)).await };
975        let resp_stream = resp_to_stream(resp);
976
977        Ok((
978            StorageWriterClient {
979                environment,
980                request_sender: tx,
981            },
982            resp_stream,
983        ))
984    }
985
986    pub fn append_rows(&mut self, row: AppendRowsRequestRows, write_stream: String) -> Result<()> {
987        let append_req = AppendRowsRequest {
988            write_stream,
989            offset: None,
990            trace_id: Uuid::new_v4().hyphenated().to_string(),
991            missing_value_interpretations: HashMap::default(),
992            rows: Some(row),
993            default_missing_value_interpretation: MissingValueInterpretation::DefaultValue as i32,
994        };
995        self.request_sender
996            .send(append_req)
997            .map_err(|e| SinkError::BigQuery(e.into()))?;
998        Ok(())
999    }
1000
1001    fn bigquery_grpc_auth_config() -> google_cloud_auth::project::Config<'static> {
1002        let mut auth_config = google_cloud_auth::project::Config::default();
1003        auth_config =
1004            auth_config.with_audience(google_cloud_bigquery::grpc::apiv1::conn_pool::AUDIENCE);
1005        auth_config =
1006            auth_config.with_scopes(&google_cloud_bigquery::grpc::apiv1::conn_pool::SCOPES);
1007        auth_config
1008    }
1009}
1010
1011fn build_protobuf_descriptor_pool(desc: &DescriptorProto) -> Result<prost_reflect::DescriptorPool> {
1012    let file_descriptor = FileDescriptorProto {
1013        message_type: vec![desc.clone()],
1014        name: Some("bigquery".to_owned()),
1015        ..Default::default()
1016    };
1017
1018    prost_reflect::DescriptorPool::from_file_descriptor_set(FileDescriptorSet {
1019        file: vec![file_descriptor],
1020    })
1021    .context("failed to build descriptor pool")
1022    .map_err(SinkError::BigQuery)
1023}
1024
1025fn build_protobuf_schema<'a>(
1026    fields: impl Iterator<Item = (&'a str, &'a DataType)>,
1027    name: String,
1028) -> Result<DescriptorProto> {
1029    let mut proto = DescriptorProto {
1030        name: Some(name),
1031        ..Default::default()
1032    };
1033    let mut struct_vec = vec![];
1034    let field_vec = fields
1035        .enumerate()
1036        .map(|(index, (name, data_type))| {
1037            let (field, des_proto) =
1038                build_protobuf_field(data_type, (index + 1) as i32, name.to_owned())?;
1039            if let Some(sv) = des_proto {
1040                struct_vec.push(sv);
1041            }
1042            Ok(field)
1043        })
1044        .collect::<Result<Vec<_>>>()?;
1045    proto.field = field_vec;
1046    proto.nested_type = struct_vec;
1047    Ok(proto)
1048}
1049
1050fn build_protobuf_field(
1051    data_type: &DataType,
1052    index: i32,
1053    name: String,
1054) -> Result<(FieldDescriptorProto, Option<DescriptorProto>)> {
1055    let mut field = FieldDescriptorProto {
1056        name: Some(name.clone()),
1057        number: Some(index),
1058        ..Default::default()
1059    };
1060    match data_type {
1061        DataType::Boolean => field.r#type = Some(field_descriptor_proto::Type::Bool.into()),
1062        DataType::Int32 => field.r#type = Some(field_descriptor_proto::Type::Int32.into()),
1063        DataType::Int16 | DataType::Int64 => {
1064            field.r#type = Some(field_descriptor_proto::Type::Int64.into())
1065        }
1066        DataType::Float64 => field.r#type = Some(field_descriptor_proto::Type::Double.into()),
1067        DataType::Decimal => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1068        DataType::Date => field.r#type = Some(field_descriptor_proto::Type::Int32.into()),
1069        DataType::Varchar => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1070        DataType::Time => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1071        DataType::Timestamp => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1072        DataType::Timestamptz => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1073        DataType::Interval => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1074        DataType::Struct(s) => {
1075            field.r#type = Some(field_descriptor_proto::Type::Message.into());
1076            let name = format!("Struct{}", name);
1077            let sub_proto = build_protobuf_schema(s.iter(), name.clone())?;
1078            field.type_name = Some(name);
1079            return Ok((field, Some(sub_proto)));
1080        }
1081        DataType::List(l) => {
1082            let (mut field, proto) = build_protobuf_field(l.elem(), index, name)?;
1083            field.label = Some(field_descriptor_proto::Label::Repeated.into());
1084            return Ok((field, proto));
1085        }
1086        DataType::Bytea => field.r#type = Some(field_descriptor_proto::Type::Bytes.into()),
1087        DataType::Jsonb => field.r#type = Some(field_descriptor_proto::Type::String.into()),
1088        DataType::Variant => {
1089            return Err(SinkError::BigQuery(anyhow::anyhow!("Don't support Variant")));
1090        }
1091        DataType::Serial => field.r#type = Some(field_descriptor_proto::Type::Int64.into()),
1092        DataType::Float32 | DataType::Int256 => {
1093            return Err(SinkError::BigQuery(anyhow::anyhow!(
1094                "Don't support Float32 and Int256"
1095            )));
1096        }
1097        DataType::Map(_) => return Err(SinkError::BigQuery(anyhow::anyhow!("Don't support Map"))),
1098        DataType::Vector(_) => {
1099            return Err(SinkError::BigQuery(anyhow::anyhow!("Don't support Vector")));
1100        }
1101    }
1102    Ok((field, None))
1103}
1104
1105#[cfg(test)]
1106mod test {
1107
1108    use std::assert_matches;
1109    use std::collections::{BTreeMap, HashMap};
1110
1111    use google_cloud_gax::conn::Environment;
1112    use google_cloud_googleapis::cloud::bigquery::storage::v1::append_rows_request::{
1113        ProtoData, Rows as AppendRowsRequestRows,
1114    };
1115    use prost_reflect::{DynamicMessage, ReflectMessage};
1116    use risingwave_common::array::StreamChunk;
1117    use risingwave_common::array::stream_chunk::StreamChunkTestExt as _;
1118    use risingwave_common::catalog::{Field, Schema};
1119    use risingwave_common::types::{DataType, StructType};
1120    use tokio::sync::mpsc;
1121
1122    use crate::connector_common::AwsAuthProps;
1123    use crate::sink::big_query::{
1124        BigQueryCommon, BigQueryConfig, BigQuerySink, BigQuerySinkWriter, CHANGE_SEQUENCE_NUMBER,
1125        CHANGE_TYPE, ChangeSequenceClock, StorageWriterClient, build_protobuf_descriptor_pool,
1126        build_protobuf_schema,
1127    };
1128
1129    #[tokio::test]
1130    async fn test_type_check() {
1131        let big_query_type_string = "ARRAY<STRUCT<v1 ARRAY<INT64>, v2 STRUCT<v1 INT64, v2 INT64>>>";
1132        let rw_datatype = DataType::list(DataType::Struct(StructType::new(vec![
1133            ("v1".to_owned(), DataType::Int64.list()),
1134            (
1135                "v2".to_owned(),
1136                DataType::Struct(StructType::new(vec![
1137                    ("v1".to_owned(), DataType::Int64),
1138                    ("v2".to_owned(), DataType::Int64),
1139                ])),
1140            ),
1141        ])));
1142        assert_eq!(
1143            BigQuerySink::get_string_and_check_support_from_datatype(&rw_datatype).unwrap(),
1144            big_query_type_string
1145        );
1146    }
1147
1148    #[tokio::test]
1149    async fn test_schema_check() {
1150        let schema = Schema {
1151            fields: vec![
1152                Field::with_name(DataType::Int64, "v1"),
1153                Field::with_name(DataType::Float64, "v2"),
1154                Field::with_name(
1155                    DataType::list(DataType::Struct(StructType::new(vec![
1156                        ("v1".to_owned(), DataType::Int64.list()),
1157                        (
1158                            "v3".to_owned(),
1159                            DataType::Struct(StructType::new(vec![
1160                                ("v1".to_owned(), DataType::Int64),
1161                                ("v2".to_owned(), DataType::Int64),
1162                            ])),
1163                        ),
1164                    ]))),
1165                    "v3",
1166                ),
1167            ],
1168        };
1169        let fields = schema
1170            .fields()
1171            .iter()
1172            .map(|f| (f.name.as_str(), &f.data_type));
1173        let desc = build_protobuf_schema(fields, "t1".to_owned()).unwrap();
1174        let pool = build_protobuf_descriptor_pool(&desc).unwrap();
1175        let t1_message = pool.get_message_by_name("t1").unwrap();
1176        assert_matches!(
1177            t1_message.get_field_by_name("v1").unwrap().kind(),
1178            prost_reflect::Kind::Int64
1179        );
1180        assert_matches!(
1181            t1_message.get_field_by_name("v2").unwrap().kind(),
1182            prost_reflect::Kind::Double
1183        );
1184        assert_matches!(
1185            t1_message.get_field_by_name("v3").unwrap().kind(),
1186            prost_reflect::Kind::Message(_)
1187        );
1188
1189        let v3_message = pool.get_message_by_name("t1.Structv3").unwrap();
1190        assert_matches!(
1191            v3_message.get_field_by_name("v1").unwrap().kind(),
1192            prost_reflect::Kind::Int64
1193        );
1194        assert!(v3_message.get_field_by_name("v1").unwrap().is_list());
1195
1196        let v3_v3_message = pool.get_message_by_name("t1.Structv3.Structv3").unwrap();
1197        assert_matches!(
1198            v3_v3_message.get_field_by_name("v1").unwrap().kind(),
1199            prost_reflect::Kind::Int64
1200        );
1201        assert_matches!(
1202            v3_v3_message.get_field_by_name("v2").unwrap().kind(),
1203            prost_reflect::Kind::Int64
1204        );
1205    }
1206
1207    #[test]
1208    fn test_decimal_type_family_compatibility() {
1209        assert!(BigQuerySink::is_decimal_type_compatible("NUMERIC"));
1210        assert!(BigQuerySink::is_decimal_type_compatible("numeric(31, 2)"));
1211        assert!(BigQuerySink::is_decimal_type_compatible("BIGNUMERIC"));
1212        assert!(BigQuerySink::is_decimal_type_compatible(
1213            "bignumeric(35, 12)"
1214        ));
1215        assert!(!BigQuerySink::is_decimal_type_compatible("STRING"));
1216    }
1217
1218    #[test]
1219    fn test_decimal_schema_check_accepts_parameterized_numeric_types() {
1220        let sink = BigQuerySink {
1221            config: BigQueryConfig {
1222                common: BigQueryCommon {
1223                    local_path: None,
1224                    s3_path: None,
1225                    project: "project".to_owned(),
1226                    dataset: "dataset".to_owned(),
1227                    table: "table".to_owned(),
1228                    auto_create: false,
1229                    credentials: None,
1230                },
1231                aws_auth_props: AwsAuthProps {
1232                    region: None,
1233                    endpoint: None,
1234                    access_key: None,
1235                    secret_key: None,
1236                    session_token: None,
1237                    arn: None,
1238                    external_id: None,
1239                    profile: None,
1240                    msk_signer_timeout_sec: None,
1241                },
1242                r#type: "append-only".to_owned(),
1243                unknown_fields: Default::default(),
1244            },
1245            schema: Schema {
1246                fields: vec![Field::with_name(DataType::Decimal, "capitalizedcost")],
1247            },
1248            pk_indices: vec![],
1249            is_append_only: true,
1250        };
1251
1252        sink.check_column_name_and_type(HashMap::from([(
1253            "capitalizedcost".to_owned(),
1254            "NUMERIC(31, 2)".to_owned(),
1255        )]))
1256        .unwrap();
1257
1258        sink.check_column_name_and_type(HashMap::from([(
1259            "capitalizedcost".to_owned(),
1260            "BIGNUMERIC(35, 12)".to_owned(),
1261        )]))
1262        .unwrap();
1263    }
1264
1265    fn sink_config(sink_type: &str) -> BigQueryConfig {
1266        BigQueryConfig::from_btreemap(BTreeMap::from(
1267            [
1268                ("bigquery.project", "project"),
1269                ("bigquery.dataset", "dataset"),
1270                ("bigquery.table", "table"),
1271                ("type", sink_type),
1272            ]
1273            .map(|(k, v)| (k.to_owned(), v.to_owned())),
1274        ))
1275        .unwrap()
1276    }
1277
1278    /// Writes `chunks` in order and decodes the rows of the sent requests.
1279    fn write_and_decode(sink_type: &str, chunks: Vec<StreamChunk>) -> Vec<DynamicMessage> {
1280        let schema = Schema {
1281            fields: vec![
1282                Field::with_name(DataType::Int64, "id"),
1283                Field::with_name(DataType::Varchar, "v"),
1284            ],
1285        };
1286        let (request_sender, mut request_rx) = mpsc::unbounded_channel();
1287        let client = StorageWriterClient {
1288            environment: Environment::Emulator(String::new()),
1289            request_sender,
1290        };
1291        let mut writer = BigQuerySinkWriter::with_client(
1292            sink_config(sink_type),
1293            schema,
1294            vec![0],
1295            sink_type == "append-only",
1296            client,
1297        )
1298        .unwrap();
1299        for chunk in chunks {
1300            writer.write_chunk(chunk).unwrap();
1301        }
1302
1303        let mut rows = vec![];
1304        while let Ok(request) = request_rx.try_recv() {
1305            let Some(AppendRowsRequestRows::ProtoRows(ProtoData {
1306                writer_schema: Some(writer_schema),
1307                rows: Some(proto_rows),
1308            })) = request.rows
1309            else {
1310                panic!("expected proto rows");
1311            };
1312            let pool =
1313                build_protobuf_descriptor_pool(&writer_schema.proto_descriptor.unwrap()).unwrap();
1314            let message = pool.get_message_by_name("table").unwrap();
1315            rows.extend(
1316                proto_rows
1317                    .serialized_rows
1318                    .iter()
1319                    .map(|row| DynamicMessage::decode(message.clone(), row.as_slice()).unwrap()),
1320            );
1321        }
1322        rows
1323    }
1324
1325    #[test]
1326    fn test_change_sequence_clock() {
1327        let clock = ChangeSequenceClock::new();
1328        assert_eq!(clock.reserve_at(100, 3), (100, 0));
1329        // The clock has not moved past the last physical time, or has moved back.
1330        assert_eq!(clock.reserve_at(100, 2), (100, 3));
1331        assert_eq!(clock.reserve_at(90, 1), (100, 5));
1332        assert_eq!(clock.reserve_at(101, 1), (101, 0));
1333    }
1334
1335    #[test]
1336    fn test_write_change_sequence_number() {
1337        let chunks = || {
1338            vec![
1339                StreamChunk::from_pretty(
1340                    "  I T
1341                    +  1 a
1342                    U- 1 a
1343                    U+ 1 b",
1344                ),
1345                StreamChunk::from_pretty(
1346                    " I T
1347                    - 1 b",
1348                ),
1349            ]
1350        };
1351        let field = |row: &DynamicMessage, name: &str| {
1352            row.get_field_by_name(name)
1353                .unwrap()
1354                .as_str()
1355                .unwrap()
1356                .to_owned()
1357        };
1358
1359        let rows = write_and_decode("upsert", chunks());
1360        let change_types: Vec<_> = rows.iter().map(|row| field(row, CHANGE_TYPE)).collect();
1361        assert_eq!(change_types, ["UPSERT", "UPSERT", "DELETE"]);
1362        let sequence_numbers: Vec<_> = rows
1363            .iter()
1364            .map(|row| {
1365                let sequence_number = field(row, CHANGE_SEQUENCE_NUMBER);
1366                let (physical, logical) = sequence_number.split_once('/').unwrap();
1367                (
1368                    u64::from_str_radix(physical, 16).unwrap(),
1369                    u64::from_str_radix(logical, 16).unwrap(),
1370                )
1371            })
1372            .collect();
1373        assert!(
1374            sequence_numbers.windows(2).all(|pair| pair[0] < pair[1]),
1375            "{sequence_numbers:?}"
1376        );
1377
1378        let rows = write_and_decode("append-only", chunks());
1379        for name in [CHANGE_TYPE, CHANGE_SEQUENCE_NUMBER] {
1380            assert!(rows[0].descriptor().get_field_by_name(name).is_none());
1381        }
1382    }
1383}