impl(bigquery): add internal proto custom stream writers

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;

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<Transport>, write_stream: String, schema: ProtoSchema) -> Self {
+        let runner = Runner::new(inner.clone());
+        let client = BigQueryWrite::from_stub::<Transport>(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<FinalizeWriteStreamResponse> {
+        self.client
+            .finalize_write_stream()
+            .set_name(&self.write_stream)
+            .send()
+            .await
+    }
+}

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<Transport>, 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<FlushRowsResponse> {
+        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<FinalizeWriteStreamResponse> {
+        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())])
+    }
+}