google_cloud_bigquery/query/
iterator.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 crate::RequestOptions;
16use crate::error::RowError;
17#[cfg(any(
18    google_cloud_unstable_bigquery_arrow,
19    google_cloud_unstable_bigquery_storage_read
20))]
21use crate::query::arrow::{ArrowReader, ArrowResponseReader, RecordBatchSource};
22use crate::query::query_handle::CachedData;
23use crate::query::{CompleteQuery, Row, Schema};
24#[cfg(any(
25    google_cloud_unstable_bigquery_arrow,
26    google_cloud_unstable_bigquery_storage_read
27))]
28use arrow::record_batch::RecordBatch;
29use google_cloud_bigquery_v2::client::JobService;
30use google_cloud_bigquery_v2::model::{GetQueryResultsRequest, JobReference};
31use std::collections::VecDeque;
32use std::sync::Arc;
33
34pub type Result<T> = std::result::Result<T, RowError>;
35/// An iterator over rows returned by a query.
36///
37/// [`CompleteQuery::read()`](crate::query::CompleteQuery::read) returns a
38/// `RowIterator`.
39///
40/// `RowIterator` yields rows from the query result set, automatically requesting
41/// subsequent pages from BigQuery as needed until the entire result set has been consumed.
42///
43/// # Example
44///
45/// ```
46/// # use google_cloud_bigquery::client::BigQuery;
47/// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
48/// let mut rows = client
49///     .query("SELECT name, state FROM `bigquery-public-data.usa_names.usa_1910_2013` LIMIT 10")
50///     .until_done()
51///     .await?
52///     .read();
53///
54/// while let Some(row) = rows.next().await.transpose()? {
55///     let name: String = row.get("name")?;
56///     let state: String = row.get("state")?;
57///     println!("{name} from {state}");
58/// }
59/// # Ok(())
60/// # }
61/// ```
62#[derive(Debug)]
63pub struct RowIterator {
64    job_service: Arc<JobService>,
65    job_ref: Option<JobReference>,
66    schema: Arc<Schema>,
67    page_token: Option<String>,
68    #[cfg(any(
69        google_cloud_unstable_bigquery_arrow,
70        google_cloud_unstable_bigquery_storage_read
71    ))]
72    current_batch: Option<Arc<RecordBatch>>,
73    #[cfg(any(
74        google_cloud_unstable_bigquery_arrow,
75        google_cloud_unstable_bigquery_storage_read
76    ))]
77    row_index: usize,
78    rows: VecDeque<wkt::Struct>,
79    page_size: Option<u32>,
80    options: RequestOptions,
81    #[cfg(any(
82        google_cloud_unstable_bigquery_arrow,
83        google_cloud_unstable_bigquery_storage_read
84    ))]
85    arrow_reader: Option<ArrowReader>,
86}
87
88impl RowIterator {
89    pub(crate) fn new(q: CompleteQuery) -> Self {
90        // DDL/DML or Arrow queries may have no TableSchema in metadata initially.
91        let schema = Arc::new(Schema::new(q.metadata.schema.unwrap_or_default()));
92        #[cfg(any(
93            google_cloud_unstable_bigquery_arrow,
94            google_cloud_unstable_bigquery_storage_read
95        ))]
96        let (rows, arrow_reader) = match q.cached_data {
97            CachedData::Rows(rows) => (rows, None),
98            CachedData::Arrow {
99                serialized_record_batch,
100                serialized_schema,
101            } => (
102                VecDeque::new(),
103                Some(ArrowReader::Response(ArrowResponseReader::new(
104                    serialized_schema,
105                    serialized_record_batch,
106                ))),
107            ),
108        };
109        #[cfg(not(any(
110            google_cloud_unstable_bigquery_arrow,
111            google_cloud_unstable_bigquery_storage_read
112        )))]
113        let rows = match q.cached_data {
114            CachedData::Rows(rows) => rows,
115            // TODO(#7035): remove once google_cloud_unstable_bigquery_arrow is stabilized
116            CachedData::Arrow { .. } => {
117                unimplemented!("Arrow IPC decoding is not yet implemented (#7035)")
118            }
119        };
120        Self {
121            job_service: q.job_service,
122            job_ref: q.job_ref,
123            schema,
124            page_token: q.page_token,
125            #[cfg(any(
126                google_cloud_unstable_bigquery_arrow,
127                google_cloud_unstable_bigquery_storage_read
128            ))]
129            current_batch: None,
130            #[cfg(any(
131                google_cloud_unstable_bigquery_arrow,
132                google_cloud_unstable_bigquery_storage_read
133            ))]
134            row_index: 0,
135            rows,
136            page_size: q.page_size,
137            options: q.options,
138            #[cfg(any(
139                google_cloud_unstable_bigquery_arrow,
140                google_cloud_unstable_bigquery_storage_read
141            ))]
142            arrow_reader,
143        }
144    }
145
146    /// Sets the maximum number of rows to buffer in memory.
147    ///
148    /// This controls page size during network fetches when streaming result
149    /// sets.
150    ///
151    /// # Example
152    ///
153    /// ```
154    /// # use google_cloud_bigquery::client::BigQuery;
155    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
156    /// let mut rows = client
157    ///     .query("SELECT 1 AS n")
158    ///     .until_done()
159    ///     .await?
160    ///     .read()
161    ///     .set_page_size(500);
162    /// # Ok(())
163    /// # }
164    /// ```
165    pub fn set_page_size(mut self, page_size: u32) -> Self {
166        self.page_size = Some(page_size);
167        self
168    }
169
170    /// Fetches the next [`Row`](crate::query::Row) from the result set.
171    ///
172    /// Returns `None` when all rows have been retrieved.
173    ///
174    /// # Example
175    ///
176    /// ```
177    /// # use google_cloud_bigquery::query::RowIterator;
178    /// # async fn sample(mut rows: RowIterator) -> anyhow::Result<()> {
179    /// while let Some(row) = rows.next().await.transpose()? {
180    ///     let msg: String = row.get("msg")?;
181    ///     println!("Message: {msg}");
182    /// }
183    /// # Ok(())
184    /// # }
185    /// ```
186    pub async fn next(&mut self) -> Option<Result<Row>> {
187        loop {
188            #[cfg(any(
189                google_cloud_unstable_bigquery_arrow,
190                google_cloud_unstable_bigquery_storage_read
191            ))]
192            {
193                if let Some(batch) = &self.current_batch {
194                    if self.row_index < batch.num_rows() {
195                        let idx = self.row_index;
196                        self.row_index += 1;
197                        return Some(Row::try_new_from_arrow(batch, idx, &self.schema));
198                    }
199                    self.current_batch = None;
200                    self.row_index = 0;
201                }
202
203                if let Some(ref mut arrow_reader) = self.arrow_reader {
204                    match arrow_reader.next_record_batch().await {
205                        Ok(Some(batch)) => {
206                            if self.schema.len() == 0 {
207                                self.schema = Arc::new(Schema::from_arrow_schema(&batch.schema()));
208                            }
209                            self.current_batch = Some(Arc::new(batch));
210                            self.row_index = 0;
211                            continue;
212                        }
213                        Ok(None) => {
214                            self.arrow_reader = None;
215                        }
216                        Err(_err) if arrow_reader.can_fallback_to_rest() => {
217                            // Fall back to JSON pagination if read stream opening failed
218                            self.arrow_reader = None;
219                        }
220                        Err(err) => {
221                            self.arrow_reader = None;
222                            return Some(Err(err));
223                        }
224                    }
225                }
226            }
227
228            if let Some(raw_row) = self.rows.pop_front() {
229                return Some(Row::try_new(raw_row, &self.schema));
230            }
231
232            if let Err(e) = self.try_fetch_page().await {
233                return Some(Err(e));
234            }
235
236            if self.rows.is_empty() && self.page_token.is_none() {
237                return None;
238            }
239        }
240    }
241
242    async fn try_fetch_page(&mut self) -> Result<()> {
243        let Some(token) = self.page_token.as_deref() else {
244            return Ok(());
245        };
246
247        let (fetched_rows, next_token) = self.fetch_page(token).await?;
248        if fetched_rows.is_empty() && next_token.as_deref() == Some(token) {
249            self.page_token = None;
250        } else {
251            self.page_token = next_token;
252            self.rows.extend(fetched_rows);
253        }
254        Ok(())
255    }
256
257    // Fetches the next page of results and the next page token.
258    async fn fetch_page(&self, token: &str) -> Result<(Vec<wkt::Struct>, Option<String>)> {
259        let job_ref = self.job_ref.as_ref().expect(
260            "only queries with a job reference should have page tokens and can fetch more pages",
261        );
262
263        let mut req = GetQueryResultsRequest::new()
264            .set_project_id(job_ref.project_id.clone())
265            .set_or_clear_max_results(self.page_size)
266            .set_job_id(job_ref.job_id.clone())
267            .set_page_token(token)
268            .set_format_options(
269                google_cloud_bigquery_v2::model::DataFormatOptions::new()
270                    .set_use_int64_timestamp(true),
271            );
272        if let Some(location) = job_ref.location.clone() {
273            req = req.set_location(location);
274        }
275
276        let res = self
277            .job_service
278            .get_query_results()
279            .with_options(self.options.clone())
280            .with_request(req)
281            .send()
282            .await?;
283
284        let page_token = if res.page_token.is_empty() {
285            None
286        } else {
287            Some(res.page_token)
288        };
289
290        Ok((res.rows, page_token))
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use crate::query::tests::{MockJobService, create_job_service};
298    use google_cloud_bigquery_v2::model::{
299        DataFormatOptions, GetQueryResultsResponse, JobReference, QueryResponse, TableFieldSchema,
300        TableSchema,
301    };
302    use google_cloud_gax::error::Error as GaxError;
303    use google_cloud_gax::error::rpc::{Code, Status};
304    use google_cloud_gax::response::Response;
305    use serde_json::{Map, json};
306    use std::sync::Arc;
307
308    type TestResult = anyhow::Result<()>;
309
310    fn create_test_schema() -> TableSchema {
311        TableSchema::new().set_fields([TableFieldSchema::new()
312            .set_name("col")
313            .set_type("STRING")
314            .set_mode("NULLABLE")])
315    }
316
317    fn create_test_row(val: &str) -> wkt::Struct {
318        Map::from_iter([("f".to_string(), json!([{ "v": val }]))])
319    }
320
321    fn create_test_job_ref() -> JobReference {
322        JobReference::new()
323            .set_project_id("test_project")
324            .set_job_id("test_job")
325    }
326
327    fn create_test_job_ref_with_location(location: &str) -> JobReference {
328        create_test_job_ref().set_location(location)
329    }
330
331    fn create_test_complete_query(
332        job_service: Arc<JobService>,
333        job_ref: Option<JobReference>,
334        rows: Vec<wkt::Struct>,
335        page_token: Option<String>,
336    ) -> CompleteQuery {
337        let mut res = QueryResponse::new()
338            .set_schema(create_test_schema())
339            .set_rows(rows);
340        if let Some(token) = page_token {
341            res = res.set_page_token(token);
342        }
343        if let Some(job_ref) = job_ref {
344            res = res.set_job_reference(job_ref);
345        }
346        CompleteQuery::from_query_response(job_service, res, None)
347    }
348
349    #[tokio::test]
350    async fn test_row_iterator_empty_no_token() -> TestResult {
351        let job_service = create_job_service(MockJobService::new());
352        let q = create_test_complete_query(job_service, Some(create_test_job_ref()), vec![], None);
353        let mut iter = q.read();
354        assert!(iter.next().await.is_none(), "{iter:?}");
355        Ok(())
356    }
357
358    #[tokio::test]
359    async fn test_row_iterator_cached_rows_only() -> TestResult {
360        let job_service = create_job_service(MockJobService::new());
361        let rows = vec![create_test_row("first"), create_test_row("second")];
362        let q = create_test_complete_query(job_service, Some(create_test_job_ref()), rows, None);
363        let mut iter = q.read();
364
365        let row1 = iter.next().await.expect("should have row 1")?;
366        assert_eq!(row1.get::<String>("col")?, "first");
367
368        let row2 = iter.next().await.expect("should have row 2")?;
369        assert_eq!(row2.get::<String>("col")?, "second");
370
371        assert!(iter.next().await.is_none(), "{iter:?}");
372        Ok(())
373    }
374
375    #[tokio::test]
376    async fn test_row_iterator_row_conversion_error() -> TestResult {
377        let job_service = create_job_service(MockJobService::new());
378        let invalid_row = Map::from_iter([("f".to_string(), json!([]))]);
379        let q = create_test_complete_query(
380            job_service,
381            Some(create_test_job_ref()),
382            vec![invalid_row],
383            None,
384        );
385        let mut iter = q.read();
386
387        let err = iter.next().await.expect("should return error").unwrap_err();
388        assert!(matches!(err, RowError::InvalidRowFormat(_)), "{err:?}");
389        assert!(iter.next().await.is_none(), "{iter:?}");
390        Ok(())
391    }
392
393    #[tokio::test]
394    async fn test_row_iterator_fetch_multiple_pages() -> TestResult {
395        let mut mock = MockJobService::new();
396        let mut seq = mockall::Sequence::new();
397
398        mock.expect_get_query_results()
399            .in_sequence(&mut seq)
400            .times(1)
401            .returning(|req, _| {
402                assert_eq!(req.page_token, "token_page_1");
403                let res = GetQueryResultsResponse::new()
404                    .set_rows(vec![
405                        create_test_row("page_1_row_1"),
406                        create_test_row("page_1_row_2"),
407                    ])
408                    .set_page_token("token_page_2");
409                Ok(Response::from(res))
410            });
411
412        mock.expect_get_query_results()
413            .in_sequence(&mut seq)
414            .times(1)
415            .returning(|req, _| {
416                assert_eq!(req.page_token, "token_page_2");
417                let res = GetQueryResultsResponse::new()
418                    .set_rows(vec![create_test_row("page_2_row_1")])
419                    .set_page_token("");
420                Ok(Response::from(res))
421            });
422
423        let job_service = create_job_service(mock);
424        let q = create_test_complete_query(
425            job_service,
426            Some(create_test_job_ref()),
427            vec![create_test_row("cached_row")],
428            Some("token_page_1".to_string()),
429        );
430        let mut iter = q.read();
431
432        let row1 = iter.next().await.expect("should have row 1")?;
433        assert_eq!(row1.get::<String>("col")?, "cached_row");
434
435        let row2 = iter.next().await.expect("should have row 2")?;
436        assert_eq!(row2.get::<String>("col")?, "page_1_row_1");
437
438        let row3 = iter.next().await.expect("should have row 3")?;
439        assert_eq!(row3.get::<String>("col")?, "page_1_row_2");
440
441        let row4 = iter.next().await.expect("should have row 4")?;
442        assert_eq!(row4.get::<String>("col")?, "page_2_row_1");
443
444        assert!(iter.next().await.is_none(), "{iter:?}");
445        Ok(())
446    }
447
448    #[tokio::test]
449    async fn test_row_iterator_fetch_page_request_parameters() -> TestResult {
450        let mut mock = MockJobService::new();
451        mock.expect_get_query_results()
452            .times(1)
453            .returning(|req, _| {
454                assert_eq!(req.project_id, "test_project");
455                assert_eq!(req.job_id, "test_job");
456                assert_eq!(req.page_token, "token_1");
457                assert_eq!(req.max_results, Some(50));
458                assert_eq!(req.location, "us-east1");
459                assert_eq!(
460                    req.format_options,
461                    Some(DataFormatOptions::new().set_use_int64_timestamp(true))
462                );
463                let res = GetQueryResultsResponse::new()
464                    .set_rows(vec![create_test_row("page_row")])
465                    .set_page_token("");
466                Ok(Response::from(res))
467            });
468
469        let job_service = create_job_service(mock);
470        let q = create_test_complete_query(
471            job_service,
472            Some(create_test_job_ref_with_location("us-east1")),
473            vec![],
474            Some("token_1".to_string()),
475        );
476        let mut iter = q.read().set_page_size(50);
477
478        let row = iter.next().await.expect("should have row")?;
479        assert_eq!(row.get::<String>("col")?, "page_row");
480
481        assert!(iter.next().await.is_none(), "{iter:?}");
482        Ok(())
483    }
484
485    #[tokio::test]
486    async fn test_row_iterator_inherits_page_size() -> TestResult {
487        let mut mock = MockJobService::new();
488        mock.expect_get_query_results()
489            .times(1)
490            .returning(|req, _| {
491                assert_eq!(req.max_results, Some(25));
492                let res = GetQueryResultsResponse::new()
493                    .set_rows(vec![create_test_row("page_row")])
494                    .set_page_token("");
495                Ok(Response::from(res))
496            });
497
498        let job_service = create_job_service(mock);
499        let res = QueryResponse::new()
500            .set_schema(create_test_schema())
501            .set_job_reference(create_test_job_ref())
502            .set_page_token("token_1");
503        let q = CompleteQuery::from_query_response(job_service, res, Some(25));
504        let mut iter = q.read();
505
506        let row = iter.next().await.expect("should have row")?;
507        assert_eq!(row.get::<String>("col")?, "page_row");
508        assert!(iter.next().await.is_none(), "{iter:?}");
509        Ok(())
510    }
511
512    #[tokio::test]
513    async fn test_row_iterator_fetch_page_rpc_error() -> TestResult {
514        let mut mock = MockJobService::new();
515        mock.expect_get_query_results().returning(|_, _| {
516            let status = Status::default()
517                .set_code(Code::Unavailable)
518                .set_message("temporary service error");
519            Err(GaxError::service(status))
520        });
521
522        let job_service = create_job_service(mock);
523        let q = create_test_complete_query(
524            job_service,
525            Some(create_test_job_ref()),
526            vec![],
527            Some("token_err".to_string()),
528        );
529        let mut iter = q.read();
530
531        let err = iter.next().await.expect("should return error").unwrap_err();
532        assert!(matches!(err, RowError::Rpc { .. }), "{err:?}");
533        assert!(
534            err.to_string().contains("temporary service error"),
535            "{err:?}"
536        );
537        Ok(())
538    }
539
540    #[tokio::test]
541    async fn test_row_iterator_empty_page_handling() -> TestResult {
542        let mut mock = MockJobService::new();
543        let mut seq = mockall::Sequence::new();
544
545        // empty page with a new token: should continue fetching
546        mock.expect_get_query_results()
547            .times(1)
548            .in_sequence(&mut seq)
549            .returning(|req, _| {
550                assert_eq!(req.page_token, "token_1");
551                let res = GetQueryResultsResponse::new()
552                    .set_rows(Vec::<wkt::Struct>::new())
553                    .set_page_token("token_2");
554                Ok(Response::from(res))
555            });
556
557        // next page returns rows and advances to token_3
558        mock.expect_get_query_results()
559            .times(1)
560            .in_sequence(&mut seq)
561            .returning(|req, _| {
562                assert_eq!(req.page_token, "token_2");
563                let res = GetQueryResultsResponse::new()
564                    .set_rows(vec![create_test_row("page2_row")])
565                    .set_page_token("token_3");
566                Ok(Response::from(res))
567            });
568
569        // empty page with the same token: should terminate pagination
570        mock.expect_get_query_results()
571            .times(1)
572            .in_sequence(&mut seq)
573            .returning(|req, _| {
574                assert_eq!(req.page_token, "token_3");
575                let res = GetQueryResultsResponse::new()
576                    .set_rows(Vec::<wkt::Struct>::new())
577                    .set_page_token("token_3");
578                Ok(Response::from(res))
579            });
580
581        let job_service = create_job_service(mock);
582        let q = create_test_complete_query(
583            job_service,
584            Some(create_test_job_ref()),
585            vec![],
586            Some("token_1".to_string()),
587        );
588        let mut iter = q.read();
589
590        let row = iter.next().await.expect("should have row")?;
591        assert_eq!(row.get::<String>("col")?, "page2_row");
592        assert!(iter.next().await.is_none(), "{iter:?}");
593        Ok(())
594    }
595
596    #[tokio::test]
597    async fn test_row_iterator_inherits_user_agent() -> TestResult {
598        let mut mock = MockJobService::new();
599        mock.expect_get_query_results()
600            .times(1)
601            .returning(|_, options| {
602                assert_eq!(options.user_agent().as_deref(), Some("my-app/1.2.3"));
603                let res = GetQueryResultsResponse::new()
604                    .set_rows(vec![create_test_row("page_row")])
605                    .set_page_token("");
606                Ok(Response::from(res))
607            });
608
609        let job_service = create_job_service(mock);
610        let res = QueryResponse::new()
611            .set_schema(create_test_schema())
612            .set_job_reference(create_test_job_ref())
613            .set_page_token("token_1");
614        let mut q = CompleteQuery::from_query_response(job_service, res, None);
615        q.options.set_user_agent("my-app/1.2.3");
616        let mut iter = q.read();
617
618        let row = iter.next().await.expect("should have row")?;
619        assert_eq!(row.get::<String>("col")?, "page_row");
620        assert!(iter.next().await.is_none(), "{iter:?}");
621        Ok(())
622    }
623
624    #[tokio::test]
625    async fn test_row_iterator_inherits_quota_project() -> TestResult {
626        let mut mock = MockJobService::new();
627        mock.expect_get_query_results()
628            .times(1)
629            .returning(|_, options| {
630                assert_eq!(options.quota_project().as_deref(), Some("test_project"));
631                let res = GetQueryResultsResponse::new()
632                    .set_rows(vec![create_test_row("page_row")])
633                    .set_page_token("");
634                Ok(Response::from(res))
635            });
636
637        let job_service = create_job_service(mock);
638        let res = QueryResponse::new()
639            .set_schema(create_test_schema())
640            .set_job_reference(create_test_job_ref())
641            .set_page_token("token_1");
642        let mut q = CompleteQuery::from_query_response(job_service, res, None);
643        q.options.set_quota_project("test_project");
644        let mut iter = q.read();
645
646        let row = iter.next().await.expect("should have row")?;
647        assert_eq!(row.get::<String>("col")?, "page_row");
648        assert!(iter.next().await.is_none(), "{iter:?}");
649        Ok(())
650    }
651
652    #[cfg(any(
653        google_cloud_unstable_bigquery_arrow,
654        google_cloud_unstable_bigquery_storage_read
655    ))]
656    #[tokio::test]
657    async fn test_row_iterator_cached_arrow_multiple_batches() -> TestResult {
658        use arrow::array::{Int64Array, StringArray};
659        use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
660        use arrow::ipc::writer::StreamWriter;
661
662        let arrow_schema = Arc::new(ArrowSchema::new(vec![
663            Field::new("col", DataType::Utf8, false),
664            Field::new("num", DataType::Int64, false),
665        ]));
666
667        let mut schema_buf = Vec::new();
668        let _ = StreamWriter::try_new(&mut schema_buf, &arrow_schema)?;
669
670        let batch1 = RecordBatch::try_new(
671            arrow_schema.clone(),
672            vec![
673                Arc::new(StringArray::from(vec!["hello", "world"])),
674                Arc::new(Int64Array::from(vec![42, 100])),
675            ],
676        )?;
677        let empty_batch = RecordBatch::new_empty(arrow_schema.clone());
678        let batch2 = RecordBatch::try_new(
679            arrow_schema.clone(),
680            vec![
681                Arc::new(StringArray::from(vec!["last"])),
682                Arc::new(Int64Array::from(vec![999])),
683            ],
684        )?;
685
686        let mut stream_buf = Vec::new();
687        let mut writer = StreamWriter::try_new(&mut stream_buf, &arrow_schema)?;
688        writer.write(&batch1)?;
689        writer.write(&empty_batch)?;
690        writer.write(&batch2)?;
691        writer.finish()?;
692        let batch_buf = stream_buf[schema_buf.len()..].to_vec();
693
694        let job_service = create_job_service(MockJobService::new());
695        let q = CompleteQuery {
696            job_service,
697            job_ref: None,
698            cached_data: CachedData::Arrow {
699                serialized_schema: schema_buf.into(),
700                serialized_record_batch: batch_buf.into(),
701            },
702            page_token: None,
703            metadata: Default::default(),
704            page_size: None,
705            options: Default::default(),
706        };
707
708        let mut iter = q.read();
709
710        let row1 = iter.next().await.expect("row 1")?;
711        assert_eq!(row1.get::<String>("col")?, "hello");
712        assert_eq!(row1.get::<i64>("num")?, 42);
713
714        let row2 = iter.next().await.expect("row 2")?;
715        assert_eq!(row2.get::<String>("col")?, "world");
716        assert_eq!(row2.get::<i64>("num")?, 100);
717
718        let row3 = iter.next().await.expect("row 3")?;
719        assert_eq!(row3.get::<String>("col")?, "last");
720        assert_eq!(row3.get::<i64>("num")?, 999);
721
722        assert!(iter.next().await.is_none(), "{iter:?}");
723        Ok(())
724    }
725
726    #[cfg(any(
727        google_cloud_unstable_bigquery_arrow,
728        google_cloud_unstable_bigquery_storage_read
729    ))]
730    #[tokio::test]
731    async fn test_row_iterator_cached_arrow_then_fetch_page() -> TestResult {
732        use arrow::array::StringArray;
733        use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
734        use arrow::ipc::writer::StreamWriter;
735
736        let mut mock = MockJobService::new();
737        mock.expect_get_query_results()
738            .times(1)
739            .returning(|req, _| {
740                assert_eq!(req.page_token, "token_page_1");
741                let res = GetQueryResultsResponse::new()
742                    .set_rows(vec![create_test_row("from_json_page")])
743                    .set_page_token("");
744                Ok(Response::from(res))
745            });
746
747        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
748            "col",
749            DataType::Utf8,
750            false,
751        )]));
752        let mut schema_buf = Vec::new();
753        let _ = StreamWriter::try_new(&mut schema_buf, &arrow_schema)?;
754
755        let batch = RecordBatch::try_new(
756            arrow_schema.clone(),
757            vec![Arc::new(StringArray::from(vec!["from_arrow_batch"]))],
758        )?;
759        let mut stream_buf = Vec::new();
760        let mut writer = StreamWriter::try_new(&mut stream_buf, &arrow_schema)?;
761        writer.write(&batch)?;
762        writer.finish()?;
763        let batch_buf = stream_buf[schema_buf.len()..].to_vec();
764
765        let job_service = create_job_service(mock);
766        let q = CompleteQuery {
767            job_service,
768            job_ref: Some(create_test_job_ref()),
769            cached_data: CachedData::Arrow {
770                serialized_schema: schema_buf.into(),
771                serialized_record_batch: batch_buf.into(),
772            },
773            page_token: Some("token_page_1".to_string()),
774            metadata: Default::default(),
775            page_size: None,
776            options: Default::default(),
777        };
778
779        let mut iter = q.read();
780
781        let row1 = iter.next().await.expect("row 1")?;
782        assert_eq!(row1.get::<String>("col")?, "from_arrow_batch");
783
784        let row2 = iter.next().await.expect("row 2")?;
785        assert_eq!(row2.get::<String>("col")?, "from_json_page");
786
787        assert!(iter.next().await.is_none(), "{iter:?}");
788        Ok(())
789    }
790
791    #[cfg(any(
792        google_cloud_unstable_bigquery_arrow,
793        google_cloud_unstable_bigquery_storage_read
794    ))]
795    fn create_test_arrow_schema_bytes() -> bytes::Bytes {
796        use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
797        use arrow::ipc::writer::StreamWriter;
798
799        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
800            "col",
801            DataType::Utf8,
802            false,
803        )]));
804        let mut schema_buf = Vec::new();
805        let _ =
806            StreamWriter::try_new(&mut schema_buf, &arrow_schema).expect("valid test arrow schema");
807        schema_buf.into()
808    }
809
810    #[cfg(any(
811        google_cloud_unstable_bigquery_arrow,
812        google_cloud_unstable_bigquery_storage_read
813    ))]
814    #[test_case::test_case(
815        bytes::Bytes::from_static(b"invalid arrow schema bytes"),
816        bytes::Bytes::new();
817        "malformed schema"
818    )]
819    #[test_case::test_case(
820        create_test_arrow_schema_bytes(),
821        bytes::Bytes::from_static(b"\xff\xff\xff\xff\x10\x00\x00\x00corrupted_batch_header");
822        "malformed record batch"
823    )]
824    #[tokio::test]
825    async fn test_row_iterator_cached_arrow_malformed_ipc(
826        serialized_schema: bytes::Bytes,
827        serialized_record_batch: bytes::Bytes,
828    ) -> TestResult {
829        let job_service = create_job_service(MockJobService::new());
830        let q = CompleteQuery {
831            job_service,
832            job_ref: None,
833            cached_data: CachedData::Arrow {
834                serialized_schema,
835                serialized_record_batch,
836            },
837            page_token: None,
838            metadata: Default::default(),
839            page_size: None,
840            options: Default::default(),
841        };
842
843        let mut iter = q.read();
844        let err = iter.next().await.expect("should return error").unwrap_err();
845        assert!(matches!(err, RowError::InvalidRowFormat(_)), "{err:?}");
846        assert!(iter.next().await.is_none(), "{iter:?}");
847        Ok(())
848    }
849}
google_cloud_bigquery/query/iterator.rs

google_cloud_bigquery/query/
iterator.rs
