1use 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;
83const 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")] #[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: 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, #[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 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 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
855struct ChangeSequenceClock {
864 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 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 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
896static 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 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 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}