diff --git a/tests/bigquery/src/table.rs b/tests/bigquery/src/table.rs index be01d8b02f..49cc3ff249 100644 --- a/tests/bigquery/src/table.rs +++ b/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> { - 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) -} diff --git a/tests/bigquery/src/writes.rs b/tests/bigquery/src/writes.rs index a9635e893a..b9a430bf1b 100644 --- a/tests/bigquery/src/writes.rs +++ b/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> { + 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(()) +} diff --git a/tests/bigquery/src/writes/arrow.rs b/tests/bigquery/src/writes/arrow.rs index c1629f9a6f..331d2d458c 100644 --- a/tests/bigquery/src/writes/arrow.rs +++ b/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 { + 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, + names: Vec<&str>, + ages: Vec, + test: &str, +) -> Result { + 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> { let mut buf = Vec::new();