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

google_cloud_bigquery/query/
iterator.rs
