google_cloud_bigquery/write/arrow/
writer_builder.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::super::generated::gapic_storage::client::BigQueryWrite;
16use super::super::transport::Transport;
17use super::{BufferedWriter, CommittedWriter, DefaultWriter, PendingWriter, Writer};
18use crate::model::write_stream::Type;
19use crate::model::{ArrowSchema, WriteStream};
20use crate::write::error::{AttachError, AttachResult};
21use crate::{Error, Result};
22use gaxi::path_parameter::{PathMismatchBuilder, try_match};
23use gaxi::routing_parameter::Segment;
24use google_cloud_gax::error::binding::BindingError;
25use std::sync::Arc;
26
27/// A builder to create a stream writer
28#[derive(Clone, Debug)]
29pub struct WriterBuilder {
30    inner: Arc<Transport>,
31    schema: ArrowSchema,
32}
33
34impl WriterBuilder {
35    pub(crate) fn new(inner: Arc<Transport>, schema: ArrowSchema) -> Self {
36        Self { inner, schema }
37    }
38
39    /// Create a writer for the [default stream] for the given table.
40    ///
41    /// # Example
42    ///
43    /// ```
44    /// # use google_cloud_bigquery::client::Write;
45    /// # async fn sample(client: Write) -> anyhow::Result<()> {
46    /// let writer = client
47    ///     .arrow(schema())
48    ///     .default("projects/my-project/datasets/my-dataset/tables/my-table")?;
49    /// # Ok(()) }
50    ///
51    /// use google_cloud_bigquery::model::ArrowSchema;
52    /// fn schema() -> ArrowSchema {
53    ///   todo!("Define your table's schema...")
54    /// }
55    /// ```
56    ///
57    /// [default stream]: https://docs.cloud.google.com/bigquery/docs/write-api#default_stream
58    pub fn default<T: Into<String>>(self, table: T) -> Result<DefaultWriter> {
59        let table = table.into();
60        validate_table(table.as_str())?;
61        let mut write_stream = table;
62        write_stream.push_str("/streams/_default");
63        Ok(DefaultWriter::new(self.inner, write_stream, self.schema))
64    }
65
66    /// Creates a pending writer for the given table.
67    ///
68    /// # Example
69    ///
70    /// ```
71    /// # use google_cloud_bigquery::client::Write;
72    /// # async fn sample(client: Write) -> anyhow::Result<()> {
73    /// let writer = client
74    ///     .arrow(schema())
75    ///     .pending("projects/my-project/datasets/my-dataset/tables/my-table")
76    ///     .await?;
77    /// # Ok(()) }
78    ///
79    /// use google_cloud_bigquery::model::ArrowSchema;
80    /// fn schema() -> ArrowSchema {
81    ///   todo!("Define your table's schema...")
82    /// }
83    /// ```
84    pub async fn pending<T: Into<String>>(self, table: T) -> Result<PendingWriter> {
85        let table = table.into();
86        validate_table(table.as_str())?;
87
88        let client = BigQueryWrite::from_stub::<Transport>(self.inner.clone());
89        let write_stream = client
90            .create_write_stream()
91            .set_parent(table)
92            .set_write_stream(WriteStream::new().set_type(Type::Pending))
93            .send()
94            .await?;
95
96        Ok(PendingWriter::new(
97            self.inner,
98            write_stream.name,
99            self.schema,
100        ))
101    }
102
103    /// Creates a committed writer for the given table.
104    ///
105    /// # Example
106    ///
107    /// ```
108    /// # use google_cloud_bigquery::client::Write;
109    /// # async fn sample(client: Write) -> anyhow::Result<()> {
110    /// let writer = client
111    ///     .arrow(schema())
112    ///     .committed("projects/my-project/datasets/my-dataset/tables/my-table")
113    ///     .await?;
114    /// # Ok(()) }
115    ///
116    /// use google_cloud_bigquery::model::ArrowSchema;
117    /// fn schema() -> ArrowSchema {
118    ///   todo!("Define your table's schema...")
119    /// }
120    /// ```
121    pub async fn committed<T: Into<String>>(self, table: T) -> Result<CommittedWriter> {
122        let table = table.into();
123        validate_table(table.as_str())?;
124
125        let client = BigQueryWrite::from_stub::<Transport>(self.inner.clone());
126        let write_stream = client
127            .create_write_stream()
128            .set_parent(table)
129            .set_write_stream(WriteStream::new().set_type(Type::Committed))
130            .send()
131            .await?;
132
133        Ok(CommittedWriter::new(
134            self.inner,
135            write_stream.name,
136            self.schema,
137        ))
138    }
139
140    /// Creates a buffered writer for the given table.
141    ///
142    /// # Example
143    ///
144    /// ```
145    /// # use google_cloud_bigquery::client::Write;
146    /// # async fn sample(client: Write) -> anyhow::Result<()> {
147    /// let writer = client
148    ///     .arrow(schema())
149    ///     .buffered("projects/my-project/datasets/my-dataset/tables/my-table")
150    ///     .await?;
151    /// # Ok(()) }
152    ///
153    /// use google_cloud_bigquery::model::ArrowSchema;
154    /// fn schema() -> ArrowSchema {
155    ///   todo!("Define your table's schema...")
156    /// }
157    /// ```
158    pub async fn buffered<T: Into<String>>(self, table: T) -> Result<BufferedWriter> {
159        let table = table.into();
160        validate_table(table.as_str())?;
161
162        let client = BigQueryWrite::from_stub::<Transport>(self.inner.clone());
163        let write_stream = client
164            .create_write_stream()
165            .set_parent(table)
166            .set_write_stream(WriteStream::new().set_type(Type::Buffered))
167            .send()
168            .await?;
169
170        Ok(BufferedWriter::new(
171            self.inner,
172            write_stream.name,
173            self.schema,
174        ))
175    }
176
177    /// Attaches the builder to an existing stream.
178    ///
179    /// # Example
180    /// ```
181    /// use google_cloud_bigquery::write::arrow::CommittedWriter;
182    /// # use google_cloud_bigquery::client::Write;
183    /// # async fn sample(client: Write) -> anyhow::Result<()> {
184    /// let writer: CommittedWriter = client
185    ///     .arrow(schema())
186    ///     .attach("projects/my-project/datasets/my_dataset/tables/my_table/streams/my_stream")
187    ///     .await?;
188    /// # Ok(())
189    /// # }
190    /// #
191    /// # use google_cloud_bigquery::model::ArrowSchema;
192    /// # fn schema() -> ArrowSchema {
193    /// #   todo!("Define your table's schema...")
194    /// # }
195    /// ```
196    pub async fn attach<U: Writer, S: Into<String>>(self, write_stream: S) -> AttachResult<U> {
197        let write_stream = write_stream.into();
198        validate_stream(write_stream.as_str())?;
199
200        let client = BigQueryWrite::from_stub::<Transport>(self.inner.clone());
201        let stream = client
202            .get_write_stream()
203            .set_name(&write_stream)
204            .send()
205            .await?;
206
207        let stream_type = stream.r#type.clone();
208        if stream_type != U::STREAM_TYPE {
209            return Err(AttachError::TypeMismatch {
210                expected: U::STREAM_TYPE,
211                actual: stream_type,
212            });
213        }
214        Ok(U::build(self.inner, write_stream, self.schema))
215    }
216}
217
218fn validate_table(table: &str) -> Result<()> {
219    let segments = &[
220        Segment::Literal("projects/"),
221        Segment::SingleWildcard,
222        Segment::Literal("/datasets/"),
223        Segment::SingleWildcard,
224        Segment::Literal("/tables/"),
225        Segment::SingleWildcard,
226    ];
227    try_match(Some(table), segments)
228        .ok_or_else(|| {
229            let builder = PathMismatchBuilder::default().maybe_add(
230                Some(table),
231                segments,
232                "table",
233                "projects/*/datasets/*/tables/*",
234            );
235            Error::binding(BindingError {
236                paths: vec![builder.build()],
237            })
238        })
239        .map(|_| ())
240}
241
242fn validate_stream(stream: &str) -> crate::Result<()> {
243    let segments = &[
244        Segment::Literal("projects/"),
245        Segment::SingleWildcard,
246        Segment::Literal("/datasets/"),
247        Segment::SingleWildcard,
248        Segment::Literal("/tables/"),
249        Segment::SingleWildcard,
250        Segment::Literal("/streams/"),
251        Segment::SingleWildcard,
252    ];
253    try_match(Some(stream), segments)
254        .ok_or_else(|| {
255            let builder = gaxi::path_parameter::PathMismatchBuilder::default();
256            let builder = builder.maybe_add(
257                Some(stream),
258                segments,
259                "write_stream",
260                "projects/*/datasets/*/tables/*/streams/*",
261            );
262            Error::binding(BindingError {
263                paths: vec![builder.build()],
264            })
265        })
266        .map(|_| ())
267}
268
269#[cfg(test)]
270mod tests {
271    use super::super::super::transport::tests::test_transport;
272    use super::*;
273    use bigquery_grpc_mock::google::cloud::bigquery::storage::v1::WriteStream as MockWriteStream;
274    use bigquery_grpc_mock::{MockBigQueryWrite, start};
275    use test_case::test_case;
276    use tokio::task::JoinHandle;
277
278    #[tokio::test]
279    async fn pending_success() -> anyhow::Result<()> {
280        let mut mock = MockBigQueryWrite::new();
281        mock.expect_create_write_stream().return_once(|req| {
282            let req = req.into_inner();
283            assert_eq!(req.parent, "projects/p/datasets/d/tables/t");
284            let ws = req.write_stream.expect("write_stream populated");
285            assert_eq!(Type::from(ws.r#type), Type::Pending);
286            Ok(gaxi::grpc::tonic::Response::new(MockWriteStream {
287                name: "projects/p/datasets/d/tables/t/streams/s".to_string(),
288                ..Default::default()
289            }))
290        });
291        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
292        let transport = Arc::new(test_transport(endpoint).await?);
293        let schema = ArrowSchema::new().set_serialized_schema("test");
294        let builder = WriterBuilder::new(transport, schema.clone());
295        let writer = builder.pending("projects/p/datasets/d/tables/t").await?;
296        assert_eq!(
297            writer.inner.write_stream,
298            "projects/p/datasets/d/tables/t/streams/s"
299        );
300        assert_eq!(writer.inner.schema, schema);
301        Ok(())
302    }
303
304    #[test_case("projects/p")]
305    #[test_case("projects/p/tables/t")]
306    #[test_case("projects/p/datasets/d/tables/")]
307    #[tokio::test]
308    async fn pending_bad_table_format(table: &str) -> anyhow::Result<()> {
309        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
310        let schema = ArrowSchema::new().set_serialized_schema("test");
311        let builder = WriterBuilder::new(transport, schema.clone());
312        let err = builder
313            .pending(table)
314            .await
315            .expect_err("should fail locally on bad format");
316        assert!(err.is_binding(), "{err:?}");
317        Ok(())
318    }
319
320    #[tokio::test]
321    async fn committed_success() -> anyhow::Result<()> {
322        let mut mock = MockBigQueryWrite::new();
323        mock.expect_create_write_stream().return_once(|req| {
324            let req = req.into_inner();
325            assert_eq!(req.parent, "projects/p/datasets/d/tables/t");
326            let ws = req.write_stream.expect("write_stream populated");
327            assert_eq!(Type::from(ws.r#type), Type::Committed);
328            Ok(gaxi::grpc::tonic::Response::new(MockWriteStream {
329                name: "projects/p/datasets/d/tables/t/streams/s".to_string(),
330                ..Default::default()
331            }))
332        });
333        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
334        let transport = Arc::new(test_transport(endpoint).await?);
335        let schema = ArrowSchema::new().set_serialized_schema("test");
336        let builder = WriterBuilder::new(transport, schema.clone());
337        let writer = builder.committed("projects/p/datasets/d/tables/t").await?;
338        assert_eq!(
339            writer.inner.write_stream,
340            "projects/p/datasets/d/tables/t/streams/s"
341        );
342        assert_eq!(writer.inner.schema, schema);
343        Ok(())
344    }
345
346    #[test_case("projects/p")]
347    #[test_case("projects/p/tables/t")]
348    #[test_case("projects/p/datasets/d/tables/")]
349    #[tokio::test]
350    async fn committed_bad_table_format(table: &str) -> anyhow::Result<()> {
351        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
352        let schema = ArrowSchema::new().set_serialized_schema("test");
353        let builder = WriterBuilder::new(transport, schema.clone());
354        let err = builder
355            .committed(table)
356            .await
357            .expect_err("should fail locally on bad format");
358        assert!(err.is_binding(), "{err:?}");
359        Ok(())
360    }
361
362    #[tokio::test]
363    async fn default() -> anyhow::Result<()> {
364        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
365        let schema = ArrowSchema::new().set_serialized_schema("test");
366        let builder = WriterBuilder::new(transport, schema.clone());
367        let writer = builder.default("projects/p/datasets/d/tables/t")?;
368        assert_eq!(
369            writer.write_stream,
370            "projects/p/datasets/d/tables/t/streams/_default"
371        );
372        assert_eq!(writer.schema, schema);
373        Ok(())
374    }
375
376    #[test_case("projects/p")]
377    #[test_case("projects/p/tables/t")]
378    #[test_case("projects/p/datasets/d/tables/")]
379    #[test_case("projects/p/instances/i/tables/t")]
380    #[test_case("projects/p/datasets/d/tables/t/streams")]
381    #[test_case("projects/p/datasets/d/tables/t/streams/_default")]
382    #[tokio::test]
383    async fn bad_table_format(table: &str) -> anyhow::Result<()> {
384        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
385        let schema = ArrowSchema::new().set_serialized_schema("test");
386        let builder = WriterBuilder::new(transport, schema.clone());
387        let err = builder
388            .default(table)
389            .expect_err("should fail locally on bad format");
390        assert!(err.is_binding(), "{err:?}");
391        Ok(())
392    }
393    #[tokio::test]
394    async fn buffered_success() -> anyhow::Result<()> {
395        let mut mock = MockBigQueryWrite::new();
396        mock.expect_create_write_stream().return_once(|req| {
397            let req = req.into_inner();
398            assert_eq!(req.parent, "projects/p/datasets/d/tables/t");
399            let ws = req.write_stream.expect("write_stream populated");
400            assert_eq!(Type::from(ws.r#type), Type::Buffered);
401            Ok(gaxi::grpc::tonic::Response::new(MockWriteStream {
402                name: "projects/p/datasets/d/tables/t/streams/s".to_string(),
403                ..Default::default()
404            }))
405        });
406        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
407        let transport = Arc::new(test_transport(endpoint).await?);
408        let schema = ArrowSchema::new().set_serialized_schema("test");
409        let builder = WriterBuilder::new(transport, schema.clone());
410        let writer = builder.buffered("projects/p/datasets/d/tables/t").await?;
411        assert_eq!(
412            writer.inner.write_stream,
413            "projects/p/datasets/d/tables/t/streams/s"
414        );
415        assert_eq!(writer.inner.schema, schema);
416        Ok(())
417    }
418
419    #[test_case("projects/p")]
420    #[test_case("projects/p/tables/t")]
421    #[test_case("projects/p/datasets/d/tables/")]
422    #[tokio::test]
423    async fn buffered_bad_table_format(table: &str) -> anyhow::Result<()> {
424        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
425        let schema = ArrowSchema::new().set_serialized_schema("test");
426        let builder = WriterBuilder::new(transport, schema.clone());
427        let err = builder
428            .buffered(table)
429            .await
430            .expect_err("should fail locally on bad format");
431        assert!(err.is_binding(), "{err:?}");
432        Ok(())
433    }
434
435    async fn attach_mock(stream_type: Type) -> anyhow::Result<(Arc<Transport>, JoinHandle<()>)> {
436        let mut mock = MockBigQueryWrite::new();
437        mock.expect_get_write_stream().return_once(move |req| {
438            let req = req.into_inner();
439            assert_eq!(req.name, "projects/p/datasets/d/tables/t/streams/s");
440            Ok(gaxi::grpc::tonic::Response::new(MockWriteStream {
441                name: "projects/p/datasets/d/tables/t/streams/s".to_string(),
442                r#type: stream_type.value().expect("known enum value"),
443                ..Default::default()
444            }))
445        });
446        let (endpoint, server) = start("0.0.0.0:0", mock).await?;
447        let transport = Arc::new(test_transport(endpoint).await?);
448        Ok((transport, server))
449    }
450
451    #[tokio::test]
452    async fn attach_committed_success() -> anyhow::Result<()> {
453        let (transport, _server) = attach_mock(Type::Committed).await?;
454        let schema = ArrowSchema::new().set_serialized_schema("test");
455        let builder = WriterBuilder::new(transport, schema.clone());
456        let writer: CommittedWriter = builder
457            .attach("projects/p/datasets/d/tables/t/streams/s")
458            .await?;
459        assert_eq!(
460            writer.inner.write_stream,
461            "projects/p/datasets/d/tables/t/streams/s"
462        );
463        assert_eq!(writer.inner.schema, schema);
464        Ok(())
465    }
466
467    #[tokio::test]
468    async fn attach_pending_success() -> anyhow::Result<()> {
469        let (transport, _server) = attach_mock(Type::Pending).await?;
470        let schema = ArrowSchema::new().set_serialized_schema("test");
471        let builder = WriterBuilder::new(transport, schema.clone());
472        let writer: PendingWriter = builder
473            .attach("projects/p/datasets/d/tables/t/streams/s")
474            .await?;
475        assert_eq!(
476            writer.inner.write_stream,
477            "projects/p/datasets/d/tables/t/streams/s"
478        );
479        assert_eq!(writer.inner.schema, schema);
480        Ok(())
481    }
482
483    #[tokio::test]
484    async fn attach_buffered_success() -> anyhow::Result<()> {
485        let (transport, _server) = attach_mock(Type::Buffered).await?;
486        let schema = ArrowSchema::new().set_serialized_schema("test");
487        let builder = WriterBuilder::new(transport, schema.clone());
488        let writer: BufferedWriter = builder
489            .attach("projects/p/datasets/d/tables/t/streams/s")
490            .await?;
491        assert_eq!(
492            writer.inner.write_stream,
493            "projects/p/datasets/d/tables/t/streams/s"
494        );
495        assert_eq!(writer.inner.schema, schema);
496        Ok(())
497    }
498
499    #[test_case("projects/p")]
500    #[test_case("projects/p/tables/t")]
501    #[test_case("projects/p/datasets/d/tables/t")]
502    #[test_case("projects/p/datasets/d/tables/t/streams/")]
503    #[tokio::test]
504    async fn attach_bad_stream_format(stream: &str) -> anyhow::Result<()> {
505        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
506        let schema = ArrowSchema::new().set_serialized_schema("test");
507        let builder = WriterBuilder::new(transport, schema.clone());
508        let err = builder
509            .attach::<CommittedWriter, _>(stream)
510            .await
511            .expect_err("should fail locally on bad format");
512        assert!(matches!(err, AttachError::Rpc { source: e } if e.is_binding()));
513        Ok(())
514    }
515
516    #[tokio::test]
517    async fn attach_stream_type_mismatch() -> anyhow::Result<()> {
518        let (transport, _server) = attach_mock(Type::Buffered).await?;
519        let schema = ArrowSchema::new().set_serialized_schema("test");
520        let builder = WriterBuilder::new(transport, schema.clone());
521        let err = builder
522            .attach::<CommittedWriter, _>("projects/p/datasets/d/tables/t/streams/s")
523            .await
524            .expect_err("should return type mismatch error");
525        assert!(matches!(err, AttachError::TypeMismatch { .. }));
526        assert!(err.to_string().contains("stream type mismatch: requested"));
527        Ok(())
528    }
529}
google_cloud_bigquery/write/arrow/writer_builder.rs

google_cloud_bigquery/write/arrow/
writer_builder.rs
