fix: table iterator should not use bqstorage when page_size is not None (#2154)

chelsea-lin · web-flow · commit e89a707b1621 · 2025-04-03T11:45:08.000-07:00
* fix: table iterator should not use bqstorage when page_size is not None

* fix dbapi cursor tests
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
@@ -1873,6 +1873,11 @@ def total_bytes_processed(self) -> Optional[int]:
         """total bytes processed from job statistics, if present."""
         return self._total_bytes_processed
 
+    @property
+    def page_size(self) -> Optional[int]:
+        """The maximum number of rows in each page of results from this request, if present."""
+        return self._page_size
+
     def _is_almost_completely_cached(self):
         """Check if all results are completely cached.
 
@@ -1924,7 +1929,7 @@ def _should_use_bqstorage(self, bqstorage_client, create_bqstorage_client):
         if self._is_almost_completely_cached():
             return False
 
-        if self.max_results is not None:
+        if self.max_results is not None or self.page_size is not None:
             return False
 
         try:
@@ -1994,7 +1999,9 @@ def _maybe_warn_max_results(
             bqstorage_client:
                 The BigQuery Storage client intended to use for downloading result rows.
         """
-        if bqstorage_client is not None and self.max_results is not None:
+        if bqstorage_client is not None and (
+            self.max_results is not None or self.page_size is not None
+        ):
             warnings.warn(
                 "Cannot use bqstorage_client if max_results is set, "
                 "reverting to fetching data with the REST endpoint.",
diff --git a/tests/unit/test_dbapi_cursor.py b/tests/unit/test_dbapi_cursor.py
@@ -161,6 +161,7 @@ def _mock_rows(
             mock_rows,
         )
         mock_rows.max_results = None
+        mock_rows.page_size = None
         type(mock_rows).job_id = mock.PropertyMock(return_value="test-job-id")
         type(mock_rows).location = mock.PropertyMock(return_value="test-location")
         type(mock_rows).num_dml_affected_rows = mock.PropertyMock(
diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
@@ -2693,6 +2693,13 @@ def test__should_use_bqstorage_returns_false_if_max_results_set(self):
         )
         self.assertFalse(result)
 
+    def test__should_use_bqstorage_returns_false_if_page_size_set(self):
+        iterator = self._make_one(page_size=10, first_page_response=None)  # not cached
+        result = iterator._should_use_bqstorage(
+            bqstorage_client=None, create_bqstorage_client=True
+        )
+        self.assertFalse(result)
+
     def test__should_use_bqstorage_returns_false_w_warning_if_missing_dependency(self):
         iterator = self._make_one(first_page_response=None)  # not cached
 

Original file line number	Diff line number	Diff line change
`@@ -161,6 +161,7 @@ def _mock_rows(`
`161`	`161`	`mock_rows,`
`162`	`162`	`)`
`163`	`163`	`mock_rows.max_results = None`
	`164`	`+ mock_rows.page_size = None`
`164`	`165`	`type(mock_rows).job_id = mock.PropertyMock(return_value="test-job-id")`
`165`	`166`	`type(mock_rows).location = mock.PropertyMock(return_value="test-location")`
`166`	`167`	`type(mock_rows).num_dml_affected_rows = mock.PropertyMock(`