diff --git a/src/bigquery/src/write/proto.rs b/src/bigquery/src/write/proto.rs index 3cad0c4d86..ba76b3783f 100644 --- a/src/bigquery/src/write/proto.rs +++ b/src/bigquery/src/write/proto.rs @@ -12,8 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod base; +mod buffered; +mod committed; mod default; +mod pending; mod writer_builder; +pub(crate) use buffered::BufferedWriter; +pub(crate) use committed::CommittedWriter; pub(crate) use default::DefaultWriter; +pub(crate) use pending::PendingWriter; pub(crate) use writer_builder::WriterBuilder; diff --git a/src/bigquery/src/write/proto/base.rs b/src/bigquery/src/write/proto/base.rs new file mode 100644 index 0000000000..fc1c23980f --- /dev/null +++ b/src/bigquery/src/write/proto/base.rs @@ -0,0 +1,63 @@ +// 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::Result; +use crate::model::append_rows_request::ProtoData; +use crate::model::{AppendRowsRequest, FinalizeWriteStreamResponse, ProtoRows, ProtoSchema}; +use crate::write::generated::gapic_storage::client::BigQueryWrite; +use crate::write::runner::Runner; +use crate::write::transport::Transport; +use std::sync::Arc; + +/// A shared internal structure for holding common state across different stream types. +/// Providing shared implementations of operations core to most write streams. +/// Specific stream behaviors should be handled individually by their respective wrapper structs (e.g. `BufferedWriter`, `CommittedWriter`, `PendingWriter`). +#[derive(Debug)] +pub(crate) struct BaseWriter { + pub(crate) runner: Runner, + pub(crate) write_stream: String, + pub(crate) schema: ProtoSchema, + pub(crate) client: BigQueryWrite, +} + +impl BaseWriter { + pub(crate) fn new(inner: Arc, write_stream: String, schema: ProtoSchema) -> Self { + let runner = Runner::new(inner.clone()); + let client = BigQueryWrite::from_stub::(inner); + Self { + runner, + write_stream, + schema, + client, + } + } + + pub(crate) fn append_request(&self, rows: ProtoRows) -> AppendRowsRequest { + AppendRowsRequest::new() + .set_write_stream(&self.write_stream) + .set_proto_rows( + ProtoData::new() + .set_writer_schema(self.schema.clone()) + .set_rows(rows), + ) + } + + pub(crate) async fn finalize(&self) -> Result { + self.client + .finalize_write_stream() + .set_name(&self.write_stream) + .send() + .await + } +} diff --git a/src/bigquery/src/write/proto/buffered.rs b/src/bigquery/src/write/proto/buffered.rs new file mode 100644 index 0000000000..933d16ddfb --- /dev/null +++ b/src/bigquery/src/write/proto/buffered.rs @@ -0,0 +1,188 @@ +// 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 super::base::BaseWriter; +use crate::Result; +use crate::model::{FinalizeWriteStreamResponse, FlushRowsResponse, ProtoRows, ProtoSchema}; +use crate::write::builder::AppendWithOffset; +use crate::write::transport::Transport; +use std::sync::Arc; + +/// A writer for a [buffered stream] using Protobuf as the data format. +/// +/// [buffered stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#buffered_type +#[derive(Debug)] +pub struct BufferedWriter { + pub(crate) inner: BaseWriter, +} + +impl BufferedWriter { + pub(crate) fn new(inner: Arc, write_stream: String, schema: ProtoSchema) -> Self { + Self { + inner: BaseWriter::new(inner, write_stream, schema), + } + } + + /// Return the full resource name of the underlying write stream. + pub fn write_stream(&self) -> &str { + &self.inner.write_stream + } + + /// Append rows to the buffered stream. + pub fn append(&self, rows: ProtoRows) -> AppendWithOffset { + AppendWithOffset::new( + self.inner.runner.req_tx.clone(), + self.inner.append_request(rows), + ) + } + + /// Flush the buffered stream, making rows up to the specified offset available for reading. + pub async fn flush(&self, offset: i64) -> Result { + self.inner + .client + .flush_rows() + .set_write_stream(&self.inner.write_stream) + .set_offset(offset) + .send() + .await + } + + /// Finalize the buffered stream, preventing further writes. + pub async fn finalize(&self) -> Result { + self.inner.finalize().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::AppendError; + use crate::write::test::*; + use bigquery_grpc_mock::{MockBigQueryWrite, start}; + use gaxi::grpc::tonic::Response as TonicResponse; + use tokio::sync::mpsc; + + #[tokio::test] + async fn request_fields() -> anyhow::Result<()> { + let transport = Arc::new(test_transport("http://ignored:1").await?); + let writer = BufferedWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + let b = writer.append(rows(1)); + assert_eq!(b.req.write_stream, write_stream()); + let data = b.req.proto_rows().expect("proto rows should be set"); + let s = data.writer_schema.as_ref().expect("schema should be set"); + assert_eq!(s.proto_descriptor.as_ref().unwrap().name, "TestMessage"); + let r = data.rows.as_ref().expect("rows should be set"); + assert_eq!(r.serialized_rows, vec![bytes::Bytes::from("1")]); + + let b = writer.append(rows(2)); + assert_eq!(b.req.write_stream, write_stream()); + let data = b.req.proto_rows().expect("proto rows should be set"); + let s = data.writer_schema.as_ref().expect("schema should be set"); + assert_eq!(s.proto_descriptor.as_ref().unwrap().name, "TestMessage"); + let r = data.rows.as_ref().expect("rows should be set"); + assert_eq!(r.serialized_rows, vec![bytes::Bytes::from("2")]); + + Ok(()) + } + + #[tokio::test] + async fn basic_success() -> anyhow::Result<()> { + let (response_tx, response_rx) = mpsc::channel(10); + + let mut mock = MockBigQueryWrite::new(); + mock.expect_append_rows() + .return_once(|_| Ok(TonicResponse::from(response_rx))); + + mock.expect_flush_rows().return_once(|req| { + assert_eq!(req.get_ref().offset, Some(3)); + assert_eq!(req.get_ref().write_stream, write_stream()); + Ok(TonicResponse::new( + bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FlushRowsResponse::default(), + )) + }); + + mock.expect_finalize_write_stream().return_once(|req| { + assert_eq!(req.get_ref().name, write_stream()); + Ok(TonicResponse::new( + bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FinalizeWriteStreamResponse::default(), + )) + }); + + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + + let writer = BufferedWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + response_tx.send(Ok(convert(&test_response(1)))).await?; + let resp = writer.append(rows(1)).send().await?; + assert_eq!(resp.offset, Some(1)); + + response_tx.send(Ok(convert(&test_response(2)))).await?; + let resp = writer.append(rows(2)).send().await?; + assert_eq!(resp.offset, Some(2)); + + response_tx.send(Ok(convert(&test_response(3)))).await?; + let resp = writer.append(rows(3)).send().await?; + assert_eq!(resp.offset, Some(3)); + + drop(response_tx); + let err = writer.append(rows(4)).send().await.expect_err("channel"); + assert!(matches!(err, AppendError::UnexpectedEndOfStream)); + + writer.flush(3).await?; + writer.finalize().await?; + + Ok(()) + } + + #[tokio::test] + async fn multiple_flushes() -> anyhow::Result<()> { + let (response_tx, response_rx) = mpsc::channel(10); + let mut mock = MockBigQueryWrite::new(); + mock.expect_append_rows() + .return_once(|_| Ok(TonicResponse::from(response_rx))); + + mock.expect_flush_rows().times(2).returning(|req| { + Ok(TonicResponse::new( + bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FlushRowsResponse { + offset: req.get_ref().offset.unwrap_or(0), + }, + )) + }); + + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + let writer = BufferedWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + response_tx.send(Ok(convert(&test_response(1)))).await?; + let _ = writer.append(rows(1)).send().await?; + let flush1 = writer.flush(1).await?; + assert_eq!(flush1.offset, 1); + + response_tx.send(Ok(convert(&test_response(2)))).await?; + let _ = writer.append(rows(2)).send().await?; + let flush2 = writer.flush(2).await?; + assert_eq!(flush2.offset, 2); + + Ok(()) + } + + fn rows(id: i64) -> ProtoRows { + ProtoRows::new().set_serialized_rows(vec![bytes::Bytes::from(id.to_string())]) + } +} diff --git a/src/bigquery/src/write/proto/committed.rs b/src/bigquery/src/write/proto/committed.rs new file mode 100644 index 0000000000..82d86f3272 --- /dev/null +++ b/src/bigquery/src/write/proto/committed.rs @@ -0,0 +1,135 @@ +// 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 super::base::BaseWriter; +use crate::Result; +use crate::model::{FinalizeWriteStreamResponse, ProtoRows, ProtoSchema}; +use crate::write::builder::AppendWithOffset; +use crate::write::transport::Transport; +use std::sync::Arc; + +/// A writer for a [committed stream] using Protobuf as the data format. +/// +/// [committed stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#committed_type +#[derive(Debug)] +pub struct CommittedWriter { + pub(crate) inner: BaseWriter, +} + +impl CommittedWriter { + pub(crate) fn new(inner: Arc, write_stream: String, schema: ProtoSchema) -> Self { + Self { + inner: BaseWriter::new(inner, write_stream, schema), + } + } + + /// Return the full resource name of the underlying write stream. + pub fn write_stream(&self) -> &str { + &self.inner.write_stream + } + + /// Append rows to the committed stream. + pub fn append(&self, rows: ProtoRows) -> AppendWithOffset { + AppendWithOffset::new( + self.inner.runner.req_tx.clone(), + self.inner.append_request(rows), + ) + } + + /// Finalize the stream, preventing further writes. + pub async fn finalize(&self) -> Result { + self.inner.finalize().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::AppendError; + use crate::write::test::*; + use bigquery_grpc_mock::{MockBigQueryWrite, start}; + use gaxi::grpc::tonic::Response as TonicResponse; + use tokio::sync::mpsc; + + #[tokio::test] + async fn request_fields() -> anyhow::Result<()> { + let transport = Arc::new(test_transport("http://ignored:1").await?); + let writer = CommittedWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + let b = writer.append(rows(1)); + assert_eq!(b.req.write_stream, write_stream()); + let data = b.req.proto_rows().expect("proto rows should be set"); + let s = data.writer_schema.as_ref().expect("schema should be set"); + assert_eq!(s.proto_descriptor.as_ref().unwrap().name, "TestMessage"); + let r = data.rows.as_ref().expect("rows should be set"); + assert_eq!(r.serialized_rows, vec![bytes::Bytes::from("1")]); + + let b = writer.append(rows(2)); + assert_eq!(b.req.write_stream, write_stream()); + let data = b.req.proto_rows().expect("proto rows should be set"); + let s = data.writer_schema.as_ref().expect("schema should be set"); + assert_eq!(s.proto_descriptor.as_ref().unwrap().name, "TestMessage"); + let r = data.rows.as_ref().expect("rows should be set"); + assert_eq!(r.serialized_rows, vec![bytes::Bytes::from("2")]); + + Ok(()) + } + + #[tokio::test] + async fn basic_success() -> anyhow::Result<()> { + let (response_tx, response_rx) = mpsc::channel(10); + + let mut mock = MockBigQueryWrite::new(); + mock.expect_append_rows() + .return_once(|_| Ok(TonicResponse::from(response_rx))); + + mock.expect_finalize_write_stream().return_once(|_| { + Ok(TonicResponse::new( + bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FinalizeWriteStreamResponse::default(), + )) + }); + + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + + let writer = CommittedWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + response_tx.send(Ok(convert(&test_response(1)))).await?; + let resp = writer.append(rows(1)).send().await?; + assert_eq!(resp.offset, Some(1)); + + response_tx.send(Ok(convert(&test_response(2)))).await?; + let resp = writer.append(rows(2)).send().await?; + assert_eq!(resp.offset, Some(2)); + + response_tx.send(Ok(convert(&test_response(3)))).await?; + let resp = writer.append(rows(3)).send().await?; + assert_eq!(resp.offset, Some(3)); + + drop(response_tx); + let err = writer.append(rows(4)).send().await.expect_err("channel"); + assert!(matches!(err, AppendError::UnexpectedEndOfStream)); + + // We can still finalize the stream even if row appends hit a closed bidirectional stream + writer.finalize().await?; + + Ok(()) + } + + fn rows(id: i64) -> ProtoRows { + ProtoRows::new().set_serialized_rows(vec![bytes::Bytes::from(id.to_string())]) + } +} diff --git a/src/bigquery/src/write/proto/pending.rs b/src/bigquery/src/write/proto/pending.rs new file mode 100644 index 0000000000..ae8ad2cf5b --- /dev/null +++ b/src/bigquery/src/write/proto/pending.rs @@ -0,0 +1,139 @@ +// 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 super::base::BaseWriter; +use crate::Result; +use crate::model::{ + BatchCommitWriteStreamsResponse, FinalizeWriteStreamResponse, ProtoRows, ProtoSchema, +}; +use crate::write::builder::AppendWithOffset; +use crate::write::transport::Transport; +use std::sync::Arc; + +/// A writer for a [pending stream] using Protobuf as the data format. +/// +/// [pending stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#pending_type +#[derive(Debug)] +pub struct PendingWriter { + pub(crate) inner: BaseWriter, +} + +impl PendingWriter { + pub(crate) fn new(inner: Arc, write_stream: String, schema: ProtoSchema) -> Self { + Self { + inner: BaseWriter::new(inner, write_stream, schema), + } + } + + /// Returns the full resource name of the underlying write stream. + pub fn write_stream(&self) -> &str { + &self.inner.write_stream + } + + /// Appends rows to the pending stream. + pub fn append(&self, rows: ProtoRows) -> AppendWithOffset { + AppendWithOffset::new( + self.inner.runner.req_tx.clone(), + self.inner.append_request(rows), + ) + } + + /// Finalizes the pending stream, preventing further writes. + pub async fn finalize(&self) -> Result { + self.inner.finalize().await + } + + /// Commits the pending stream to the table. + pub async fn commit(&self) -> Result { + // Extract the parent table path from the stream name: + // "projects/p/datasets/d/tables/t/streams/s" -> "projects/p/datasets/d/tables/t" + let parent = self + .inner + .write_stream + .split_once("/streams/") + .map_or(self.inner.write_stream.as_str(), |(p, _)| p); + + self.inner + .client + .batch_commit_write_streams() + .set_parent(parent) + .set_write_streams(vec![self.inner.write_stream.clone()]) + .send() + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::write::test::*; + use bigquery_grpc_mock::{MockBigQueryWrite, start}; + use gaxi::grpc::tonic::Response as TonicResponse; + use tokio::sync::mpsc; + + #[tokio::test] + async fn request_fields() -> anyhow::Result<()> { + let transport = Arc::new(test_transport("http://ignored:1").await?); + let writer = PendingWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + let b = writer.append(rows(1)); + assert_eq!(b.req.write_stream, write_stream()); + let data = b.req.proto_rows().expect("proto rows should be set"); + let s = data.writer_schema.as_ref().expect("schema should be set"); + assert_eq!(s.proto_descriptor.as_ref().unwrap().name, "TestMessage"); + let r = data.rows.as_ref().expect("rows should be set"); + assert_eq!(r.serialized_rows, vec![bytes::Bytes::from("1")]); + + Ok(()) + } + + #[tokio::test] + async fn basic_success() -> anyhow::Result<()> { + let (response_tx, response_rx) = mpsc::channel(10); + + let mut mock = MockBigQueryWrite::new(); + mock.expect_append_rows() + .return_once(|_| Ok(TonicResponse::from(response_rx))); + + mock.expect_finalize_write_stream() + .return_once(|_| Ok(TonicResponse::new( + bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FinalizeWriteStreamResponse::default() + ))); + + mock.expect_batch_commit_write_streams() + .return_once(|_| Ok(TonicResponse::new( + bigquery_grpc_mock::google::cloud::bigquery::storage::v1::BatchCommitWriteStreamsResponse::default() + ))); + + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + + let writer = PendingWriter::new(transport, write_stream(), proto_schema()); + assert_eq!(writer.write_stream(), write_stream()); + + response_tx.send(Ok(convert(&test_response(1)))).await?; + let resp = writer.append(rows(1)).send().await?; + assert_eq!(resp.offset, Some(1)); + + writer.finalize().await?; + writer.commit().await?; + + Ok(()) + } + + fn rows(id: i64) -> ProtoRows { + ProtoRows::new().set_serialized_rows(vec![bytes::Bytes::from(id.to_string())]) + } +} diff --git a/src/bigquery/src/write/proto/writer_builder.rs b/src/bigquery/src/write/proto/writer_builder.rs index b06b09ee16..1192401902 100644 --- a/src/bigquery/src/write/proto/writer_builder.rs +++ b/src/bigquery/src/write/proto/writer_builder.rs @@ -12,11 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +use super::super::generated::gapic_storage::client::BigQueryWrite; use super::super::transport::Transport; use super::super::validate::validate_table; -use super::DefaultWriter; +use super::{BufferedWriter, CommittedWriter, DefaultWriter, PendingWriter}; use crate::Result; -use crate::model::ProtoSchema; +use crate::model::write_stream::Type; +use crate::model::{ProtoSchema, WriteStream}; use std::sync::Arc; /// A builder to create a protobuf stream writer @@ -41,14 +43,162 @@ impl WriterBuilder { write_stream.push_str("/streams/_default"); Ok(DefaultWriter::new(self.inner, write_stream, self.schema)) } + + /// Creates a writer for a [pending stream] for the given table. + /// + /// [pending stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#pending_type + pub async fn pending>(self, table: T) -> Result { + let table = table.into(); + validate_table(table.as_str())?; + + let client = BigQueryWrite::from_stub::(self.inner.clone()); + let write_stream = client + .create_write_stream() + .set_parent(table) + .set_write_stream(WriteStream::new().set_type(Type::Pending)) + .send() + .await?; + + Ok(PendingWriter::new( + self.inner, + write_stream.name, + self.schema, + )) + } + + /// Creates a writer for a [committed stream] for the given table. + /// + /// [committed stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#committed_type + pub async fn committed>(self, table: T) -> Result { + let table = table.into(); + validate_table(table.as_str())?; + + let client = BigQueryWrite::from_stub::(self.inner.clone()); + let write_stream = client + .create_write_stream() + .set_parent(table) + .set_write_stream(WriteStream::new().set_type(Type::Committed)) + .send() + .await?; + + Ok(CommittedWriter::new( + self.inner, + write_stream.name, + self.schema, + )) + } + + /// Creates a writer for a [buffered stream] for the given table. + /// + /// [buffered stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#buffered_type + pub async fn buffered>(self, table: T) -> Result { + let table = table.into(); + validate_table(table.as_str())?; + + let client = BigQueryWrite::from_stub::(self.inner.clone()); + let write_stream = client + .create_write_stream() + .set_parent(table) + .set_write_stream(WriteStream::new().set_type(Type::Buffered)) + .send() + .await?; + + Ok(BufferedWriter::new( + self.inner, + write_stream.name, + self.schema, + )) + } } #[cfg(test)] mod tests { use super::*; use crate::write::test::*; + use bigquery_grpc_mock::google::cloud::bigquery::storage::v1::WriteStream as MockWriteStream; + use bigquery_grpc_mock::{MockBigQueryWrite, start}; use test_case::test_case; + #[tokio::test] + async fn pending_success() -> anyhow::Result<()> { + let mut mock = MockBigQueryWrite::new(); + mock.expect_create_write_stream().return_once(|req| { + let req = req.into_inner(); + assert_eq!(req.parent, "projects/p/datasets/d/tables/t"); + let ws = req.write_stream.expect("write_stream populated"); + assert_eq!(Type::from(ws.r#type), Type::Pending); + Ok(gaxi::grpc::tonic::Response::new(MockWriteStream { + name: "projects/p/datasets/d/tables/t/streams/s".to_string(), + ..Default::default() + })) + }); + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + let builder = WriterBuilder::new(transport, proto_schema()); + let writer = builder.pending("projects/p/datasets/d/tables/t").await?; + assert_eq!( + writer.inner.write_stream, + "projects/p/datasets/d/tables/t/streams/s" + ); + assert_eq!(writer.inner.schema, proto_schema()); + Ok(()) + } + + #[test_case("projects/p")] + #[test_case("projects/p/tables/t")] + #[test_case("projects/p/datasets/d/tables/")] + #[tokio::test] + async fn pending_bad_table_format(table: &str) -> anyhow::Result<()> { + let transport = Arc::new(test_transport("http://ignored:1").await?); + let builder = WriterBuilder::new(transport, proto_schema()); + let err = builder + .pending(table) + .await + .expect_err("should fail locally on bad format"); + assert!(err.is_binding(), "{err:?}"); + Ok(()) + } + + #[tokio::test] + async fn committed_success() -> anyhow::Result<()> { + let mut mock = MockBigQueryWrite::new(); + mock.expect_create_write_stream().return_once(|req| { + let req = req.into_inner(); + assert_eq!(req.parent, "projects/p/datasets/d/tables/t"); + let ws = req.write_stream.expect("write_stream populated"); + assert_eq!(Type::from(ws.r#type), Type::Committed); + Ok(gaxi::grpc::tonic::Response::new(MockWriteStream { + name: "projects/p/datasets/d/tables/t/streams/s".to_string(), + ..Default::default() + })) + }); + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + let builder = WriterBuilder::new(transport, proto_schema()); + let writer = builder.committed("projects/p/datasets/d/tables/t").await?; + assert_eq!( + writer.inner.write_stream, + "projects/p/datasets/d/tables/t/streams/s" + ); + assert_eq!(writer.inner.schema, proto_schema()); + Ok(()) + } + + #[test_case("projects/p")] + #[test_case("projects/p/tables/t")] + #[test_case("projects/p/datasets/d/tables/")] + #[tokio::test] + async fn committed_bad_table_format(table: &str) -> anyhow::Result<()> { + let transport = Arc::new(test_transport("http://ignored:1").await?); + let builder = WriterBuilder::new(transport, proto_schema()); + let err = builder + .committed(table) + .await + .expect_err("should fail locally on bad format"); + assert!(err.is_binding(), "{err:?}"); + Ok(()) + } + #[tokio::test] async fn default() -> anyhow::Result<()> { let transport = Arc::new(test_transport("http://ignored:1").await?); @@ -78,4 +228,44 @@ mod tests { assert!(err.is_binding(), "{err:?}"); Ok(()) } + + #[tokio::test] + async fn buffered_success() -> anyhow::Result<()> { + let mut mock = MockBigQueryWrite::new(); + mock.expect_create_write_stream().return_once(|req| { + let req = req.into_inner(); + assert_eq!(req.parent, "projects/p/datasets/d/tables/t"); + let ws = req.write_stream.expect("write_stream populated"); + assert_eq!(Type::from(ws.r#type), Type::Buffered); + Ok(gaxi::grpc::tonic::Response::new(MockWriteStream { + name: "projects/p/datasets/d/tables/t/streams/s".to_string(), + ..Default::default() + })) + }); + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + let transport = Arc::new(test_transport(endpoint).await?); + let builder = WriterBuilder::new(transport, proto_schema()); + let writer = builder.buffered("projects/p/datasets/d/tables/t").await?; + assert_eq!( + writer.inner.write_stream, + "projects/p/datasets/d/tables/t/streams/s" + ); + assert_eq!(writer.inner.schema, proto_schema()); + Ok(()) + } + + #[test_case("projects/p")] + #[test_case("projects/p/tables/t")] + #[test_case("projects/p/datasets/d/tables/")] + #[tokio::test] + async fn buffered_bad_table_format(table: &str) -> anyhow::Result<()> { + let transport = Arc::new(test_transport("http://ignored:1").await?); + let builder = WriterBuilder::new(transport, proto_schema()); + let err = builder + .buffered(table) + .await + .expect_err("should fail locally on bad format"); + assert!(err.is_binding(), "{err:?}"); + Ok(()) + } }