diff --git a/tests/bigquery/src/lib.rs b/tests/bigquery/src/lib.rs index 04d851ff4f..29a4b1f59d 100644 --- a/tests/bigquery/src/lib.rs +++ b/tests/bigquery/src/lib.rs @@ -28,6 +28,6 @@ pub use query::{ }; pub use reads::read_rows; pub use reads::run_reads; -pub use writes::run_writes; +pub use writes::{run_writes, run_writes_flaky}; pub use bigquery_samples::{INSTANCE_LABEL, random_id_suffix}; diff --git a/tests/bigquery/src/writes.rs b/tests/bigquery/src/writes.rs index c78ae34c16..c3ea8c63de 100644 --- a/tests/bigquery/src/writes.rs +++ b/tests/bigquery/src/writes.rs @@ -13,6 +13,7 @@ // limitations under the License. mod arrow; +mod flaky; use anyhow::Result; use bigquery_samples::{ @@ -21,7 +22,7 @@ use bigquery_samples::{ use google_cloud_bigquery::client::{BigQuery, Write}; use google_cloud_bigquery::query::FromRow; use google_cloud_bigquery_v2::client::{DatasetService, TableService}; -use google_cloud_bigquery_v2::model::{TableFieldSchema, TableSchema}; +use google_cloud_bigquery_v2::model::{Dataset, DatasetReference, TableFieldSchema, TableSchema}; use google_cloud_test_utils::runtime_config::project_id; pub async fn run_writes() -> Result<()> { @@ -75,6 +76,82 @@ pub async fn run_writes() -> Result<()> { result } +pub async fn run_writes_flaky() -> Result<()> { + let project_id = project_id()?; + let dataset_service = DatasetService::builder().with_tracing().build().await?; + cleanup_stale_datasets(&dataset_service, &project_id).await?; + + let dataset_id = format!("rust_bq_flaky_{}", bigquery_samples::random_id_suffix()); + let _ = dataset_service + .insert_dataset() + .set_project_id(&project_id) + .set_dataset( + Dataset::new() + .set_dataset_reference(DatasetReference::new().set_dataset_id(&dataset_id)) + .set_location(flaky::FLAKY_REGION) + .set_labels([(bigquery_samples::INSTANCE_LABEL, "true")]), + ) + .send() + .await?; + + let table_service = TableService::builder().with_tracing().build().await?; + let schema = TableSchema::new().set_fields([ + TableFieldSchema::new().set_name("name").set_type("STRING"), + TableFieldSchema::new().set_name("age").set_type("INTEGER"), + TableFieldSchema::new().set_name("test").set_type("STRING"), + ]); + + let result = async { + let client = Write::builder().build().await?; + flaky::reconnect_on_close_sequential( + &client, + &table_service, + &project_id, + &dataset_id, + schema.clone(), + ) + .await?; + flaky::reconnect_on_close_parallel( + &client, + &table_service, + &project_id, + &dataset_id, + schema.clone(), + ) + .await?; + flaky::initial_connect_failure_sequential( + &client, + &table_service, + &project_id, + &dataset_id, + schema.clone(), + ) + .await?; + flaky::initial_connect_failure_parallel( + &client, + &table_service, + &project_id, + &dataset_id, + schema.clone(), + ) + .await?; + flaky::reconnect_on_close_default( + &client, + &table_service, + &project_id, + &dataset_id, + schema, + ) + .await?; + + Ok(()) + } + .await; + + let _ = delete_dataset(&dataset_service, &project_id, &dataset_id).await; + result +} + #[derive(FromRow, Debug, PartialEq)] pub(crate) struct WriteUserRecord { pub(crate) name: String, @@ -82,6 +159,11 @@ pub(crate) struct WriteUserRecord { pub(crate) test: String, } +#[derive(FromRow, Debug, PartialEq)] +pub(crate) struct WriteCountRecord { + pub(crate) count: i64, +} + pub(crate) async fn read_writes_table( project_id: &str, dataset_id: &str, @@ -106,3 +188,30 @@ pub(crate) async fn read_writes_table( } Ok(users) } + +pub(crate) async fn count_writes_table( + project_id: &str, + dataset_id: &str, + table_id: &str, + test_filter: &str, + location: Option<&str>, +) -> Result { + let client = BigQuery::builder().build().await?; + let query = format!( + "SELECT COUNT(*) as count FROM `{project_id}.{dataset_id}.{table_id}` WHERE test = '{test_filter}'" + ); + let mut builder = client + .query(query) + .with_project_id(project_id) + .set_labels(vec![(bigquery_samples::INSTANCE_LABEL, "true")]); + if let Some(loc) = location { + builder = builder.set_location(loc); + } + let mut rows = builder.until_done().await?.read(); + + if let Some(row) = rows.next().await { + let count_row: WriteCountRecord = row?.try_into()?; + return Ok(count_row.count); + } + Ok(0) +} diff --git a/tests/bigquery/src/writes/arrow.rs b/tests/bigquery/src/writes/arrow.rs index d20311632d..18b30c41aa 100644 --- a/tests/bigquery/src/writes/arrow.rs +++ b/tests/bigquery/src/writes/arrow.rs @@ -398,14 +398,14 @@ pub async fn multiplex( Ok(()) } -struct ArrowSerializer { +pub(crate) struct ArrowSerializer { schema: Arc, writer: StreamWriter>, test: &'static str, } impl ArrowSerializer { - fn new(test: &'static str) -> Result { + pub(crate) fn new(test: &'static str) -> Result { let schema = Arc::new(Schema::new(vec![ Field::new("name", DataType::Utf8, false), Field::new("age", DataType::Int64, false), @@ -419,12 +419,12 @@ impl ArrowSerializer { }) } - fn schema(&mut self) -> ArrowSchema { + pub(crate) fn schema(&mut self) -> ArrowSchema { let buf = std::mem::take(self.writer.get_mut()); ArrowSchema::new().set_serialized_schema(buf) } - fn batch(&mut self, names: Vec<&str>, ages: Vec) -> Result { + pub(crate) fn batch(&mut self, names: Vec<&str>, ages: Vec) -> Result { let batch = { let name = StringArray::from(names); let age = Int64Array::from(ages); @@ -438,4 +438,17 @@ impl ArrowSerializer { let buf = std::mem::take(self.writer.get_mut()); Ok(ArrowRecordBatch::new().set_serialized_record_batch(buf)) } + + pub(crate) fn generate_batch( + &mut self, + count: usize, + start_row: usize, + ) -> Result { + let names: Vec = (0..count) + .map(|i| format!("user_{}", start_row + i)) + .collect(); + let name_slices: Vec<&str> = names.iter().map(|s| s.as_str()).collect(); + let ages: Vec = (0..count).map(|i| (start_row + i) as i64).collect(); + self.batch(name_slices, ages) + } } diff --git a/tests/bigquery/src/writes/flaky.rs b/tests/bigquery/src/writes/flaky.rs new file mode 100644 index 0000000000..7d012c234a --- /dev/null +++ b/tests/bigquery/src/writes/flaky.rs @@ -0,0 +1,326 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::writes::arrow::ArrowSerializer; +use crate::writes::count_writes_table; +use anyhow::Result; +use google_cloud_bigquery::client::Write; +use google_cloud_bigquery_v2::client::TableService; +use google_cloud_bigquery_v2::model::{Table, TableReference, TableSchema}; +use std::sync::Arc; + +pub const FLAKY_REGION: &str = "us-east7"; +const ITERATIONS: usize = 50; +const ROWS_PER_BATCH: usize = 10; +const TOTAL_ROWS: usize = ITERATIONS * ROWS_PER_BATCH; // 500 + +/// Tests sequential appends to a table with the `_reconnect_on_close` suffix in `us-east7`. +/// +/// BigQuery drops the connection every 10 requests on this table to exercise +/// reconnect and retry logic during sequential writes. +pub async fn reconnect_on_close_sequential( + client: &Write, + table_service: &TableService, + project_id: &str, + dataset_id: &str, + schema: TableSchema, +) -> Result<()> { + let table_id = format!( + "{}_reconnect_on_close", + bigquery_samples::random_id_suffix() + ); + create_flaky_table(table_service, project_id, dataset_id, &table_id, schema).await?; + + let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"); + let mut serializer = ArrowSerializer::new("reconnect_seq")?; + + let writer = client.arrow(serializer.schema()).pending(table).await?; + + for i in 0..ITERATIONS { + let offset = (i * ROWS_PER_BATCH) as i64; + let batch = serializer.generate_batch(ROWS_PER_BATCH, i * ROWS_PER_BATCH)?; + let resp = writer.append(batch).set_offset(offset).send().await?; + assert_eq!(resp.offset, Some(offset)); + } + + let finalize_resp = writer.finalize().await?; + assert_eq!(finalize_resp.row_count, TOTAL_ROWS as i64); + + let commit_resp = writer.commit().await?; + assert!( + commit_resp.stream_errors.is_empty(), + "unexpected stream errors on commit: {:?}", + commit_resp.stream_errors + ); + + let count = count_writes_table( + project_id, + dataset_id, + &table_id, + "reconnect_seq", + Some(FLAKY_REGION), + ) + .await?; + assert_eq!(count, TOTAL_ROWS as i64); + + Ok(()) +} + +/// Tests parallel appends to a table with the `_reconnect_on_close` suffix in `us-east7`. +/// +/// BigQuery drops the connection every 10 requests to exercise concurrent +/// retry and reconnect handling. +pub async fn reconnect_on_close_parallel( + client: &Write, + table_service: &TableService, + project_id: &str, + dataset_id: &str, + schema: TableSchema, +) -> Result<()> { + let table_id = format!( + "{}_reconnect_on_close_parallel", + bigquery_samples::random_id_suffix() + ); + create_flaky_table(table_service, project_id, dataset_id, &table_id, schema).await?; + + let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"); + let mut serializer = ArrowSerializer::new("reconnect_par")?; + + // Pre-generate all batches + let mut batches = Vec::with_capacity(ITERATIONS); + for i in 0..ITERATIONS { + let offset = (i * ROWS_PER_BATCH) as i64; + let batch = serializer.generate_batch(ROWS_PER_BATCH, i * ROWS_PER_BATCH)?; + batches.push((offset, batch)); + } + + let writer = Arc::new(client.arrow(serializer.schema()).pending(table).await?); + + let mut handles = Vec::with_capacity(ITERATIONS); + for (offset, batch) in batches { + let writer = writer.clone(); + handles.push(tokio::spawn(async move { + let resp = writer.append(batch).set_offset(offset).send().await?; + assert_eq!(resp.offset, Some(offset)); + Ok::<(), anyhow::Error>(()) + })); + } + + for handle in handles { + handle.await??; + } + + let finalize_resp = writer.finalize().await?; + assert_eq!(finalize_resp.row_count, TOTAL_ROWS as i64); + + let commit_resp = writer.commit().await?; + assert!( + commit_resp.stream_errors.is_empty(), + "unexpected stream errors on commit: {:?}", + commit_resp.stream_errors + ); + + let count = count_writes_table( + project_id, + dataset_id, + &table_id, + "reconnect_par", + Some(FLAKY_REGION), + ) + .await?; + assert_eq!(count, TOTAL_ROWS as i64); + + Ok(()) +} + +/// Tests sequential appends to a table with the `_initial_connect_failure` suffix in `us-east7`. +/// +/// BigQuery fails the initial connection more frequently on this table to exercise +/// stream connection retry behavior. +pub async fn initial_connect_failure_sequential( + client: &Write, + table_service: &TableService, + project_id: &str, + dataset_id: &str, + schema: TableSchema, +) -> Result<()> { + let table_id = format!( + "{}_initial_connect_failure", + bigquery_samples::random_id_suffix() + ); + create_flaky_table(table_service, project_id, dataset_id, &table_id, schema).await?; + + let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"); + let mut serializer = ArrowSerializer::new("init_fail_seq")?; + + let writer = client.arrow(serializer.schema()).pending(table).await?; + + for i in 0..ITERATIONS { + let offset = (i * ROWS_PER_BATCH) as i64; + let batch = serializer.generate_batch(ROWS_PER_BATCH, i * ROWS_PER_BATCH)?; + let resp = writer.append(batch).set_offset(offset).send().await?; + assert_eq!(resp.offset, Some(offset)); + } + + let finalize_resp = writer.finalize().await?; + assert_eq!(finalize_resp.row_count, TOTAL_ROWS as i64); + + let commit_resp = writer.commit().await?; + assert!( + commit_resp.stream_errors.is_empty(), + "unexpected stream errors on commit: {:?}", + commit_resp.stream_errors + ); + + let count = count_writes_table( + project_id, + dataset_id, + &table_id, + "init_fail_seq", + Some(FLAKY_REGION), + ) + .await?; + assert_eq!(count, TOTAL_ROWS as i64); + + Ok(()) +} + +/// Tests parallel appends to a table with the `_initial_connect_failure` suffix in `us-east7`. +/// +/// BigQuery fails the initial connection more frequently on this table to exercise +/// concurrent stream connection retry behavior. +pub async fn initial_connect_failure_parallel( + client: &Write, + table_service: &TableService, + project_id: &str, + dataset_id: &str, + schema: TableSchema, +) -> Result<()> { + let table_id = format!( + "{}_initial_connect_failure_parallel", + bigquery_samples::random_id_suffix() + ); + create_flaky_table(table_service, project_id, dataset_id, &table_id, schema).await?; + + let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"); + let mut serializer = ArrowSerializer::new("init_fail_par")?; + + // Pre-generate all batches + let mut batches = Vec::with_capacity(ITERATIONS); + for i in 0..ITERATIONS { + let offset = (i * ROWS_PER_BATCH) as i64; + let batch = serializer.generate_batch(ROWS_PER_BATCH, i * ROWS_PER_BATCH)?; + batches.push((offset, batch)); + } + + let writer = Arc::new(client.arrow(serializer.schema()).pending(table).await?); + + let mut handles = Vec::with_capacity(ITERATIONS); + for (offset, batch) in batches { + let writer = writer.clone(); + handles.push(tokio::spawn(async move { + let resp = writer.append(batch).set_offset(offset).send().await?; + assert_eq!(resp.offset, Some(offset)); + Ok::<(), anyhow::Error>(()) + })); + } + + for handle in handles { + handle.await??; + } + + let finalize_resp = writer.finalize().await?; + assert_eq!(finalize_resp.row_count, TOTAL_ROWS as i64); + + let commit_resp = writer.commit().await?; + assert!( + commit_resp.stream_errors.is_empty(), + "unexpected stream errors on commit: {:?}", + commit_resp.stream_errors + ); + + let count = count_writes_table( + project_id, + dataset_id, + &table_id, + "init_fail_par", + Some(FLAKY_REGION), + ) + .await?; + assert_eq!(count, TOTAL_ROWS as i64); + + Ok(()) +} + +/// Tests default stream appends to a table with the `_reconnect_on_close` suffix in `us-east7`. +pub async fn reconnect_on_close_default( + client: &Write, + table_service: &TableService, + project_id: &str, + dataset_id: &str, + schema: TableSchema, +) -> Result<()> { + let table_id = format!("{}_reconnect_default", bigquery_samples::random_id_suffix()); + create_flaky_table(table_service, project_id, dataset_id, &table_id, schema).await?; + + let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"); + let mut serializer = ArrowSerializer::new("reconnect_def")?; + + let writer = client.arrow(serializer.schema()).default(table).await?; + + for i in 0..ITERATIONS { + let batch = serializer.generate_batch(ROWS_PER_BATCH, i * ROWS_PER_BATCH)?; + let _ = writer.append(batch).send().await?; + } + + let count = count_writes_table( + project_id, + dataset_id, + &table_id, + "reconnect_def", + Some(FLAKY_REGION), + ) + .await?; + assert_eq!(count, TOTAL_ROWS as i64); + + Ok(()) +} + +async fn create_flaky_table( + table_service: &TableService, + project_id: &str, + dataset_id: &str, + table_id: &str, + schema: TableSchema, +) -> Result<()> { + println!("CREATING FLAKY TABLE WITH ID: {table_id} IN LOCATION: {FLAKY_REGION}"); + table_service + .insert_table() + .set_project_id(project_id) + .set_dataset_id(dataset_id) + .set_table( + Table::new() + .set_table_reference( + TableReference::new() + .set_project_id(project_id) + .set_dataset_id(dataset_id) + .set_table_id(table_id), + ) + .set_location(FLAKY_REGION) + .set_schema(schema), + ) + .send() + .await?; + Ok(()) +} diff --git a/tests/bigquery/tests/driver.rs b/tests/bigquery/tests/driver.rs index 69a2ab1185..e5c84879c8 100644 --- a/tests/bigquery/tests/driver.rs +++ b/tests/bigquery/tests/driver.rs @@ -97,6 +97,14 @@ mod bigquery { .inspect_err(anydump) } + #[tokio::test] + async fn run_writes_flaky() -> anyhow::Result<()> { + let _guard = enable_tracing(); + integration_tests_bigquery::run_writes_flaky() + .await + .inspect_err(anydump) + } + #[tokio::test] async fn run_job_service_poller_heavy() -> anyhow::Result<()> { let _guard = enable_tracing();