test(bigquery): add integration tests for pending stream

tests/bigquery/src/table.rs

@@ -12,10 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use super::INSTANCE_LABEL;
-use crate::query::UserRecord;
 use anyhow::Result;
-use google_cloud_bigquery::client::BigQuery;
 use google_cloud_bigquery_v2::client::TableService;
 use google_cloud_bigquery_v2::model::{Table, TableFieldSchema, TableReference, TableSchema};
 
@@ -49,25 +46,3 @@ pub(crate) async fn create_table(
 
     Ok(())
 }
-
-pub(crate) async fn read_table(
-    project_id: &str,
-    dataset_id: &str,
-    table_id: &str,
-) -> Result<Vec<UserRecord>> {
-    let client = BigQuery::builder().build().await?;
-    let query = format!("SELECT * FROM `{project_id}.{dataset_id}.{table_id}` ORDER BY name");
-    let mut rows = client
-        .query(query)
-        .with_project_id(project_id)
-        .set_labels(vec![(INSTANCE_LABEL, "true")])
-        .until_done()
-        .await?
-        .read();
-
-    let mut users = Vec::new();
-    while let Some(row) = rows.next().await {
-        users.push(row?.try_into()?);
-    }
-    Ok(users)
-}

tests/bigquery/src/writes.rs

@@ -15,10 +15,12 @@
 mod arrow;
 
 use crate::dataset::{cleanup_stale_datasets, create_dataset, delete_dataset, random_dataset_id};
-use crate::query::UserRecord;
-use crate::table::{create_table, read_table};
 use anyhow::Result;
+use google_cloud_bigquery::FromRow;
+use google_cloud_bigquery::client::BigQuery;
 use google_cloud_bigquery_v2::client::{DatasetService, TableService};
+use google_cloud_bigquery_v2::model::{Table, TableFieldSchema, TableReference, TableSchema};
+use google_cloud_bigquery_write::client::Write;
 use google_cloud_test_utils::runtime_config::project_id;
 
 pub async fn run_writes() -> Result<()> {
@@ -33,11 +35,77 @@ pub async fn run_writes() -> Result<()> {
     let table_id = "writes";
 
     let result = async {
-        create_table(&table_service, &project_id, &dataset_id, table_id).await?;
-        arrow::basic(&project_id, &dataset_id, table_id).await
+        create_writes_table(&table_service, &project_id, &dataset_id, table_id).await?;
+        let client = Write::builder().build().await?;
+        arrow::basic(&client, &project_id, &dataset_id, table_id).await?;
+        arrow::pending(&client, &project_id, &dataset_id, table_id).await
     }
     .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,
+}
+
+pub(crate) async fn read_writes_table(
+    project_id: &str,
+    dataset_id: &str,
+    table_id: &str,
+    test_filter: &str,
+) -> Result<Vec<WriteUserRecord>> {
+    let client = BigQuery::builder().build().await?;
+    let query = format!(
+        "SELECT * FROM `{project_id}.{dataset_id}.{table_id}` WHERE test = '{test_filter}' ORDER BY name"
+    );
+    let mut rows = client
+        .query(query)
+        .with_project_id(project_id)
+        .set_labels(vec![(crate::INSTANCE_LABEL, "true")])
+        .until_done()
+        .await?
+        .read();
+
+    let mut users = Vec::new();
+    while let Some(row) = rows.next().await {
+        users.push(row?.try_into()?);
+    }
+    Ok(users)
+}
+
+pub(crate) async fn create_writes_table(
+    table_service: &TableService,
+    project_id: &str,
+    dataset_id: &str,
+    table_id: &str,
+) -> anyhow::Result<()> {
+    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"),
+    ]);
+
+    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_schema(schema),
+        )
+        .send()
+        .await?;
+
+    Ok(())
+}

tests/bigquery/src/writes/arrow.rs

@@ -12,73 +12,162 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use super::*;
+use crate::writes::{WriteUserRecord, read_writes_table};
 use ::arrow::array::{Int64Array, StringArray};
 use ::arrow::datatypes::{DataType, Field, Schema};
 use ::arrow::ipc::writer::StreamWriter;
 use ::arrow::record_batch::RecordBatch;
+use anyhow::Result;
 use google_cloud_bigquery_write::client::Write;
 use google_cloud_bigquery_write::model::{ArrowRecordBatch, ArrowSchema};
 use std::sync::Arc;
 
-pub async fn basic(project_id: &str, dataset_id: &str, table_id: &str) -> Result<()> {
+pub async fn basic(
+    client: &Write,
+    project_id: &str,
+    dataset_id: &str,
+    table_id: &str,
+) -> Result<()> {
     let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}");
-
-    // Create a Schema
-    let arrow_schema = Arc::new(Schema::new(vec![
-        Field::new("name", DataType::Utf8, false),
-        Field::new("age", DataType::Int64, false),
-    ]));
-    let schema_buf = serialize_schema(&arrow_schema)?;
-    let schema_len = schema_buf.len();
+    let schema = create_test_schema();
 
     // Create a writer for the default stream
-    let client = Write::builder().build().await?;
-    let schema = ArrowSchema::new().set_serialized_schema(schema_buf);
-    let writer = client.arrow(schema).default(table)?;
-
-    // Create a RecordBatch
-    let name = StringArray::from(vec!["Alice", "Bob"]);
-    let age = Int64Array::from(vec![25, 28]);
-    let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(name), Arc::new(age)])?;
-    let batch_buf = serialize_batch(&batch, schema_len)?;
+    let writer = client
+        .arrow(ArrowSchema::new().set_serialized_schema(serialize_schema(&schema)?))
+        .default(table)?;
 
