google_cloud_bigquery/write/arrow/
default.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::append_builder::Append;
16use super::super::runner::Runner;
17use super::super::transport::Transport;
18use crate::model::append_rows_request::ArrowData;
19use crate::model::{AppendRowsRequest, ArrowRecordBatch, ArrowSchema};
20use std::sync::Arc;
21
22/// A writer for the [default stream]
23///
24/// [default stream]: https://docs.cloud.google.com/bigquery/docs/write-api#default_stream
25#[derive(Debug)]
26pub struct DefaultWriter {
27    // TODO(#5744) - support multiplexed connections
28    runner: Runner,
29    pub(crate) write_stream: String,
30    pub(crate) schema: ArrowSchema,
31}
32
33impl DefaultWriter {
34    pub(crate) fn new(inner: Arc<Transport>, write_stream: String, schema: ArrowSchema) -> Self {
35        let runner = Runner::new(inner);
36        Self {
37            runner,
38            write_stream,
39            schema,
40        }
41    }
42
43    /// Append rows to the stream.
44    pub fn append(&self, rows: ArrowRecordBatch) -> Append {
45        // TODO(#5744) - send optimization
46        let req = AppendRowsRequest::new()
47            .set_write_stream(&self.write_stream)
48            .set_arrow_rows(
49                ArrowData::new()
50                    .set_writer_schema(self.schema.clone())
51                    .set_rows(rows),
52            );
53        Append::new(self.runner.req_tx.clone(), req)
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::super::super::runner::tests::*;
60    use super::super::super::transport::tests::*;
61    use super::*;
62    use crate::error::AppendError;
63    use bigquery_grpc_mock::{MockBigQueryWrite, start};
64    use gaxi::grpc::tonic::Response as TonicResponse;
65    use tokio::sync::mpsc;
66
67    #[tokio::test]
68    async fn request_fields() -> anyhow::Result<()> {
69        let transport = Arc::new(test_transport("http://ignored:1".to_string()).await?);
70        let writer = DefaultWriter::new(transport, write_stream(), schema());
71
72        let b = writer.append(rows(1));
73        assert_eq!(b.req.write_stream, write_stream());
74        let data = b.req.arrow_rows().expect("arrow rows should be set");
75        let s = data.writer_schema.as_ref().expect("schema should be set");
76        assert_eq!(s.serialized_schema, "test");
77        let r = data.rows.as_ref().expect("rows should be set");
78        assert_eq!(r.serialized_record_batch, "1");
79
80        let b = writer.append(rows(2));
81        assert_eq!(b.req.write_stream, write_stream());
82        let data = b.req.arrow_rows().expect("arrow rows should be set");
83        let s = data.writer_schema.as_ref().expect("schema should be set");
84        assert_eq!(s.serialized_schema, "test");
85        let r = data.rows.as_ref().expect("rows should be set");
86        assert_eq!(r.serialized_record_batch, "2");
87
88        Ok(())
89    }
90
91    #[tokio::test]
92    async fn basic_success() -> anyhow::Result<()> {
93        let (response_tx, response_rx) = mpsc::channel(10);
94
95        let mut mock = MockBigQueryWrite::new();
96        mock.expect_append_rows()
97            .return_once(|_| Ok(TonicResponse::from(response_rx)));
98        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
99        let transport = Arc::new(test_transport(endpoint).await?);
100
101        let writer = DefaultWriter::new(transport, write_stream(), schema());
102
103        response_tx.send(Ok(convert(&test_response(1)))).await?;
104        let resp = writer.append(rows(1)).send().await?;
105        assert_eq!(resp.offset, Some(1));
106
107        response_tx.send(Ok(convert(&test_response(2)))).await?;
108        let resp = writer.append(rows(2)).send().await?;
109        assert_eq!(resp.offset, Some(2));
110
111        response_tx.send(Ok(convert(&test_response(3)))).await?;
112        let resp = writer.append(rows(3)).send().await?;
113        assert_eq!(resp.offset, Some(3));
114
115        drop(response_tx);
116        let err = writer.append(rows(4)).send().await.expect_err("channel");
117        assert!(matches!(err, AppendError::UnexpectedEndOfStream));
118
119        Ok(())
120    }
121
122    fn write_stream() -> String {
123        "projects/p/datasets/d/tables/t/streams/_default".to_string()
124    }
125
126    fn schema() -> ArrowSchema {
127        ArrowSchema::new().set_serialized_schema("test")
128    }
129
130    fn rows(id: i64) -> ArrowRecordBatch {
131        ArrowRecordBatch::new().set_serialized_record_batch(id.to_string())
132    }
133}
google_cloud_bigquery/write/arrow/default.rs

google_cloud_bigquery/write/arrow/
default.rs
