google_cloud_bigquery/write/arrow/
pending.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::{
18    ArrowRecordBatch, ArrowSchema, BatchCommitWriteStreamsResponse, FinalizeWriteStreamResponse,
19};
20use crate::write::append_builder::AppendWithOffset;
21use crate::write::transport::Transport;
22use std::sync::Arc;
23
24/// A writer for a [pending stream].
25///
26/// [pending stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#pending_type
27#[derive(Debug)]
28pub struct PendingWriter {
29    pub(crate) inner: BaseWriter,
30}
31
32impl PendingWriter {
33    pub(crate) fn new(inner: Arc<Transport>, write_stream: String, schema: ArrowSchema) -> Self {
34        Self {
35            inner: BaseWriter::new(inner, write_stream, schema),
36        }
37    }
38
39    /// Returns the full resource name of the underlying write stream.
40    pub fn write_stream(&self) -> &str {
41        &self.inner.write_stream
42    }
43
44    /// Appends rows to the pending stream.
45    pub fn append(&self, rows: ArrowRecordBatch) -> AppendWithOffset {
46        AppendWithOffset::new(
47            self.inner.runner.req_tx.clone(),
48            self.inner.append_request(rows),
49        )
50    }
51
52    /// Finalizes the pending stream, preventing further writes.
53    pub async fn finalize(&self) -> Result<FinalizeWriteStreamResponse> {
54        self.inner.finalize().await
55    }
56
57    /// Commits the pending stream to the table.
58    pub async fn commit(&self) -> Result<BatchCommitWriteStreamsResponse> {
59        // Extract the parent table path from the stream name:
60        // "projects/p/datasets/d/tables/t/streams/s" -> "projects/p/datasets/d/tables/t"
61        let parent = self
62            .inner
63            .write_stream
64            .split_once("/streams/")
65            .map_or(self.inner.write_stream.as_str(), |(p, _)| p)
66            .to_string();
67
68        self.inner
69            .client
70            .batch_commit_write_streams()
71            .set_parent(parent)
72            .set_write_streams(vec![self.inner.write_stream.clone()])
73            .send()
74            .await
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::super::super::runner::tests::*;
81    use super::super::super::transport::tests::*;
82    use super::*;
83    use bigquery_grpc_mock::{MockBigQueryWrite, start};
84    use gaxi::grpc::tonic::Response as TonicResponse;
85    use tokio::sync::mpsc;
86
87    #[tokio::test]
88    async fn request_fields() -> anyhow::Result<()> {
89        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
90        let writer = PendingWriter::new(transport, write_stream(), schema());
91        assert_eq!(writer.write_stream(), write_stream());
92
93        let b = writer.append(rows(1));
94        assert_eq!(b.req.write_stream, write_stream());
95        let data = b.req.arrow_rows().expect("arrow rows should be set");
96        let s = data.writer_schema.as_ref().expect("schema should be set");
97        assert_eq!(s.serialized_schema, "test");
98        let r = data.rows.as_ref().expect("rows should be set");
99        assert_eq!(r.serialized_record_batch, "1");
100
101        Ok(())
102    }
103
104    #[tokio::test]
105    async fn basic_success() -> anyhow::Result<()> {
106        let (response_tx, response_rx) = mpsc::channel(10);
107
108        let mut mock = MockBigQueryWrite::new();
109        mock.expect_append_rows()
110            .return_once(|_| Ok(TonicResponse::from(response_rx)));
111
112        mock.expect_finalize_write_stream()
113            .return_once(|_| Ok(TonicResponse::new(
114                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::FinalizeWriteStreamResponse::default()
115            )));
116
117        mock.expect_batch_commit_write_streams()
118            .return_once(|_| Ok(TonicResponse::new(
119                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::BatchCommitWriteStreamsResponse::default()
120            )));
121
122        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
123        let transport = Arc::new(test_transport(endpoint).await?);
124
125        let writer = PendingWriter::new(transport, write_stream(), schema());
126        assert_eq!(writer.write_stream(), write_stream());
127
128        response_tx.send(Ok(convert(&test_response(1)))).await?;
129        let resp = writer.append(rows(1)).send().await?;
130        assert_eq!(resp.offset, Some(1));
131
132        writer.finalize().await?;
133        writer.commit().await?;
134
135        Ok(())
136    }
137
138    fn write_stream() -> String {
139        "projects/p/datasets/d/tables/t/streams/s".to_string()
140    }
141
142    fn schema() -> ArrowSchema {
143        ArrowSchema::new().set_serialized_schema("test")
144    }
145
146    fn rows(id: i64) -> ArrowRecordBatch {
147        ArrowRecordBatch::new().set_serialized_record_batch(id.to_string())
148    }
149}
google_cloud_bigquery/write/arrow/pending.rs

google_cloud_bigquery/write/arrow/
pending.rs
