google_cloud_bigquery/write/arrow/
committed.rs

1// Copyright 2026 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::base::BaseWriter;
16use crate::Result;
17use crate::model::{ArrowRecordBatch, ArrowSchema, FinalizeWriteStreamResponse};
18use crate::write::append_builder::AppendWithOffset;
19use crate::write::transport::Transport;
20use std::sync::Arc;
21
22/// A writer for a [committed stream].
23///
24/// [committed stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#committed_type
25#[derive(Debug)]
26pub struct CommittedWriter {
27    pub(crate) inner: BaseWriter,
28}
29
30impl CommittedWriter {
31    pub(crate) fn new(inner: Arc<Transport>, write_stream: String, schema: ArrowSchema) -> Self {
32        Self {
33            inner: BaseWriter::new(inner, write_stream, schema),
34        }
35    }
36
37    /// Return the full resource name of the underlying write stream.
38    pub fn write_stream(&self) -> &str {
39        &self.inner.write_stream
40    }
41
42    /// Append rows to the committed stream.
43    pub fn append(&self, rows: ArrowRecordBatch) -> AppendWithOffset {
44        AppendWithOffset::new(
45            self.inner.runner.req_tx.clone(),
46            self.inner.append_request(rows),
47        )
48    }
49
50    /// Finalize the stream, preventing further writes.
51    pub async fn finalize(&self) -> Result<FinalizeWriteStreamResponse> {
52        self.inner.finalize().await
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::super::super::runner::tests::*;
59    use super::super::super::transport::tests::*;
60    use super::*;
61    use crate::error::AppendError;
62    use bigquery_grpc_mock::{MockBigQueryWrite, start};
63    use gaxi::grpc::tonic::Response as TonicResponse;
64    use tokio::sync::mpsc;
65
66    #[tokio::test]
67    async fn request_fields() -> anyhow::Result<()> {
68        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
69        let writer = CommittedWriter::new(transport, write_stream(), schema());
70        assert_eq!(writer.write_stream(), write_stream());
71
72        let b = writer.append(rows(1));
73        assert_eq!(b.req.write_stream, write_stream());
74        let data = b.req.arrow_rows().expect("arrow rows should be set");
75        let s = data.writer_schema.as_ref().expect("schema should be set");
76        assert_eq!(s.serialized_schema, "test");
77        let r = data.rows.as_ref().expect("rows should be set");
78        assert_eq!(r.serialized_record_batch, "1");
79
80        let b = writer.append(rows(2));
81        assert_eq!(b.req.write_stream, write_stream());
82        let data = b.req.arrow_rows().expect("arrow rows should be set");
83        let s = data.writer_schema.as_ref().expect("schema should be set");
84        assert_eq!(s.serialized_schema, "test");
85        let r = data.rows.as_ref().expect("rows should be set");
86        assert_eq!(r.serialized_record_batch, "2");
87
88        Ok(())
89    }
90
91    #[tokio::test]
92    async fn basic_success() -> anyhow::Result<()> {
93        let (response_tx, response_rx) = mpsc::channel(10);
94
95        let mut mock = MockBigQueryWrite::new();
96        mock.expect_append_rows()
97            .return_once(|_| Ok(TonicResponse::from(response_rx)));
98
99        mock.expect_finalize_write_stream().return_once(|_| {
100            Ok(TonicResponse::new(
101                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FinalizeWriteStreamResponse::default(),
102            ))
103        });
104
105        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
106        let transport = Arc::new(test_transport(endpoint).await?);
107
108        let writer = CommittedWriter::new(transport, write_stream(), schema());
109        assert_eq!(writer.write_stream(), write_stream());
110
111        response_tx.send(Ok(convert(&test_response(1)))).await?;
112        let resp = writer.append(rows(1)).send().await?;
113        assert_eq!(resp.offset, Some(1));
114
115        response_tx.send(Ok(convert(&test_response(2)))).await?;
116        let resp = writer.append(rows(2)).send().await?;
117        assert_eq!(resp.offset, Some(2));
118
119        response_tx.send(Ok(convert(&test_response(3)))).await?;
120        let resp = writer.append(rows(3)).send().await?;
121        assert_eq!(resp.offset, Some(3));
122
123        drop(response_tx);
124        let err = writer.append(rows(4)).send().await.expect_err("channel");
125        assert!(matches!(err, AppendError::UnexpectedEndOfStream));
126
127        // We can still finalize the stream even if row appends hit a closed bidirectional stream
128        writer.finalize().await?;
129
130        Ok(())
131    }
132
133    fn write_stream() -> String {
134        "projects/p/datasets/d/tables/t/streams/s".to_string()
135    }
136
137    fn schema() -> ArrowSchema {
138        ArrowSchema::new().set_serialized_schema("test")
139    }
140
141    fn rows(id: i64) -> ArrowRecordBatch {
142        ArrowRecordBatch::new().set_serialized_record_batch(id.to_string())
143    }
144}
google_cloud_bigquery/write/arrow/committed.rs

google_cloud_bigquery/write/arrow/
committed.rs
