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

google_cloud_bigquery/write/
pending.rs
