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::client_builder::ClientBuilder;
16use super::pool::{StreamPool, StreamPoolOptions};
17use super::retry_policy::RetryOptions;
18use super::stream_type::{ApplicationCreatedStream, DefaultStream};
19use super::transport::Transport;
20use super::writer_builder::WriterBuilder;
21use crate::ClientBuilderResult as BuilderResult;
22use std::collections::HashMap;
23use std::sync::{Arc, Mutex};
24
25/// A client for BigQuery Storage Write API.
26#[derive(Debug)]
27pub struct Write {
28    inner: Arc<Transport>,
29    pools: Arc<Mutex<HashMap<String, Arc<StreamPool>>>>,
30    pool_options: StreamPoolOptions,
31    retry_options: RetryOptions,
32}
33
34impl Write {
35    /// Creates a new [ClientBuilder].
36    pub fn builder() -> ClientBuilder {
37        ClientBuilder::new()
38    }
39
40    pub(crate) async fn new(builder: ClientBuilder) -> BuilderResult<Self> {
41        let inner = Arc::new(Transport::new(builder.config).await?);
42        let pools = Arc::new(Mutex::new(HashMap::new()));
43        Ok(Self {
44            inner,
45            pools,
46            pool_options: builder.pool_options,
47            retry_options: builder.retry_options,
48        })
49    }
50
51    /// Opens the [default stream] for the given table.
52    ///
53    /// # Example
54    /// ```
55    /// # use google_cloud_bigquery::client::Write;
56    /// # async fn sample(client: Write) -> anyhow::Result<()> {
57    /// let writer = client
58    ///     .open_default_stream("projects/my-project/datasets/my-dataset/tables/my-table")
59    ///     .build_arrow(schema())
60    ///     .await?;
61    /// # Ok(()) }
62    ///
63    /// use google_cloud_bigquery::model::ArrowSchema;
64    /// fn schema() -> ArrowSchema {
65    ///   todo!("Define your table's schema...")
66    /// }
67    /// ```
68    ///
69    /// [default stream]: https://docs.cloud.google.com/bigquery/docs/write-api#default_stream
70    pub fn open_default_stream<T: Into<String>>(&self, table: T) -> WriterBuilder<DefaultStream> {
71        WriterBuilder::new_open_default(
72            self.inner.clone(),
73            self.pools.clone(),
74            self.pool_options.clone(),
75            self.retry_options.clone(),
76            table.into(),
77        )
78    }
79
80    /// Creates a new [application-created stream] for the given table.
81    ///
82    /// The stream type `S` can be inferred from the variable's writer type
83    /// annotation
84    /// ([`PendingWriter`][crate::write::PendingWriter],
85    /// [`CommittedWriter`][crate::write::CommittedWriter], or
86    /// [`BufferedWriter`][crate::write::BufferedWriter]) or specified explicitly via turbofish
87    /// (`create_stream::<PendingStream, _>(...)`).
88    ///
89    /// # Example
90    /// ```
91    /// use google_cloud_bigquery::write::PendingWriter;
92    /// use google_cloud_bigquery::write::format::Arrow;
93    /// # use google_cloud_bigquery::client::Write;
94    /// # async fn sample(client: Write) -> anyhow::Result<()> {
95    /// let writer: PendingWriter<Arrow> = client
96    ///     .create_stream("projects/my-project/datasets/my-dataset/tables/my-table")
97    ///     .build_arrow(schema())
98    ///     .await?;
99    /// # Ok(()) }
100    ///
101    /// use google_cloud_bigquery::model::ArrowSchema;
102    /// fn schema() -> ArrowSchema {
103    ///   todo!("Define your table's schema...")
104    /// }
105    /// ```
106    ///
107    /// [application-created stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#application-created_streams
108    pub fn create_stream<S: ApplicationCreatedStream, T: Into<String>>(
109        &self,
110        table: T,
111    ) -> WriterBuilder<S> {
112        WriterBuilder::new_create(self.inner.clone(), self.retry_options.clone(), table.into())
113    }
114
115    /// Attaches to an existing [application-created stream].
116    ///
117    /// The stream type `S` can be inferred from the variable's writer type
118    /// annotation
119    /// ([`PendingWriter`][crate::write::PendingWriter],
120    /// [`CommittedWriter`][crate::write::CommittedWriter], or
121    /// [`BufferedWriter`][crate::write::BufferedWriter]) or specified explicitly via turbofish
122    /// (`attach_to_stream::<PendingStream, _>(...)`).
123    ///
124    /// # Example
125    /// ```
126    /// use google_cloud_bigquery::write::CommittedWriter;
127    /// use google_cloud_bigquery::write::format::Arrow;
128    /// # use google_cloud_bigquery::client::Write;
129    /// # async fn sample(client: Write) -> anyhow::Result<()> {
130    /// let writer: CommittedWriter<Arrow> = client
131    ///     .attach_to_stream("projects/my-project/datasets/my_dataset/tables/my_table/streams/my_stream")
132    ///     .build_arrow(schema())
133    ///     .await?;
134    /// # Ok(())
135    /// # }
136    /// #
137    /// # use google_cloud_bigquery::model::ArrowSchema;
138    /// # fn schema() -> ArrowSchema {
139    /// #   todo!("Define your table's schema...")
140    /// # }
141    /// ```
142    ///
143    /// [application-created stream]: https://docs.cloud.google.com/bigquery/docs/write-api-grpc#application-created_streams
144    pub fn attach_to_stream<S: ApplicationCreatedStream, T: Into<String>>(
145        &self,
146        write_stream: T,
147    ) -> WriterBuilder<S> {
148        WriterBuilder::new_attach(
149            self.inner.clone(),
150            self.retry_options.clone(),
151            write_stream.into(),
152        )
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::super::error::AppendError;
159    use super::*;
160    use crate::model::{ArrowRecordBatch, ArrowSchema, ProtoRows, ProtoSchema};
161    use bigquery_grpc_mock::{MockBigQueryWrite, start};
162    use gaxi::grpc::tonic::Status as TonicStatus;
163    use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
164
165    #[tokio::test]
166    async fn arrow() -> anyhow::Result<()> {
167        let mut mock = MockBigQueryWrite::new();
168        mock.expect_append_rows()
169            .return_once(|_| Err(TonicStatus::failed_precondition("fail")));
170        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
171        let client = Write::builder()
172            .with_endpoint(endpoint)
173            .with_credentials(Anonymous::new().build())
174            .build()
175            .await?;
176        let writer = client
177            .open_default_stream("projects/p/datasets/d/tables/t")
178            .build_arrow(ArrowSchema::new())
179            .await?;
180        let err = writer
181            .append(ArrowRecordBatch::new())
182            .send()
183            .await
184            .expect_err("write should fail");
185        assert!(matches!(err, AppendError::Rpc { source: _ }));
186
187        Ok(())
188    }
189
190    #[tokio::test]
191    async fn proto() -> anyhow::Result<()> {
192        let mut mock = MockBigQueryWrite::new();
193        mock.expect_append_rows()
194            .return_once(|_| Err(TonicStatus::failed_precondition("fail")));
195        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
196        let client = Write::builder()
197            .with_endpoint(endpoint)
198            .with_credentials(Anonymous::new().build())
199            .build()
200            .await?;
201        let writer = client
202            .open_default_stream("projects/p/datasets/d/tables/t")
203            .build_proto(ProtoSchema::new())
204            .await?;
205        let err = writer
206            .append(ProtoRows::new())
207            .send()
208            .await
209            .expect_err("write should fail");
210        assert!(matches!(err, AppendError::Rpc { source: _ }));
211
212        Ok(())
213    }
214
215    #[tokio::test]
216    async fn multiplexing() -> anyhow::Result<()> {
217        let mut mock = MockBigQueryWrite::new();
218        mock.expect_get_write_stream().times(2).returning(|req| {
219            let name = req.into_inner().name;
220            Ok(gaxi::grpc::tonic::Response::new(
221                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::WriteStream {
222                    name,
223                    location: "us".to_string(),
224                    ..Default::default()
225                },
226            ))
227        });
228        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
229        let client = Write::builder()
230            .with_endpoint(endpoint)
231            .with_credentials(Anonymous::new().build())
232            .build()
233            .await?;
234        let multiplexed_writer1 = client
235            .open_default_stream("projects/p/datasets/d/tables/t1")
236            .with_multiplexing(true)
237            .build_arrow(ArrowSchema::new())
238            .await?;
239        let multiplexed_writer2 = client
240            .open_default_stream("projects/p/datasets/d/tables/t2")
241            .with_multiplexing(true)
242            .build_arrow(ArrowSchema::new())
243            .await?;
244        assert!(Arc::ptr_eq(
245            &multiplexed_writer1.inner.pool,
246            &multiplexed_writer2.inner.pool
247        ));
248
249        let standalone_writer = client
250            .open_default_stream("projects/p/datasets/d/tables/t3")
251            .with_multiplexing(false)
252            .build_arrow(ArrowSchema::new())
253            .await?;
254        assert!(!Arc::ptr_eq(
255            &multiplexed_writer1.inner.pool,
256            &standalone_writer.inner.pool
257        ));
258
259        Ok(())
260    }
261
262    #[tokio::test]
263    async fn format_isolation() -> anyhow::Result<()> {
264        let mut mock = MockBigQueryWrite::new();
265        mock.expect_get_write_stream().times(2).returning(|req| {
266            let name = req.into_inner().name;
267            Ok(gaxi::grpc::tonic::Response::new(
268                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::WriteStream {
269                    name,
270                    location: "us".to_string(),
271                    ..Default::default()
272                },
273            ))
274        });
275        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
276        let client = Write::builder()
277            .with_endpoint(endpoint)
278            .with_credentials(Anonymous::new().build())
279            .build()
280            .await?;
281        let arrow_writer = client
282            .open_default_stream("projects/p/datasets/d/tables/t1")
283            .with_multiplexing(true)
284            .build_arrow(ArrowSchema::new())
285            .await?;
286        let proto_writer = client
287            .open_default_stream("projects/p/datasets/d/tables/t2")
288            .with_multiplexing(true)
289            .build_proto(ProtoSchema::new())
290            .await?;
291
292        // Different formats receive distinct connection pools.
293        assert!(!Arc::ptr_eq(
294            &arrow_writer.inner.pool,
295            &proto_writer.inner.pool
296        ));
297
298        Ok(())
299    }
300
301    #[tokio::test]
302    async fn location_isolation() -> anyhow::Result<()> {
303        let mut mock = MockBigQueryWrite::new();
304        mock.expect_get_write_stream().times(3).returning(|req| {
305            let name = req.into_inner().name;
306            let location = if name.contains("t1") {
307                "us".to_string()
308            } else if name.contains("t2") {
309                "eu".to_string()
310            } else {
311                "us".to_string()
312            };
313            Ok(gaxi::grpc::tonic::Response::new(
314                bigquery_grpc_mock::google::cloud::bigquery::storage::v1::WriteStream {
315                    name,
316                    location,
317                    ..Default::default()
318                },
319            ))
320        });
321        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
322        let client = Write::builder()
323            .with_endpoint(endpoint)
324            .with_credentials(Anonymous::new().build())
325            .build()
326            .await?;
327        let us_writer1 = client
328            .open_default_stream("projects/p/datasets/d/tables/t1")
329            .with_multiplexing(true)
330            .build_arrow(ArrowSchema::new())
331            .await?;
332        let eu_writer = client
333            .open_default_stream("projects/p/datasets/d/tables/t2")
334            .with_multiplexing(true)
335            .build_arrow(ArrowSchema::new())
336            .await?;
337
338        // Different locations receive distinct connection pools.
339        assert!(!Arc::ptr_eq(&us_writer1.inner.pool, &eu_writer.inner.pool));
340
341        let us_writer2 = client
342            .open_default_stream("projects/p/datasets/d/tables/t3")
343            .with_multiplexing(true)
344            .build_arrow(ArrowSchema::new())
345            .await?;
346
347        // Same location shares the connection pool.
348        assert!(Arc::ptr_eq(&us_writer1.inner.pool, &us_writer2.inner.pool));
349
350        Ok(())
351    }
352
353    #[tokio::test]
354    async fn retry_options() -> anyhow::Result<()> {
355        let client = Write::builder()
356            .with_credentials(Anonymous::new().build())
357            .build()
358            .await?;
359        let writer = client
360            .open_default_stream("projects/p/datasets/d/tables/t")
361            .build_arrow(ArrowSchema::new())
362            .await?;
363
364        // The writer uses the client's policies, not a fresh set of defaults.
365        let options = &writer.inner.options;
366        assert!(Arc::ptr_eq(
367            &client.retry_options.retry_policy,
368            &options.retry_policy
369        ));
370        assert!(Arc::ptr_eq(
371            &client.retry_options.backoff_policy,
372            &options.backoff_policy
373        ));
374        assert_eq!(
375            client.retry_options.attempt_timeout,
376            options.attempt_timeout
377        );
378
379        Ok(())
380    }
381}
google_cloud_bigquery/write/client.rs

google_cloud_bigquery/write/
client.rs
