google_cloud_bigquery/write/
client.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::arrow::WriterBuilder as ArrowWriterBuilder;
16use super::client_builder::ClientBuilder;
17use super::transport::Transport;
18use crate::ClientBuilderResult as BuilderResult;
19use crate::model::ArrowSchema;
20use std::sync::Arc;
21
22/// A client for BigQuery Storage Write API.
23#[derive(Debug)]
24pub struct Write {
25    #[allow(unused)]
26    inner: Arc<Transport>,
27}
28
29impl Write {
30    /// Creates a new [ClientBuilder].
31    pub fn builder() -> ClientBuilder {
32        ClientBuilder::new()
33    }
34
35    pub(crate) async fn new(builder: ClientBuilder) -> BuilderResult<Self> {
36        let transport = Transport::new(builder.config).await?;
37        Ok(Self {
38            inner: Arc::new(transport),
39        })
40    }
41
42    /// Create a writer using [Arrow] as the data format.
43    ///
44    /// # Example
45    /// ```
46    /// # use google_cloud_bigquery::client::Write;
47    /// # async fn sample(client: Write) -> anyhow::Result<()> {
48    /// let writer = client
49    ///   .arrow(schema())
50    ///   .default("projects/my-project/datasets/my-dataset/tables/my-table")?;
51    /// # Ok(()) }
52    ///
53    /// use google_cloud_bigquery::model::ArrowSchema;
54    /// fn schema() -> ArrowSchema {
55    ///   todo!("Define your table's schema...")
56    /// }
57    /// ```
58    ///
59    /// [arrow]: https://arrow.apache.org/
60    pub fn arrow(&self, schema: ArrowSchema) -> ArrowWriterBuilder {
61        ArrowWriterBuilder::new(self.inner.clone(), schema)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::super::error::AppendError;
68    use super::*;
69    use crate::model::{ArrowRecordBatch, ArrowSchema};
70    use bigquery_grpc_mock::{MockBigQueryWrite, start};
71    use gaxi::grpc::tonic::Status as TonicStatus;
72    use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
73
74    #[tokio::test]
75    async fn arrow() -> anyhow::Result<()> {
76        let mut mock = MockBigQueryWrite::new();
77        mock.expect_append_rows()
78            .return_once(|_| Err(TonicStatus::failed_precondition("fail")));
79        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
80        let client = Write::builder()
81            .with_endpoint(endpoint)
82            .with_credentials(Anonymous::new().build())
83            .build()
84            .await?;
85        let writer = client
86            .arrow(ArrowSchema::new())
87            .default("projects/p/datasets/d/tables/t")?;
88        let err = writer
89            .append(ArrowRecordBatch::new())
90            .send()
91            .await
92            .expect_err("write should fail");
93        assert!(matches!(err, AppendError::Rpc { source: _ }));
94
95        Ok(())
96    }
97}
google_cloud_bigquery/write/client.rs

google_cloud_bigquery/write/
client.rs
