google_cloud_bigquery/write/arrow/
buffered.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, FlushRowsResponse};
18use crate::write::append_builder::AppendWithOffset;
19use crate::write::transport::Transport;
20use std::sync::Arc;
21
22/// A writer for a [buffered stream].
23///
24/// [buffered stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#buffered_type
25#[derive(Debug)]
26pub struct BufferedWriter {
27    pub(crate) inner: BaseWriter,
28}
29
30impl BufferedWriter {
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 buffered 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    /// Flush the buffered stream, making rows up to the specified offset available for reading.
51    pub async fn flush(&self, offset: i64) -> Result<FlushRowsResponse> {
52        self.inner
53            .client
54            .flush_rows()
55            .set_write_stream(&self.inner.write_stream)
56            .set_offset(offset)
57            .send()
58            .await
59    }
60
61    /// Finalize the buffered stream, preventing further writes.
62    pub async fn finalize(&self) -> Result<FinalizeWriteStreamResponse> {
63        self.inner.finalize().await
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::super::super::runner::tests::*;
70    use super::super::super::transport::tests::*;
71    use super::*;
72    use crate::error::AppendError;
73    use bigquery_grpc_mock::{MockBigQueryWrite, start};
74    use gaxi::grpc::tonic::Response as TonicResponse;
75    use tokio::sync::mpsc;
76
77    #[tokio::test]
78    async fn request_fields() -> anyhow::Result<()> {
79        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
80        let writer = BufferedWriter::new(transport, write_stream(), schema());
81        assert_eq!(writer.write_stream(), write_stream());
82
83        let b = writer.append(rows(1));
84        assert_eq!(b.req.write_stream, write_stream());
85        let data = b.req.arrow_rows().expect("arrow rows should be set");
86        let s = data.writer_schema.as_ref().expect("schema should be set");
87        assert_eq!(s.serialized_schema, "test");
88        let r = data.rows.as_ref().expect("rows should be set");
89        assert_eq!(r.serialized_record_batch, "1");
90
91        let b = writer.append(rows(2));
92        assert_eq!(b.req.write_stream, write_stream());
93        let data = b.req.arrow_rows().expect("arrow rows should be set");
94        let s = data.writer_schema.as_ref().expect("schema should be set");
95        assert_eq!(s.serialized_schema, "test");
96        let r = data.rows.as_ref().expect("rows should be set");
97        assert_eq!(r.serialized_record_batch, "2");
98
99        Ok(())
100    }
101
102    #[tokio::test]
103    async fn basic_success() -> anyhow::Result<()> {
104        let (response_tx, response_rx) = mpsc::channel(10);
105
106        let mut mock = MockBigQueryWrite::new();
107        mock.expect_append_rows()
108            .return_once(|_| Ok(TonicResponse::from(response_rx)));
109
110        mock.expect_flush_rows()
111            .return_once(|req| {
112                assert_eq!(req.get_ref().offset, Some(3));
113                assert_eq!(req.get_ref().write_stream, write_stream());
114                Ok(TonicResponse::new(
115                    bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FlushRowsResponse::default()
116                ))
117            });
118
119        mock.expect_finalize_write_stream()
120            .return_once(|req| {
121                assert_eq!(req.get_ref().name, write_stream());
122                Ok(TonicResponse::new(
123                    bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FinalizeWriteStreamResponse::default()
124                ))
125            });
126
127        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
128        let transport = Arc::new(test_transport(endpoint).await?);
129
130        let writer = BufferedWriter::new(transport, write_stream(), schema());
131        assert_eq!(writer.write_stream(), write_stream());
132
133        response_tx.send(Ok(convert(&test_response(1)))).await?;
134        let resp = writer.append(rows(1)).send().await?;
135        assert_eq!(resp.offset, Some(1));
136
137        response_tx.send(Ok(convert(&test_response(2)))).await?;
138        let resp = writer.append(rows(2)).send().await?;
139        assert_eq!(resp.offset, Some(2));
140
141        response_tx.send(Ok(convert(&test_response(3)))).await?;
142        let resp = writer.append(rows(3)).send().await?;
143        assert_eq!(resp.offset, Some(3));
144
145        drop(response_tx);
146        let err = writer.append(rows(4)).send().await.expect_err("channel");
147        assert!(matches!(err, AppendError::UnexpectedEndOfStream));
148
149        writer.flush(3).await?;
150        writer.finalize().await?;
151
152        Ok(())
153    }
154
155    #[tokio::test]
156    async fn multiple_flushes() -> anyhow::Result<()> {
157        let (response_tx, response_rx) = mpsc::channel(10);
158        let mut mock = MockBigQueryWrite::new();
159        mock.expect_append_rows()
160            .return_once(|_| Ok(TonicResponse::from(response_rx)));
161
162        mock.expect_flush_rows().times(2).returning(|req| {
163            Ok(TonicResponse::new(
164                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FlushRowsResponse {
165                    offset: req.get_ref().offset.unwrap_or(0),
166                },
167            ))
168        });
169
170        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
171        let transport = Arc::new(test_transport(endpoint).await?);
172        let writer = BufferedWriter::new(transport, write_stream(), schema());
173        assert_eq!(writer.write_stream(), write_stream());
174
175        response_tx.send(Ok(convert(&test_response(1)))).await?;
176        let _ = writer.append(rows(1)).send().await?;
177        let flush1 = writer.flush(1).await?;
178        assert_eq!(flush1.offset, 1);
179
180        response_tx.send(Ok(convert(&test_response(2)))).await?;
181        let _ = writer.append(rows(2)).send().await?;
182        let flush2 = writer.flush(2).await?;
183        assert_eq!(flush2.offset, 2);
184
185        Ok(())
186    }
187
188    fn write_stream() -> String {
189        "projects/p/datasets/d/tables/t/streams/s".to_string()
190    }
191
192    fn schema() -> ArrowSchema {
193        ArrowSchema::new().set_serialized_schema("test")
194    }
195
196    fn rows(id: i64) -> ArrowRecordBatch {
197        ArrowRecordBatch::new().set_serialized_record_batch(id.to_string())
198    }
199}
google_cloud_bigquery/write/arrow/buffered.rs

google_cloud_bigquery/write/arrow/
buffered.rs
