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::error::RowError;
16use crate::query::{CompleteQuery, Row, Schema};
17use google_cloud_bigquery_v2::client::JobService;
18use google_cloud_bigquery_v2::model::{GetQueryResultsRequest, JobReference};
19use std::collections::VecDeque;
20use std::sync::Arc;
21
22pub type Result<T> = std::result::Result<T, RowError>;
23/// An iterator over rows returned by a query.
24///
25/// [`CompleteQuery::read()`](crate::query::CompleteQuery::read) returns a
26/// `RowIterator`.
27///
28/// `RowIterator` yields rows from the query result set, automatically requesting
29/// subsequent pages from BigQuery as needed until the entire result set has been consumed.
30///
31/// # Example
32///
33/// ```
34/// # use google_cloud_bigquery::client::BigQuery;
35/// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
36/// let mut rows = client
37///     .query("SELECT name, state FROM `bigquery-public-data.usa_names.usa_1910_2013` LIMIT 10")
38///     .until_done()
39///     .await?
40///     .read();
41///
42/// while let Some(row) = rows.next().await.transpose()? {
43///     let name: String = row.get("name")?;
44///     let state: String = row.get("state")?;
45///     println!("{name} from {state}");
46/// }
47/// # Ok(())
48/// # }
49/// ```
50#[derive(Debug)]
51pub struct RowIterator {
52    job_service: Arc<JobService>,
53    job_ref: Option<JobReference>,
54    schema: Arc<Schema>,
55    page_token: Option<String>,
56    rows: VecDeque<wkt::Struct>,
57    page_size: Option<u32>,
58}
59
60impl RowIterator {
61    pub(crate) fn new(q: CompleteQuery) -> Self {
62        // DDL/DML queries have no schema.
63        let schema = Arc::new(Schema::new(q.metadata.schema.unwrap_or_default()));
64        Self {
65            job_service: q.job_service,
66            job_ref: q.job_ref,
67            schema,
68            page_token: q.page_token,
69            rows: q.cached_rows,
70            page_size: q.page_size,
71        }
72    }
73
74    /// Sets the maximum number of rows to buffer in memory.
75    ///
76    /// This controls page size during network fetches when streaming result
77    /// sets.
78    ///
79    /// # Example
80    ///
81    /// ```
82    /// # use google_cloud_bigquery::client::BigQuery;
83    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
84    /// let mut rows = client
85    ///     .query("SELECT 1 AS n")
86    ///     .until_done()
87    ///     .await?
88    ///     .read()
89    ///     .set_page_size(500);
90    /// # Ok(())
91    /// # }
92    /// ```
93    pub fn set_page_size(mut self, page_size: u32) -> Self {
94        self.page_size = Some(page_size);
95        self
96    }
97
98    /// Fetches the next [`Row`](crate::query::Row) from the result set.
99    ///
100    /// Returns `None` when all rows have been retrieved.
101    ///
102    /// # Example
103    ///
104    /// ```
105    /// # use google_cloud_bigquery::query::RowIterator;
106    /// # async fn sample(mut rows: RowIterator) -> anyhow::Result<()> {
107    /// while let Some(row) = rows.next().await.transpose()? {
108    ///     let msg: String = row.get("msg")?;
109    ///     println!("Message: {msg}");
110    /// }
111    /// # Ok(())
112    /// # }
113    /// ```
114    pub async fn next(&mut self) -> Option<Result<Row>> {
115        loop {
116            if let Some(raw_row) = self.rows.pop_front() {
117                return Some(Row::try_new(raw_row, &self.schema));
118            }
119
120            if let Err(e) = self.try_fetch_page().await {
121                return Some(Err(e));
122            }
123
124            if self.rows.is_empty() && self.page_token.is_none() {
125                return None;
126            }
127        }
128    }
129
130    async fn try_fetch_page(&mut self) -> Result<()> {
131        let Some(token) = self.page_token.as_deref() else {
132            return Ok(());
133        };
134
135        let (fetched_rows, next_token) = self.fetch_page(token).await?;
136        if fetched_rows.is_empty() && next_token.as_deref() == Some(token) {
137            self.page_token = None;
138        } else {
139            self.page_token = next_token;
140            self.rows.extend(fetched_rows);
141        }
142        Ok(())
143    }
144
145    // Fetches the next page of results and the next page token.
146    async fn fetch_page(&self, token: &str) -> Result<(Vec<wkt::Struct>, Option<String>)> {
147        let job_ref = self.job_ref.as_ref().expect(
148            "only queries with a job reference should have page tokens and can fetch more pages",
149        );
150
151        let mut req = GetQueryResultsRequest::new()
152            .set_project_id(job_ref.project_id.clone())
153            .set_or_clear_max_results(self.page_size)
154            .set_job_id(job_ref.job_id.clone())
155            .set_page_token(token)
156            .set_format_options(
157                google_cloud_bigquery_v2::model::DataFormatOptions::new()
158                    .set_use_int64_timestamp(true),
159            );
160        if let Some(location) = job_ref.location.clone() {
161            req = req.set_location(location);
162        }
163
164        let res = self
165            .job_service
166            .get_query_results()
167            .with_request(req)
168            .send()
169            .await?;
170
171        let page_token = if res.page_token.is_empty() {
172            None
173        } else {
174            Some(res.page_token)
175        };
176
177        Ok((res.rows, page_token))
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    use crate::query::tests::{MockJobService, create_job_service};
185    use google_cloud_bigquery_v2::model::{
186        DataFormatOptions, GetQueryResultsResponse, JobReference, QueryResponse, TableFieldSchema,
187        TableSchema,
188    };
189    use google_cloud_gax::error::Error as GaxError;
190    use google_cloud_gax::error::rpc::{Code, Status};
191    use google_cloud_gax::response::Response;
192    use serde_json::{Map, json};
193    use std::sync::Arc;
194
195    type TestResult = anyhow::Result<()>;
196
197    fn create_test_schema() -> TableSchema {
198        TableSchema::new().set_fields([TableFieldSchema::new()
199            .set_name("col")
200            .set_type("STRING")
201            .set_mode("NULLABLE")])
202    }
203
204    fn create_test_row(val: &str) -> wkt::Struct {
205        Map::from_iter([("f".to_string(), json!([{ "v": val }]))])
206    }
207
208    fn create_test_job_ref() -> JobReference {
209        JobReference::new()
210            .set_project_id("test_project")
211            .set_job_id("test_job")
212    }
213
214    fn create_test_job_ref_with_location(location: &str) -> JobReference {
215        create_test_job_ref().set_location(location)
216    }
217
218    fn create_test_complete_query(
219        job_service: Arc<JobService>,
220        job_ref: Option<JobReference>,
221        rows: Vec<wkt::Struct>,
222        page_token: Option<String>,
223    ) -> CompleteQuery {
224        let mut res = QueryResponse::new()
225            .set_schema(create_test_schema())
226            .set_rows(rows);
227        if let Some(token) = page_token {
228            res = res.set_page_token(token);
229        }
230        if let Some(job_ref) = job_ref {
231            res = res.set_job_reference(job_ref);
232        }
233        CompleteQuery::from_query_response(job_service, res, None)
234    }
235
236    #[tokio::test]
237    async fn test_row_iterator_empty_no_token() -> TestResult {
238        let job_service = create_job_service(MockJobService::new());
239        let q = create_test_complete_query(job_service, Some(create_test_job_ref()), vec![], None);
240        let mut iter = q.read();
241        assert!(iter.next().await.is_none(), "{iter:?}");
242        Ok(())
243    }
244
245    #[tokio::test]
246    async fn test_row_iterator_cached_rows_only() -> TestResult {
247        let job_service = create_job_service(MockJobService::new());
248        let rows = vec![create_test_row("first"), create_test_row("second")];
249        let q = create_test_complete_query(job_service, Some(create_test_job_ref()), rows, None);
250        let mut iter = q.read();
251
252        let row1 = iter.next().await.expect("should have row 1")?;
253        assert_eq!(row1.get::<String, _>("col")?, "first");
254
255        let row2 = iter.next().await.expect("should have row 2")?;
256        assert_eq!(row2.get::<String, _>("col")?, "second");
257
258        assert!(iter.next().await.is_none(), "{iter:?}");
259        Ok(())
260    }
261
262    #[tokio::test]
263    async fn test_row_iterator_row_conversion_error() -> TestResult {
264        let job_service = create_job_service(MockJobService::new());
265        let invalid_row = Map::from_iter([("f".to_string(), json!([]))]);
266        let q = create_test_complete_query(
267            job_service,
268            Some(create_test_job_ref()),
269            vec![invalid_row],
270            None,
271        );
272        let mut iter = q.read();
273
274        let err = iter.next().await.expect("should return error").unwrap_err();
275        assert!(matches!(err, RowError::InvalidRowFormat(_)), "{err:?}");
276        assert!(iter.next().await.is_none(), "{iter:?}");
277        Ok(())
278    }
279
280    #[tokio::test]
281    async fn test_row_iterator_fetch_multiple_pages() -> TestResult {
282        let mut mock = MockJobService::new();
283        let mut seq = mockall::Sequence::new();
284
285        mock.expect_get_query_results()
286            .in_sequence(&mut seq)
287            .times(1)
288            .returning(|req, _| {
289                assert_eq!(req.page_token, "token_page_1");
290                let res = GetQueryResultsResponse::new()
291                    .set_rows(vec![
292                        create_test_row("page_1_row_1"),
293                        create_test_row("page_1_row_2"),
294                    ])
295                    .set_page_token("token_page_2");
296                Ok(Response::from(res))
297            });
298
299        mock.expect_get_query_results()
300            .in_sequence(&mut seq)
301            .times(1)
302            .returning(|req, _| {
303                assert_eq!(req.page_token, "token_page_2");
304                let res = GetQueryResultsResponse::new()
305                    .set_rows(vec![create_test_row("page_2_row_1")])
306                    .set_page_token("");
307                Ok(Response::from(res))
308            });
309
310        let job_service = create_job_service(mock);
311        let q = create_test_complete_query(
312            job_service,
313            Some(create_test_job_ref()),
314            vec![create_test_row("cached_row")],
315            Some("token_page_1".to_string()),
316        );
317        let mut iter = q.read();
318
319        let row1 = iter.next().await.expect("should have row 1")?;
320        assert_eq!(row1.get::<String, _>("col")?, "cached_row");
321
322        let row2 = iter.next().await.expect("should have row 2")?;
323        assert_eq!(row2.get::<String, _>("col")?, "page_1_row_1");
324
325        let row3 = iter.next().await.expect("should have row 3")?;
326        assert_eq!(row3.get::<String, _>("col")?, "page_1_row_2");
327
328        let row4 = iter.next().await.expect("should have row 4")?;
329        assert_eq!(row4.get::<String, _>("col")?, "page_2_row_1");
330
331        assert!(iter.next().await.is_none(), "{iter:?}");
332        Ok(())
333    }
334
335    #[tokio::test]
336    async fn test_row_iterator_fetch_page_request_parameters() -> TestResult {
337        let mut mock = MockJobService::new();
338        mock.expect_get_query_results()
339            .times(1)
340            .returning(|req, _| {
341                assert_eq!(req.project_id, "test_project");
342                assert_eq!(req.job_id, "test_job");
343                assert_eq!(req.page_token, "token_1");
344                assert_eq!(req.max_results, Some(50));
345                assert_eq!(req.location, "us-east1");
346                assert_eq!(
347                    req.format_options,
348                    Some(DataFormatOptions::new().set_use_int64_timestamp(true))
349                );
350                let res = GetQueryResultsResponse::new()
351                    .set_rows(vec![create_test_row("page_row")])
352                    .set_page_token("");
353                Ok(Response::from(res))
354            });
355
356        let job_service = create_job_service(mock);
357        let q = create_test_complete_query(
358            job_service,
359            Some(create_test_job_ref_with_location("us-east1")),
360            vec![],
361            Some("token_1".to_string()),
362        );
363        let mut iter = q.read().set_page_size(50);
364
365        let row = iter.next().await.expect("should have row")?;
366        assert_eq!(row.get::<String, _>("col")?, "page_row");
367
368        assert!(iter.next().await.is_none(), "{iter:?}");
369        Ok(())
370    }
371
372    #[tokio::test]
373    async fn test_row_iterator_inherits_page_size() -> TestResult {
374        let mut mock = MockJobService::new();
375        mock.expect_get_query_results()
376            .times(1)
377            .returning(|req, _| {
378                assert_eq!(req.max_results, Some(25));
379                let res = GetQueryResultsResponse::new()
380                    .set_rows(vec![create_test_row("page_row")])
381                    .set_page_token("");
382                Ok(Response::from(res))
383            });
384
385        let job_service = create_job_service(mock);
386        let res = QueryResponse::new()
387            .set_schema(create_test_schema())
388            .set_job_reference(create_test_job_ref())
389            .set_page_token("token_1");
390        let q = CompleteQuery::from_query_response(job_service, res, Some(25));
391        let mut iter = q.read();
392
393        let row = iter.next().await.expect("should have row")?;
394        assert_eq!(row.get::<String, _>("col")?, "page_row");
395        assert!(iter.next().await.is_none(), "{iter:?}");
396        Ok(())
397    }
398
399    #[tokio::test]
400    async fn test_row_iterator_fetch_page_rpc_error() -> TestResult {
401        let mut mock = MockJobService::new();
402        mock.expect_get_query_results().returning(|_, _| {
403            let status = Status::default()
404                .set_code(Code::Unavailable)
405                .set_message("temporary service error");
406            Err(GaxError::service(status))
407        });
408
409        let job_service = create_job_service(mock);
410        let q = create_test_complete_query(
411            job_service,
412            Some(create_test_job_ref()),
413            vec![],
414            Some("token_err".to_string()),
415        );
416        let mut iter = q.read();
417
418        let err = iter.next().await.expect("should return error").unwrap_err();
419        assert!(matches!(err, RowError::Rpc { .. }), "{err:?}");
420        assert!(
421            err.to_string().contains("temporary service error"),
422            "{err:?}"
423        );
424        Ok(())
425    }
426
427    #[tokio::test]
428    async fn test_row_iterator_empty_page_handling() -> TestResult {
429        let mut mock = MockJobService::new();
430        let mut seq = mockall::Sequence::new();
431
432        // empty page with a new token: should continue fetching
433        mock.expect_get_query_results()
434            .times(1)
435            .in_sequence(&mut seq)
436            .returning(|req, _| {
437                assert_eq!(req.page_token, "token_1");
438                let res = GetQueryResultsResponse::new()
439                    .set_rows(Vec::<wkt::Struct>::new())
440                    .set_page_token("token_2");
441                Ok(Response::from(res))
442            });
443
444        // next page returns rows and advances to token_3
445        mock.expect_get_query_results()
446            .times(1)
447            .in_sequence(&mut seq)
448            .returning(|req, _| {
449                assert_eq!(req.page_token, "token_2");
450                let res = GetQueryResultsResponse::new()
451                    .set_rows(vec![create_test_row("page2_row")])
452                    .set_page_token("token_3");
453                Ok(Response::from(res))
454            });
455
456        // empty page with the same token: should terminate pagination
457        mock.expect_get_query_results()
458            .times(1)
459            .in_sequence(&mut seq)
460            .returning(|req, _| {
461                assert_eq!(req.page_token, "token_3");
462                let res = GetQueryResultsResponse::new()
463                    .set_rows(Vec::<wkt::Struct>::new())
464                    .set_page_token("token_3");
465                Ok(Response::from(res))
466            });
467
468        let job_service = create_job_service(mock);
469        let q = create_test_complete_query(
470            job_service,
471            Some(create_test_job_ref()),
472            vec![],
473            Some("token_1".to_string()),
474        );
475        let mut iter = q.read();
476
477        let row = iter.next().await.expect("should have row")?;
478        assert_eq!(row.get::<String, _>("col")?, "page2_row");
479        assert!(iter.next().await.is_none(), "{iter:?}");
480        Ok(())
481    }
482}
google_cloud_bigquery/query/iterator.rs

google_cloud_bigquery/query/
iterator.rs