-    // Write the batch
-    let rows = ArrowRecordBatch::new().set_serialized_record_batch(batch_buf);
-    let _ = writer.append(rows).send().await?;
+    // Write the batches
+    let batch1 = create_test_batch(schema.clone(), vec!["Alice", "Bob"], vec![25, 28], "basic")?;
+    let _ = writer.append(batch1).send().await?;
 
-    // Create a second RecordBatch
-    let name = StringArray::from(vec!["Charlie"]);
-    let age = Int64Array::from(vec![31]);
-    let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(name), Arc::new(age)])?;
-    let batch_buf = serialize_batch(&batch, schema_len)?;
-
-    // Write the second batch
-    let rows = ArrowRecordBatch::new().set_serialized_record_batch(batch_buf);
-    let _ = writer.append(rows).send().await?;
+    let batch2 = create_test_batch(schema.clone(), vec!["Charlie"], vec![31], "basic")?;
+    let _ = writer.append(batch2).send().await?;
 
     // Verify the writes
-    let users = read_table(project_id, dataset_id, table_id).await?;
+    let users = read_writes_table(project_id, dataset_id, table_id, "basic").await?;
     assert_eq!(
         users,
         vec![
-            UserRecord {
+            WriteUserRecord {
                 name: "Alice".to_string(),
                 age: 25,
+                test: "basic".to_string()
             },
-            UserRecord {
+            WriteUserRecord {
                 name: "Bob".to_string(),
                 age: 28,
+                test: "basic".to_string()
             },
-            UserRecord {
+            WriteUserRecord {
                 name: "Charlie".to_string(),
                 age: 31,
+                test: "basic".to_string()
+            },
+        ]
+    );
+
+    Ok(())
+}
+
+pub async fn pending(
+    client: &Write,
+    project_id: &str,
+    dataset_id: &str,
+    table_id: &str,
+) -> Result<()> {
+    let table = format!("projects/{project_id}/datasets/{dataset_id}/tables/{table_id}");
+    let schema = create_test_schema();
+
+    // Create a writer for a pending stream
+    let writer = client
+        .arrow(ArrowSchema::new().set_serialized_schema(serialize_schema(&schema)?))
+        .pending(table)
+        .await?;
+
+    // Write the batches
+    let batch1 = create_test_batch(
+        schema.clone(),
+        vec!["David", "Eve"],
+        vec![42, 38],
+        "pending",
+    )?;
+    let _ = writer.append(batch1).set_offset(0).send().await?;
+
+    let batch2 = create_test_batch(schema.clone(), vec!["Frank"], vec![55], "pending")?;
+    let _ = writer.append(batch2).set_offset(2).send().await?;
+
+    // Finalize the stream
+    writer.finalize().await?;
+
+    // Verify no writes have been committed yet
+    let users = read_writes_table(project_id, dataset_id, table_id, "pending").await?;
+    assert!(users.is_empty(), "{users:?}");
+
+    // Verify that appending to a finalized stream fails
+    let batch3 = create_test_batch(schema.clone(), vec!["Ghost"], vec![99], "pending")?;
+    let _err = writer
+        .append(batch3)
+        .set_offset(3)
+        .send()
+        .await
+        .expect_err("Appending to a finalized stream should fail");
+    // Commit the stream
+    writer.commit().await?;
+
+    // Verify the writes
+    let users = read_writes_table(project_id, dataset_id, table_id, "pending").await?;
+    assert_eq!(
+        users,
+        vec![
+            WriteUserRecord {
+                name: "David".to_string(),
+                age: 42,
+                test: "pending".to_string()
+            },
+            WriteUserRecord {
+                name: "Eve".to_string(),
+                age: 38,
+                test: "pending".to_string()
+            },
+            WriteUserRecord {
+                name: "Frank".to_string(),
+                age: 55,
+                test: "pending".to_string()
             },
         ]
     );
 
     Ok(())
 }
+fn create_test_schema() -> Arc<Schema> {
+    Arc::new(Schema::new(vec![
+        Field::new("name", DataType::Utf8, false),
+        Field::new("age", DataType::Int64, false),
+        Field::new("test", DataType::Utf8, false),
+    ]))
+}
+
+fn create_test_batch(
+    schema: Arc<Schema>,
+    names: Vec<&str>,
+    ages: Vec<i64>,
+    test: &str,
+) -> Result<ArrowRecordBatch> {
+    let schema_buf = serialize_schema(&schema)?;
+    let schema_len = schema_buf.len();
+
+    let name = StringArray::from(names);
+    let age = Int64Array::from(ages);
+    let test_col = StringArray::from(vec![test; age.len()]);
+
+    let batch = RecordBatch::try_new(
+        schema,
+        vec![Arc::new(name), Arc::new(age), Arc::new(test_col)],
+    )?;
+    let batch_buf = serialize_batch(&batch, schema_len)?;
+
+    Ok(ArrowRecordBatch::new().set_serialized_record_batch(batch_buf))
+}
 
 fn serialize_schema(schema: &Schema) -> Result<Vec<u8>> {
     let mut buf = Vec::new();