google-cloud-bigquery 0.18.0

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::error::RowError;
use crate::query::{CompleteQuery, Row, Schema};
use google_cloud_bigquery_v2::client::JobService;
use google_cloud_bigquery_v2::model::{GetQueryResultsRequest, JobReference};
use std::collections::VecDeque;
use std::sync::Arc;

pub type Result<T> = std::result::Result<T, RowError>;
/// An iterator over rows returned by a query.
///
/// [`CompleteQuery::read()`](crate::query::CompleteQuery::read) returns a
/// `RowIterator`.
///
/// `RowIterator` yields rows from the query result set, automatically requesting
/// subsequent pages from BigQuery as needed until the entire result set has been consumed.
///
/// # Example
///
/// ```
/// # use google_cloud_bigquery::client::BigQuery;
/// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
/// let mut rows = client
///     .query("SELECT name, state FROM `bigquery-public-data.usa_names.usa_1910_2013` LIMIT 10")
///     .until_done()
///     .await?
///     .read();
///
/// while let Some(row) = rows.next().await.transpose()? {
///     let name: String = row.get("name")?;
///     let state: String = row.get("state")?;
///     println!("{name} from {state}");
/// }
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
pub struct RowIterator {
    job_service: Arc<JobService>,
    job_ref: Option<JobReference>,
    schema: Arc<Schema>,
    page_token: Option<String>,
    rows: VecDeque<wkt::Struct>,
    page_size: Option<u32>,
}

impl RowIterator {
    pub(crate) fn new(q: CompleteQuery) -> Self {
        // DDL/DML queries have no schema.
        let schema = Arc::new(Schema::new(q.metadata.schema.unwrap_or_default()));
        Self {
            job_service: q.job_service,
            job_ref: q.job_ref,
            schema,
            page_token: q.page_token,
            rows: q.cached_rows,
            page_size: q.page_size,
        }
    }

    /// Sets the maximum number of rows to buffer in memory.
    ///
    /// This controls page size during network fetches when streaming result
    /// sets.
    ///
    /// # Example
    ///
    /// ```
    /// # use google_cloud_bigquery::client::BigQuery;
    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
    /// let mut rows = client
    ///     .query("SELECT 1 AS n")
    ///     .until_done()
    ///     .await?
    ///     .read()
    ///     .set_page_size(500);
    /// # Ok(())
    /// # }
    /// ```
    pub fn set_page_size(mut self, page_size: u32) -> Self {
        self.page_size = Some(page_size);
        self
    }

