google_cloud_bigquery/write/
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 super::error::CommitError;
17use super::format::DataFormat;
18use crate::Result;
19use crate::model::{BatchCommitWriteStreamsResponse, FinalizeWriteStreamResponse};
20use crate::write::builder::AppendWithOffset;
21use crate::write::transport::Transport;
22use std::sync::Arc;
23
24/// A writer for a [pending stream].
25///
26/// In a pending stream, records are buffered in a pending state until you
27/// commit the stream. When you commit a stream, all of the pending data
28/// becomes available for reading atomically. Use this type for batch
29/// workloads, as an alternative to BigQuery load jobs.
30///
31/// [pending stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#pending_type
32#[derive(Debug)]
33pub struct PendingWriter<F> {
34    pub(crate) inner: BaseWriter<F>,
35}
36
37impl<F> PendingWriter<F>
38where
39    F: DataFormat,
40{
41    pub(crate) fn new(inner: Arc<Transport>, write_stream: String, format: F) -> Self {
42        Self {
43            inner: BaseWriter::new(inner, write_stream, format),
44        }
45    }
46
47    /// Returns the full resource name of the underlying write stream.
48    pub fn write_stream(&self) -> &str {
49        &self.inner.write_stream
50    }
51
52    /// Appends rows to the pending stream.
53    pub fn append(&self, rows: F::Rows) -> AppendWithOffset {
54        AppendWithOffset::new(
55            self.inner.runner.req_tx.clone(),
56            self.inner.append_request(rows),
57        )
58    }
59
60    /// Finalizes the pending stream, preventing further writes.
61    pub async fn finalize(&self) -> Result<FinalizeWriteStreamResponse> {
62        self.inner.finalize().await
63    }
64
65    /// Commits the pending stream to the table.
66    ///
67    /// The stream must be finalized via [`finalize`][Self::finalize] before
68    /// calling `commit`.
69    pub async fn commit(
70        &self,
71    ) -> std::result::Result<BatchCommitWriteStreamsResponse, CommitError> {
72        // Extract the parent table path from the stream name:
73        // "projects/p/datasets/d/tables/t/streams/s" -> "projects/p/datasets/d/tables/t"
74        let parent = self
75            .inner
76            .write_stream
77            .split_once("/streams/")
78            .map_or(self.inner.write_stream.as_str(), |(p, _)| p)
79            .to_string();
80
81        let resp = self
82            .inner
83            .client
84            .batch_commit_write_streams()
85            .set_parent(parent)
86            .set_write_streams(vec![self.inner.write_stream.clone()])
87            .send()
88            .await?;
89
90        if !resp.stream_errors.is_empty() || resp.commit_time.is_none() {
91            return Err(CommitError::FailedTransaction {
92                stream_errors: resp.stream_errors,
93            });
94        }
95
96        Ok(resp)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::model::storage_error::StorageErrorCode;
104    use crate::write::test::*;
105    use bigquery_grpc_mock::google::cloud::bigquery::storage::v1;
106    use bigquery_grpc_mock::{MockBigQueryWrite, start};
107    use gaxi::grpc::tonic::{Response as TonicResponse, Status as TonicStatus};
108    use google_cloud_gax::error::rpc::Code;
109    use tokio::sync::mpsc;
110
111    #[tokio::test]
112    async fn basic_success() -> anyhow::Result<()> {
113        let (response_tx, response_rx) = mpsc::channel(10);
114
115        let mut mock = MockBigQueryWrite::new();
116        mock.expect_append_rows()
117            .return_once(|_| Ok(TonicResponse::from(response_rx)));
118
119        mock.expect_finalize_write_stream().return_once(|_| {
120            Ok(TonicResponse::new(
121                v1::FinalizeWriteStreamResponse::default(),
122            ))
123        });
124
125        mock.expect_batch_commit_write_streams().return_once(|_| {
126            Ok(TonicResponse::new(v1::BatchCommitWriteStreamsResponse {
127                commit_time: Some(prost_types::Timestamp {
128                    seconds: 1_700_000_000,
129                    nanos: 0,
130                }),
131                stream_errors: Vec::new(),
132            }))
133        });
134
135        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
136        let transport = Arc::new(test_transport(endpoint).await?);
137
138        let writer = PendingWriter::new(transport, write_stream(), format());
139        assert_eq!(writer.write_stream(), write_stream());
140
141        response_tx.send(Ok(convert(&test_response(1)))).await?;
142        let resp = writer.append(rows(1)).send().await?;
143        assert_eq!(resp.offset, Some(1));
144
145        writer.finalize().await?;
146        let commit_resp = writer.commit().await?;
147        assert!(commit_resp.commit_time.is_some());
148
149        Ok(())
150    }
151
152    #[tokio::test]
153    async fn commit_stream_errors() -> anyhow::Result<()> {
154        let mut mock = MockBigQueryWrite::new();
155        mock.expect_batch_commit_write_streams().return_once(|_| {
156            Ok(TonicResponse::new(v1::BatchCommitWriteStreamsResponse {
157                commit_time: None,
158                stream_errors: vec![v1::StorageError {
159                    code: v1::storage_error::StorageErrorCode::InvalidStreamState as i32,
160                    entity: write_stream(),
161                    error_message: "stream is not finalized".to_string(),
162                }],
163            }))
164        });
165
166        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
167        let transport = Arc::new(test_transport(endpoint).await?);
168        let writer = PendingWriter::new(transport, write_stream(), format());
169
170        let err = writer.commit().await.expect_err("commit should fail");
171        let CommitError::FailedTransaction { stream_errors } = err else {
172            anyhow::bail!("expected FailedTransaction, got: {err:?}");
173        };
174        assert_eq!(stream_errors.len(), 1);
175        assert_eq!(stream_errors[0].code, StorageErrorCode::InvalidStreamState);
176        assert_eq!(stream_errors[0].entity, write_stream());
177        assert_eq!(stream_errors[0].error_message, "stream is not finalized");
178
179        Ok(())
180    }
181
182    #[tokio::test]
183    async fn commit_missing_commit_time() -> anyhow::Result<()> {
184        let mut mock = MockBigQueryWrite::new();
185        mock.expect_batch_commit_write_streams().return_once(|_| {
186            Ok(TonicResponse::new(
187                v1::BatchCommitWriteStreamsResponse::default(),
188            ))
189        });
190
191        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
192        let transport = Arc::new(test_transport(endpoint).await?);
193        let writer = PendingWriter::new(transport, write_stream(), format());
194
195        let err = writer.commit().await.expect_err("commit should fail");
196        let CommitError::FailedTransaction { stream_errors } = err else {
197            anyhow::bail!("expected FailedTransaction, got: {err:?}");
198        };
199        assert!(stream_errors.is_empty());
200
201        Ok(())
202    }
203
204    #[tokio::test]
205    async fn commit_rpc_error() -> anyhow::Result<()> {
206        let mut mock = MockBigQueryWrite::new();
207        mock.expect_batch_commit_write_streams()
208            .return_once(|_| Err(TonicStatus::failed_precondition("fail")));
209
210        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
211        let transport = Arc::new(test_transport(endpoint).await?);
212        let writer = PendingWriter::new(transport, write_stream(), format());
213
214        let err = writer.commit().await.expect_err("commit should fail");
215        let CommitError::Rpc { source } = err else {
216            anyhow::bail!("expected Rpc error, got: {err:?}");
217        };
218        let status = source.status().expect("error should have status");
219        assert_eq!(status.code, Code::FailedPrecondition);
220        assert_eq!(status.message, "fail");
221
222        Ok(())
223    }
224}
google_cloud_bigquery/write/pending.rs

google_cloud_bigquery/write/
pending.rs
