test(bigquery): add Storage Write API tests for flaky tables

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};

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,13 +76,94 @@ 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,
     pub(crate) age: i64,
     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<i64> {
+    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)
+}

tests/bigquery/src/writes/arrow.rs

@@ -398,14 +398,14 @@ pub async fn multiplex(
     Ok(())
 }
 
-struct ArrowSerializer {
+pub(crate) struct ArrowSerializer {
     schema: Arc<Schema>,
     writer: StreamWriter<Vec<u8>>,
     test: &'static str,
 }
 
 impl ArrowSerializer {
-    fn new(test: &'static str) -> Result<Self> {
+    pub(crate) fn new(test: &'static str) -> Result<Self> {
         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<i64>) -> Result<ArrowRecordBatch> {
+    pub(crate) fn batch(&mut self, names: Vec<&str>, ages: Vec<i64>) -> Result<ArrowRecordBatch> {
         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<ArrowRecordBatch> {
+        let names: Vec<String> = (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<i64> = (0..count).map(|i| (start_row + i) as i64).collect();
+        self.batch(name_slices, ages)
+    }
 }