    /// Fetches the next [`Row`](crate::query::Row) from the result set.
    ///
    /// Returns `None` when all rows have been retrieved.
    ///
    /// # Example
    ///
    /// ```
    /// # use google_cloud_bigquery::query::RowIterator;
    /// # async fn sample(mut rows: RowIterator) -> anyhow::Result<()> {
    /// while let Some(row) = rows.next().await.transpose()? {
    ///     let msg: String = row.get("msg")?;
    ///     println!("Message: {msg}");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn next(&mut self) -> Option<Result<Row>> {
        loop {
            if let Some(raw_row) = self.rows.pop_front() {
                return Some(Row::try_new(raw_row, &self.schema));
            }

            if let Err(e) = self.try_fetch_page().await {
                return Some(Err(e));
            }

            if self.rows.is_empty() && self.page_token.is_none() {
                return None;
            }
        }
    }

    async fn try_fetch_page(&mut self) -> Result<()> {
        let Some(token) = self.page_token.as_deref() else {
            return Ok(());
        };

        let (fetched_rows, next_token) = self.fetch_page(token).await?;
        if fetched_rows.is_empty() && next_token.as_deref() == Some(token) {
            self.page_token = None;
        } else {
            self.page_token = next_token;
            self.rows.extend(fetched_rows);
        }
        Ok(())
    }

    // Fetches the next page of results and the next page token.
    async fn fetch_page(&self, token: &str) -> Result<(Vec<wkt::Struct>, Option<String>)> {
        let job_ref = self.job_ref.as_ref().expect(
            "only queries with a job reference should have page tokens and can fetch more pages",
        );

        let mut req = GetQueryResultsRequest::new()
            .set_project_id(job_ref.project_id.clone())
            .set_or_clear_max_results(self.page_size)
            .set_job_id(job_ref.job_id.clone())
            .set_page_token(token)
            .set_format_options(
                google_cloud_bigquery_v2::model::DataFormatOptions::new()
                    .set_use_int64_timestamp(true),
            );
        if let Some(location) = job_ref.location.clone() {
            req = req.set_location(location);
        }

        let res = self
            .job_service
            .get_query_results()
            .with_request(req)
            .send()
            .await?;

        let page_token = if res.page_token.is_empty() {
            None
        } else {
            Some(res.page_token)
        };

        Ok((res.rows, page_token))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query::tests::{MockJobService, create_job_service};
    use google_cloud_bigquery_v2::model::{
        DataFormatOptions, GetQueryResultsResponse, JobReference, QueryResponse, TableFieldSchema,
        TableSchema,
    };
    use google_cloud_gax::error::Error as GaxError;
    use google_cloud_gax::error::rpc::{Code, Status};
    use google_cloud_gax::response::Response;
    use serde_json::{Map, json};
    use std::sync::Arc;

    type TestResult = anyhow::Result<()>;

    fn create_test_schema() -> TableSchema {
        TableSchema::new().set_fields([TableFieldSchema::new()
            .set_name("col")
            .set_type("STRING")
            .set_mode("NULLABLE")])
    }

    fn create_test_row(val: &str) -> wkt::Struct {
        Map::from_iter([("f".to_string(), json!([{ "v": val }]))])
    }

    fn create_test_job_ref() -> JobReference {
        JobReference::new()
            .set_project_id("test_project")
            .set_job_id("test_job")
    }

    fn create_test_job_ref_with_location(location: &str) -> JobReference {
        create_test_job_ref().set_location(location)
    }

    fn create_test_complete_query(
        job_service: Arc<JobService>,
        job_ref: Option<JobReference>,
        rows: Vec<wkt::Struct>,
        page_token: Option<String>,
    ) -> CompleteQuery {
        let mut res = QueryResponse::new()
            .set_schema(create_test_schema())
            .set_rows(rows);
        if let Some(token) = page_token {
            res = res.set_page_token(token);
        }
        if let Some(job_ref) = job_ref {
            res = res.set_job_reference(job_ref);
        }
        CompleteQuery::from_query_response(job_service, res, None)
    }

    #[tokio::test]
    async fn test_row_iterator_empty_no_token() -> TestResult {
        let job_service = create_job_service(MockJobService::new());
        let q = create_test_complete_query(job_service, Some(create_test_job_ref()), vec![], None);
        let mut iter = q.read();
        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_cached_rows_only() -> TestResult {
        let job_service = create_job_service(MockJobService::new());
        let rows = vec![create_test_row("first"), create_test_row("second")];
        let q = create_test_complete_query(job_service, Some(create_test_job_ref()), rows, None);
        let mut iter = q.read();

        let row1 = iter.next().await.expect("should have row 1")?;
        assert_eq!(row1.get::<String, _>("col")?, "first");

        let row2 = iter.next().await.expect("should have row 2")?;
        assert_eq!(row2.get::<String, _>("col")?, "second");

        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_row_conversion_error() -> TestResult {
        let job_service = create_job_service(MockJobService::new());
        let invalid_row = Map::from_iter([("f".to_string(), json!([]))]);
        let q = create_test_complete_query(
            job_service,
            Some(create_test_job_ref()),
            vec![invalid_row],
            None,
        );
        let mut iter = q.read();

        let err = iter.next().await.expect("should return error").unwrap_err();
        assert!(matches!(err, RowError::InvalidRowFormat(_)), "{err:?}");
        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_fetch_multiple_pages() -> TestResult {
        let mut mock = MockJobService::new();
        let mut seq = mockall::Sequence::new();

        mock.expect_get_query_results()
            .in_sequence(&mut seq)
            .times(1)
            .returning(|req, _| {
                assert_eq!(req.page_token, "token_page_1");
                let res = GetQueryResultsResponse::new()
                    .set_rows(vec![
                        create_test_row("page_1_row_1"),
                        create_test_row("page_1_row_2"),
                    ])
                    .set_page_token("token_page_2");
                Ok(Response::from(res))
            });

        mock.expect_get_query_results()
            .in_sequence(&mut seq)
            .times(1)
            .returning(|req, _| {
                assert_eq!(req.page_token, "token_page_2");
                let res = GetQueryResultsResponse::new()
                    .set_rows(vec![create_test_row("page_2_row_1")])
                    .set_page_token("");
                Ok(Response::from(res))
            });

        let job_service = create_job_service(mock);
        let q = create_test_complete_query(
            job_service,
            Some(create_test_job_ref()),
            vec![create_test_row("cached_row")],
            Some("token_page_1".to_string()),
        );
        let mut iter = q.read();

        let row1 = iter.next().await.expect("should have row 1")?;
        assert_eq!(row1.get::<String, _>("col")?, "cached_row");

        let row2 = iter.next().await.expect("should have row 2")?;
        assert_eq!(row2.get::<String, _>("col")?, "page_1_row_1");

        let row3 = iter.next().await.expect("should have row 3")?;
        assert_eq!(row3.get::<String, _>("col")?, "page_1_row_2");

        let row4 = iter.next().await.expect("should have row 4")?;
        assert_eq!(row4.get::<String, _>("col")?, "page_2_row_1");

        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_fetch_page_request_parameters() -> TestResult {
        let mut mock = MockJobService::new();
        mock.expect_get_query_results()
            .times(1)
            .returning(|req, _| {
                assert_eq!(req.project_id, "test_project");
                assert_eq!(req.job_id, "test_job");
                assert_eq!(req.page_token, "token_1");
                assert_eq!(req.max_results, Some(50));
                assert_eq!(req.location, "us-east1");
                assert_eq!(
                    req.format_options,
                    Some(DataFormatOptions::new().set_use_int64_timestamp(true))
                );
                let res = GetQueryResultsResponse::new()
                    .set_rows(vec![create_test_row("page_row")])
                    .set_page_token("");
                Ok(Response::from(res))
            });

        let job_service = create_job_service(mock);
        let q = create_test_complete_query(
            job_service,
            Some(create_test_job_ref_with_location("us-east1")),
            vec![],
            Some("token_1".to_string()),
        );
        let mut iter = q.read().set_page_size(50);

        let row = iter.next().await.expect("should have row")?;
        assert_eq!(row.get::<String, _>("col")?, "page_row");

        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_inherits_page_size() -> TestResult {
        let mut mock = MockJobService::new();
        mock.expect_get_query_results()
            .times(1)
            .returning(|req, _| {
                assert_eq!(req.max_results, Some(25));
                let res = GetQueryResultsResponse::new()
                    .set_rows(vec![create_test_row("page_row")])
                    .set_page_token("");
                Ok(Response::from(res))
            });

        let job_service = create_job_service(mock);
        let res = QueryResponse::new()
            .set_schema(create_test_schema())
            .set_job_reference(create_test_job_ref())
            .set_page_token("token_1");
        let q = CompleteQuery::from_query_response(job_service, res, Some(25));
        let mut iter = q.read();

        let row = iter.next().await.expect("should have row")?;
        assert_eq!(row.get::<String, _>("col")?, "page_row");
        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_fetch_page_rpc_error() -> TestResult {
        let mut mock = MockJobService::new();
        mock.expect_get_query_results().returning(|_, _| {
            let status = Status::default()
                .set_code(Code::Unavailable)
                .set_message("temporary service error");
            Err(GaxError::service(status))
        });

        let job_service = create_job_service(mock);
        let q = create_test_complete_query(
            job_service,
            Some(create_test_job_ref()),
            vec![],
            Some("token_err".to_string()),
        );
        let mut iter = q.read();

        let err = iter.next().await.expect("should return error").unwrap_err();
        assert!(matches!(err, RowError::Rpc { .. }), "{err:?}");
        assert!(
            err.to_string().contains("temporary service error"),
            "{err:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_row_iterator_empty_page_handling() -> TestResult {
        let mut mock = MockJobService::new();
        let mut seq = mockall::Sequence::new();

        // empty page with a new token: should continue fetching
        mock.expect_get_query_results()
            .times(1)
            .in_sequence(&mut seq)
            .returning(|req, _| {
                assert_eq!(req.page_token, "token_1");
                let res = GetQueryResultsResponse::new()
                    .set_rows(Vec::<wkt::Struct>::new())
                    .set_page_token("token_2");
                Ok(Response::from(res))
            });

        // next page returns rows and advances to token_3
        mock.expect_get_query_results()
            .times(1)
            .in_sequence(&mut seq)
            .returning(|req, _| {
                assert_eq!(req.page_token, "token_2");
                let res = GetQueryResultsResponse::new()
                    .set_rows(vec![create_test_row("page2_row")])
                    .set_page_token("token_3");
                Ok(Response::from(res))
            });

        // empty page with the same token: should terminate pagination
        mock.expect_get_query_results()
            .times(1)
            .in_sequence(&mut seq)
            .returning(|req, _| {
                assert_eq!(req.page_token, "token_3");
                let res = GetQueryResultsResponse::new()
                    .set_rows(Vec::<wkt::Struct>::new())
                    .set_page_token("token_3");
                Ok(Response::from(res))
            });

        let job_service = create_job_service(mock);
        let q = create_test_complete_query(
            job_service,
            Some(create_test_job_ref()),
            vec![],
            Some("token_1".to_string()),
        );
        let mut iter = q.read();

        let row = iter.next().await.expect("should have row")?;
        assert_eq!(row.get::<String, _>("col")?, "page2_row");
        assert!(iter.next().await.is_none(), "{iter:?}");
        Ok(())
    }
}