From f08dac3c2ff1d723b7a65ce3c6ade699716a6b78 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 7 Dec 2023 17:59:42 -0800 Subject: [PATCH 1/2] fix: keep `RowIterator.total_rows` populated after iteration This was being reset in some cases when the rows were all available in the first page of results. --- google/cloud/bigquery/table.py | 4 ++-- tests/unit/test_table.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index dca9f7962..fedb38fd5 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -2997,9 +2997,9 @@ def _rows_page_start(iterator, page, response): page._columns = _row_iterator_page_columns(iterator._schema, response) total_rows = response.get("totalRows") + # Don't reset total_rows is it's not present in the next API response. if total_rows is not None: - total_rows = int(total_rows) - iterator._total_rows = total_rows + iterator._total_rows = int(total_rows) # pylint: enable=unused-argument diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 85f335dd1..9b3d4fe84 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -2201,9 +2201,18 @@ def test_iterate_with_cached_first_page(self): path = "/foo" api_request = mock.Mock(return_value={"rows": rows}) row_iterator = self._make_one( - _mock_client(), api_request, path, schema, first_page_response=first_page + _mock_client(), + api_request, + path, + schema, + first_page_response=first_page, + total_rows=4, ) + self.assertEqual(row_iterator.total_rows, 4) rows = list(row_iterator) + # Total rows should be maintained, even though subsequent API calls + # don't include it. + self.assertEqual(row_iterator.total_rows, 4) self.assertEqual(len(rows), 4) self.assertEqual(rows[0].age, 27) self.assertEqual(rows[1].age, 28) From a376bd6b2d39c99bfc2f8c7f783bfe01a905ed1a Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 8 Dec 2023 15:04:04 -0600 Subject: [PATCH 2/2] Update google/cloud/bigquery/table.py Co-authored-by: Anthonios Partheniou --- google/cloud/bigquery/table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index fedb38fd5..70e601714 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -2997,7 +2997,7 @@ def _rows_page_start(iterator, page, response): page._columns = _row_iterator_page_columns(iterator._schema, response) total_rows = response.get("totalRows") - # Don't reset total_rows is it's not present in the next API response. + # Don't reset total_rows if it's not present in the next API response. if total_rows is not None: iterator._total_rows = int(total_rows)