Repository: googleapis/python-bigquery
Pull request #1942 (2024-06-03 16:29:14)

perf: if `page_size` or `max_results` is set on `QueryJob.result()`, use to download first page of results
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes internal issue 344008814
 🦕



Patch:
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
index 4234767f..1c222f2d 100644
--- a/google/cloud/bigquery/client.py
+++ b/google/cloud/bigquery/client.py
@@ -109,6 +109,7 @@ from google.cloud.bigquery.retry import (
     DEFAULT_RETRY,
     DEFAULT_TIMEOUT,
     DEFAULT_GET_JOB_TIMEOUT,
+    POLLING_DEFAULT_VALUE,
 )
 from google.cloud.bigquery.routine import Routine
 from google.cloud.bigquery.routine import RoutineReference
@@ -1963,6 +1964,7 @@ class Client(ClientWithProject):
         timeout_ms: Optional[int] = None,
         location: Optional[str] = None,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
+        page_size: int = 0,
     ) -> _QueryResults:
         """Get the query results object for a query job.
 
@@ -1981,13 +1983,16 @@ class Client(ClientWithProject):
                 before using ``retry``. If set, this connection timeout may be
                 increased to a minimum value. This prevents retries on what
                 would otherwise be a successful response.
+            page_size (int):
+                Maximum number of rows in a single response. See maxResults in
+                the jobs.getQueryResults REST API.
 
         Returns:
             google.cloud.bigquery.query._QueryResults:
                 A new ``_QueryResults`` instance.
         """
 
-        extra_params: Dict[str, Any] = {"maxResults": 0}
+        extra_params: Dict[str, Any] = {"maxResults": page_size}
 
         if timeout is not None:
             if not isinstance(timeout, (int, float)):
@@ -1995,6 +2000,9 @@ class Client(ClientWithProject):
             else:
                 timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
 
+        if page_size > 0:
+            extra_params["formatOptions.useInt64Timestamp"] = True
+
         if project is None:
             project = self.project
 
@@ -3504,7 +3512,7 @@ class Client(ClientWithProject):
         location: Optional[str] = None,
         project: Optional[str] = None,
         api_timeout: TimeoutType = DEFAULT_TIMEOUT,
-        wait_timeout: TimeoutType = None,
+        wait_timeout: Union[Optional[float], object] = POLLING_DEFAULT_VALUE,
         retry: retries.Retry = DEFAULT_RETRY,
         job_retry: retries.Retry = DEFAULT_JOB_RETRY,
         page_size: Optional[int] = None,
@@ -3538,10 +3546,12 @@ class Client(ClientWithProject):
             api_timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
-            wait_timeout (Optional[float]):
+            wait_timeout (Optional[Union[float, object]]):
                 The number of seconds to wait for the query to finish. If the
                 query doesn't finish before this timeout, the client attempts
-                to cancel the query.
+                to cancel the query. If unset, the underlying REST API calls
+                have timeouts, but we still wait indefinitely for the job to
+                finish.
             retry (Optional[google.api_core.retry.Retry]):
                 How to retry the RPC.  This only applies to making RPC
                 calls.  It isn't used to retry failed jobs.  This has
@@ -4128,11 +4138,6 @@ class Client(ClientWithProject):
         if start_index is not None:
             params["startIndex"] = start_index
 
-        # We don't call jobs.query with a page size, so if the user explicitly
-        # requests a certain size, invalidate the cache.
-        if page_size is not None:
-            first_page_response = None
-
         params["formatOptions.useInt64Timestamp"] = True
         row_iterator = RowIterator(
             client=self,
diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py
index 25b89c3d..a8530271 100644
--- a/google/cloud/bigquery/job/query.py
+++ b/google/cloud/bigquery/job/query.py
@@ -1385,7 +1385,10 @@ class QueryJob(_AsyncJob):
             raise
 
     def _reload_query_results(
-        self, retry: "retries.Retry" = DEFAULT_RETRY, timeout: Optional[float] = None
+        self,
+        retry: "retries.Retry" = DEFAULT_RETRY,
+        timeout: Optional[float] = None,
+        page_size: int = 0,
     ):
         """Refresh the cached query results unless already cached and complete.
 
@@ -1395,6 +1398,9 @@ class QueryJob(_AsyncJob):
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
+            page_size (int):
+                Maximum number of rows in a single response. See maxResults in
+                the jobs.getQueryResults REST API.
         """
         # Optimization: avoid a call to jobs.getQueryResults if it's already
         # been fetched, e.g. from jobs.query first page of results.
@@ -1425,7 +1431,14 @@ class QueryJob(_AsyncJob):
 
         # If an explicit timeout is not given, fall back to the transport timeout
         # stored in _blocking_poll() in the process of polling for job completion.
-        transport_timeout = timeout if timeout is not None else self._transport_timeout
+        if timeout is not None:
+            transport_timeout = timeout
+        else:
+            transport_timeout = self._transport_timeout
+
+            # Handle PollingJob._DEFAULT_VALUE.
+            if not isinstance(transport_timeout, (float, int)):
+                transport_timeout = None
 
         self._query_results = self._client._get_query_results(
             self.job_id,
@@ -1434,6 +1447,7 @@ class QueryJob(_AsyncJob):
             timeout_ms=timeout_ms,
             location=self.location,
             timeout=transport_timeout,
+            page_size=page_size,
         )
 
     def result(  # type: ignore  # (incompatible with supertype)
@@ -1515,11 +1529,25 @@ class QueryJob(_AsyncJob):
                 # actually correspond to a finished query job.
             )
 
+        # Setting max_results should be equivalent to setting page_size with
+        # regards to allowing the user to tune how many results to download
+        # while we wait for the query to finish. See internal issue:
+        # 344008814.
+        if page_size is None and max_results is not None:
+            page_size = max_results
+
         # When timeout has default sentinel value ``object()``, do not pass
         # anything to invoke default timeouts in subsequent calls.
-        kwargs: Dict[str, Union[_helpers.TimeoutType, object]] = {}
+        done_kwargs: Dict[str, Union[_helpers.TimeoutType, object]] = {}
+        reload_query_results_kwargs: Dict[str, Union[_helpers.TimeoutType, object]] = {}
+        list_rows_kwargs: Dict[str, Union[_helpers.TimeoutType, object]] = {}
         if type(timeout) is not object:
-            kwargs["timeout"] = timeout
+            done_kwargs["timeout"] = timeout
+            list_rows_kwargs["timeout"] = timeout
+            reload_query_results_kwargs["timeout"] = timeout
+
+        if page_size is not None:
+            reload_query_results_kwargs["page_size"] = page_size
 
         try:
             retry_do_query = getattr(self, "_retry_do_query", None)
@@ -1562,7 +1590,7 @@ class QueryJob(_AsyncJob):
                 # rateLimitExceeded errors are ambiguous. We want to know if
                 # the query job failed and not just the call to
                 # jobs.getQueryResults.
-                if self.done(retry=retry, **kwargs):
+                if self.done(retry=retry, **done_kwargs):
                     # If it's already failed, we might as well stop.
                     job_failed_exception = self.exception()
                     if job_failed_exception is not None:
@@ -1599,14 +1627,16 @@ class QueryJob(_AsyncJob):
                         # response from the REST API. This ensures we aren't
                         # making any extra API calls if the previous loop
                         # iteration fetched the finished job.
-                        self._reload_query_results(retry=retry, **kwargs)
+                        self._reload_query_results(
+                            retry=retry, **reload_query_results_kwargs
+                        )
                         return True
 
                 # Call jobs.getQueryResults with max results set to 0 just to
                 # wait for the query to finish. Unlike most methods,
                 # jobs.getQueryResults hangs as long as it can to ensure we
                 # know when the query has finished as soon as possible.
-                self._reload_query_results(retry=retry, **kwargs)
+                self._reload_query_results(retry=retry, **reload_query_results_kwargs)
 
                 # Even if the query is finished now according to
                 # jobs.getQueryResults, we'll want to reload the job status if
@@ -1679,8 +1709,9 @@ class QueryJob(_AsyncJob):
         # We know that there's at least 1 row, so only treat the response from
         # jobs.getQueryResults / jobs.query as the first page of the
         # RowIterator response if there are any rows in it. This prevents us
-        # from stopping the iteration early because we're missing rows and
-        # there's no next page token.
+        # from stopping the iteration early in the cases where we set
+        # maxResults=0. In that case, we're missing rows and there's no next
+        # page token.
         first_page_response = self._query_results._properties
         if "rows" not in first_page_response:
             first_page_response = None
@@ -1699,7 +1730,7 @@ class QueryJob(_AsyncJob):
             query_id=self.query_id,
             first_page_response=first_page_response,
             num_dml_affected_rows=self._query_results.num_dml_affected_rows,
-            **kwargs,
+            **list_rows_kwargs,
         )
         rows._preserve_order = _contains_order_by(self.query)
         return rows
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
index 57fc0d2b..faf827be 100644
--- a/google/cloud/bigquery/table.py
+++ b/google/cloud/bigquery/table.py
@@ -1759,7 +1759,7 @@ class RowIterator(HTTPIterator):
         if self._page_size is not None:
             if self.page_number and "startIndex" in params:
                 del params["startIndex"]
-            params["maxResults"] = self._page_size
+
         return self.api_request(
             method=self._HTTP_METHOD, path=self.path, query_params=params
         )


Test patch:
diff --git a/tests/unit/job/test_query.py b/tests/unit/job/test_query.py
index c7b2c5f9..66055dee 100644
--- a/tests/unit/job/test_query.py
+++ b/tests/unit/job/test_query.py
@@ -28,6 +28,7 @@ import requests
 from google.cloud.bigquery.client import _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS
 import google.cloud.bigquery._job_helpers
 import google.cloud.bigquery.query
+import google.cloud.bigquery.retry
 from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT
 from google.cloud.bigquery.table import _EmptyRowIterator
 
@@ -841,6 +842,22 @@ class TestQueryJob(_Base):
         assert isinstance(job.search_stats, SearchStats)
         assert job.search_stats.mode == "INDEX_USAGE_MODE_UNSPECIFIED"
 
+    def test_reload_query_results_uses_transport_timeout(self):
+        conn = make_connection({})
+        client = _make_client(self.PROJECT, connection=conn)
+        job = self._make_one(self.JOB_ID, self.QUERY, client)
+        job._transport_timeout = 123
+
+        job._reload_query_results()
+
+        query_results_path = f"/projects/{self.PROJECT}/queries/{self.JOB_ID}"
+        conn.api_request.assert_called_once_with(
+            method="GET",
+            path=query_results_path,
+            query_params={"maxResults": 0},
+            timeout=123,
+        )
+
     def test_result_reloads_job_state_until_done(self):
         """Verify that result() doesn't return until state == 'DONE'.
 
@@ -1053,7 +1070,7 @@ class TestQueryJob(_Base):
             method="GET",
             path=query_results_path,
             query_params={"maxResults": 0, "location": "EU"},
-            timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
+            timeout=None,
         )
         query_results_page_call = mock.call(
             method="GET",
@@ -1139,114 +1156,92 @@ class TestQueryJob(_Base):
             timeout=None,
         )
 
-    def test_result_with_done_jobs_query_response_and_page_size_invalidates_cache(self):
-        """We don't call jobs.query with a page size, so if the user explicitly
-        requests a certain size, invalidate the cache.
-        """
-        # Arrange
-        job_resource = self._make_resource(
-            started=True, ended=True, location="asia-northeast1"
-        )
-        query_resource_done = {
+    def test_result_with_max_results(self):
+        from google.cloud.bigquery.table import RowIterator
+
+        query_resource = {
             "jobComplete": True,
             "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
             "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
-            "rows": [{"f": [{"v": "abc"}]}],
-            "pageToken": "initial-page-token-shouldnt-be-used",
-            "totalRows": "4",
+            "totalRows": "10",
+            "pageToken": "first-page-token",
+            "rows": [
+                {"f": [{"v": "abc"}]},
+                {"f": [{"v": "def"}]},
+                {"f": [{"v": "ghi"}]},
+                {"f": [{"v": "jkl"}]},
+                {"f": [{"v": "mno"}]},
+                {"f": [{"v": "pqr"}]},
+                # Pretend these are very large rows, so the API doesn't return
+                # all of the rows we asked for in the first response.
+            ],
         }
         query_page_resource = {
-            "totalRows": 4,
-            "pageToken": "some-page-token",
+            "totalRows": "10",
+            "pageToken": None,
             "rows": [
-                {"f": [{"v": "row1"}]},
-                {"f": [{"v": "row2"}]},
-                {"f": [{"v": "row3"}]},
+                {"f": [{"v": "stu"}]},
+                {"f": [{"v": "vwx"}]},
+                {"f": [{"v": "yz0"}]},
             ],
         }
-        query_page_resource_2 = {"totalRows": 4, "rows": [{"f": [{"v": "row4"}]}]}
-        conn = make_connection(job_resource, query_page_resource, query_page_resource_2)
-        client = _make_client(self.PROJECT, connection=conn)
-        job = google.cloud.bigquery._job_helpers._to_query_job(
-            client,
-            "SELECT col1 FROM table",
-            request_config=None,
-            query_response=query_resource_done,
+        job_resource_running = self._make_resource(
+            started=True, ended=False, location="US"
         )
-        # We want job.result() to refresh the job state, so the conversion is
-        # always "PENDING", even if the job is finished.
-        assert job.state == "PENDING"
+        job_resource_done = self._make_resource(started=True, ended=True, location="US")
+        conn = make_connection(job_resource_done, query_resource, query_page_resource)
+        client = _make_client(self.PROJECT, connection=conn)
+        job = self._get_target_class().from_api_repr(job_resource_running, client)
 
-        # Act
-        result = job.result(page_size=3)
+        max_results = 9
+        result = job.result(max_results=max_results)
 
-        # Assert
-        actual_rows = list(result)
-        self.assertEqual(len(actual_rows), 4)
+        self.assertIsInstance(result, RowIterator)
+        self.assertEqual(result.total_rows, 10)
+
+        rows = list(result)
 
+        self.assertEqual(len(rows), 9)
+        jobs_get_path = f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}"
+        jobs_get_call = mock.call(
+            method="GET",
+            path=jobs_get_path,
+            query_params={"projection": "full", "location": "US"},
+            timeout=DEFAULT_GET_JOB_TIMEOUT,
+        )
         query_results_path = f"/projects/{self.PROJECT}/queries/{self.JOB_ID}"
-        query_page_1_call = mock.call(
+        query_page_waiting_call = mock.call(
             method="GET",
             path=query_results_path,
             query_params={
-                "maxResults": 3,
-                "fields": _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS,
-                "location": "asia-northeast1",
+                # Waiting for the results should set maxResults and cache the
+                # first page if page_size is set. This allows customers to
+                # more finely tune when we fallback to the BQ Storage API.
+                # See internal issue: 344008814.
+                "maxResults": max_results,
                 "formatOptions.useInt64Timestamp": True,
+                "location": "US",
             },
             timeout=None,
         )
         query_page_2_call = mock.call(
+            timeout=None,
             method="GET",
             path=query_results_path,
             query_params={
-                "pageToken": "some-page-token",
+                "pageToken": "first-page-token",
                 "maxResults": 3,
                 "fields": _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS,
-                "location": "asia-northeast1",
+                "location": "US",
                 "formatOptions.useInt64Timestamp": True,
             },
-            timeout=None,
         )
-        conn.api_request.assert_has_calls([query_page_1_call, query_page_2_call])
-
-    def test_result_with_max_results(self):
-        from google.cloud.bigquery.table import RowIterator
-
-        query_resource = {
-            "jobComplete": True,
-            "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
-            "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
-            "totalRows": "5",
-        }
-        query_page_resource = {
-            "totalRows": "5",
-            "pageToken": None,
-            "rows": [
-                {"f": [{"v": "abc"}]},
-                {"f": [{"v": "def"}]},
-                {"f": [{"v": "ghi"}]},
-            ],
-        }
-        connection = make_connection(query_resource, query_page_resource)
-        client = _make_client(self.PROJECT, connection=connection)
-        resource = self._make_resource(ended=True)
-        job = self._get_target_class().from_api_repr(resource, client)
-
-        max_results = 3
-
-        result = job.result(max_results=max_results)
-
-        self.assertIsInstance(result, RowIterator)
-        self.assertEqual(result.total_rows, 5)
-
-        rows = list(result)
-
-        self.assertEqual(len(rows), 3)
-        self.assertEqual(len(connection.api_request.call_args_list), 2)
-        query_page_request = connection.api_request.call_args_list[1]
-        self.assertEqual(
-            query_page_request[1]["query_params"]["maxResults"], max_results
+        # Waiting for the results should set maxResults and cache the
+        # first page if max_results is set. This allows customers to
+        # more finely tune when we fallback to the BQ Storage API.
+        # See internal issue: 344008814.
+        conn.api_request.assert_has_calls(
+            [jobs_get_call, query_page_waiting_call, query_page_2_call]
         )
 
     def test_result_w_custom_retry(self):
@@ -1469,63 +1464,85 @@ class TestQueryJob(_Base):
             "jobComplete": True,
             "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
             "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
-            "totalRows": "4",
-        }
-        job_resource = self._make_resource(started=True, ended=True, location="US")
-        q_config = job_resource["configuration"]["query"]
-        q_config["destinationTable"] = {
-            "projectId": self.PROJECT,
-            "datasetId": self.DS_ID,
-            "tableId": self.TABLE_ID,
-        }
-        query_page_resource = {
-            "totalRows": 4,
-            "pageToken": "some-page-token",
+            "totalRows": "10",
             "rows": [
                 {"f": [{"v": "row1"}]},
                 {"f": [{"v": "row2"}]},
                 {"f": [{"v": "row3"}]},
+                {"f": [{"v": "row4"}]},
+                {"f": [{"v": "row5"}]},
+                {"f": [{"v": "row6"}]},
+                {"f": [{"v": "row7"}]},
+                {"f": [{"v": "row8"}]},
+                {"f": [{"v": "row9"}]},
             ],
+            "pageToken": "first-page-token",
         }
-        query_page_resource_2 = {"totalRows": 4, "rows": [{"f": [{"v": "row4"}]}]}
+        job_resource_running = self._make_resource(
+            started=True, ended=False, location="US"
+        )
+        job_resource_done = self._make_resource(started=True, ended=True, location="US")
+        destination_table = {
+            "projectId": self.PROJECT,
+            "datasetId": self.DS_ID,
+            "tableId": self.TABLE_ID,
+        }
+        q_config = job_resource_done["configuration"]["query"]
+        q_config["destinationTable"] = destination_table
+        query_page_resource_2 = {"totalRows": 10, "rows": [{"f": [{"v": "row10"}]}]}
         conn = make_connection(
-            query_results_resource, query_page_resource, query_page_resource_2
+            job_resource_running,
+            query_results_resource,
+            job_resource_done,
+            query_page_resource_2,
         )
         client = _make_client(self.PROJECT, connection=conn)
-        job = self._get_target_class().from_api_repr(job_resource, client)
+        job = self._get_target_class().from_api_repr(job_resource_running, client)
 
         # Act
-        result = job.result(page_size=3)
+        result = job.result(page_size=9)
 
         # Assert
         actual_rows = list(result)
-        self.assertEqual(len(actual_rows), 4)
+        self.assertEqual(len(actual_rows), 10)
 
+        jobs_get_path = f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}"
+        jobs_get_call = mock.call(
+            method="GET",
+            path=jobs_get_path,
+            query_params={"projection": "full", "location": "US"},
+            timeout=DEFAULT_GET_JOB_TIMEOUT,
+        )
         query_results_path = f"/projects/{self.PROJECT}/queries/{self.JOB_ID}"
-        query_page_1_call = mock.call(
+        query_page_waiting_call = mock.call(
             method="GET",
             path=query_results_path,
             query_params={
-                "maxResults": 3,
-                "fields": _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS,
+                # Waiting for the results should set maxResults and cache the
+                # first page if page_size is set. This allows customers to
+                # more finely tune when we fallback to the BQ Storage API.
+                # See internal issue: 344008814.
+                "maxResults": 9,
                 "location": "US",
                 "formatOptions.useInt64Timestamp": True,
             },
             timeout=None,
         )
         query_page_2_call = mock.call(
+            timeout=None,
             method="GET",
             path=query_results_path,
             query_params={
-                "pageToken": "some-page-token",
-                "maxResults": 3,
+                "pageToken": "first-page-token",
+                "maxResults": 9,
                 "fields": _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS,
                 "location": "US",
                 "formatOptions.useInt64Timestamp": True,
             },
-            timeout=None,
         )
-        conn.api_request.assert_has_calls([query_page_1_call, query_page_2_call])
+        conn.api_request.assert_has_calls(
+            [jobs_get_call, query_page_waiting_call, jobs_get_call, query_page_2_call]
+        )
 
     def test_result_with_start_index(self):
         from google.cloud.bigquery.table import RowIterator
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index ed5575f6..cd336b73 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -58,6 +58,7 @@ from google.cloud import bigquery
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery import exceptions
 from google.cloud.bigquery import ParquetOptions
+import google.cloud.bigquery.retry
 from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
 import google.cloud.bigquery.table
 
@@ -5444,6 +5445,127 @@ class TestClient(unittest.TestCase):
         sent = req["data"]
         self.assertEqual(sent["location"], "not-the-client-location")
 
+    def test_query_and_wait_w_max_results(self):
+        query = "select count(*) from `bigquery-public-data.usa_names.usa_1910_2013`"
+        jobs_query_response = {
+            "jobComplete": True,
+        }
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(jobs_query_response)
+
+        _ = client.query_and_wait(query, max_results=11)
+
+        # Verify the request we send is to jobs.query.
+        conn.api_request.assert_called_once()
+        _, req = conn.api_request.call_args
+        self.assertEqual(req["method"], "POST")
+        self.assertEqual(req["path"], f"/projects/{self.PROJECT}/queries")
+        sent = req["data"]
+        self.assertTrue(sent["formatOptions"]["useInt64Timestamp"])
+        self.assertTrue(sent["maxResults"], 11)
+
+    def test_query_and_wait_w_page_size(self):
+        query = "select count(*) from `bigquery-public-data.usa_names.usa_1910_2013`"
+        jobs_query_response = {
+            "jobComplete": True,
+        }
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(jobs_query_response)
+
+        _ = client.query_and_wait(query, page_size=11)
+
+        # Verify the request we send is to jobs.query.
+        conn.api_request.assert_called_once()
+        _, req = conn.api_request.call_args
+        self.assertEqual(req["method"], "POST")
+        self.assertEqual(req["path"], f"/projects/{self.PROJECT}/queries")
+        sent = req["data"]
+        self.assertTrue(sent["formatOptions"]["useInt64Timestamp"])
+        self.assertTrue(sent["maxResults"], 11)
+
+    def test_query_and_wait_w_page_size_multiple_requests(self):
+        """
+        For queries that last longer than the intial (about 10s) call to
+        jobs.query, we should still pass through the page size to the
+        subsequent calls to jobs.getQueryResults.
+
+        See internal issue 344008814.
+        """
+        query = "select count(*) from `bigquery-public-data.usa_names.usa_1910_2013`"
+        job_reference = {
+            "projectId": "my-jobs-project",
+            "location": "my-jobs-location",
+            "jobId": "my-jobs-id",
+        }
+        jobs_query_response = {
+            "jobComplete": False,
+            "jobReference": job_reference,
+        }
+        jobs_get_response = {
+            "jobReference": job_reference,
+            "status": {"state": "DONE"},
+        }
+        get_query_results_response = {
+            "jobComplete": True,
+        }
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(
+            jobs_query_response,
+            jobs_get_response,
+            get_query_results_response,
+        )
+
+        _ = client.query_and_wait(query, page_size=11)
+
+        conn.api_request.assert_has_calls(
+            [
+                # Verify the request we send is to jobs.query.
+                mock.call(
+                    method="POST",
+                    path=f"/projects/{self.PROJECT}/queries",
+                    data={
+                        "useLegacySql": False,
+                        "query": query,
+                        "formatOptions": {"useInt64Timestamp": True},
+                        "maxResults": 11,
+                        "requestId": mock.ANY,
+                    },
+                    timeout=None,
+                ),
+                # jobs.get: Check if the job has finished.
+                mock.call(
+                    method="GET",
+                    path="/projects/my-jobs-project/jobs/my-jobs-id",
+                    query_params={
+                        "projection": "full",
+                        "location": "my-jobs-location",
+                    },
+                    timeout=google.cloud.bigquery.retry.DEFAULT_GET_JOB_TIMEOUT,
+                ),
+                # jobs.getQueryResults: wait for the query / fetch first page
+                mock.call(
+                    method="GET",
+                    path="/projects/my-jobs-project/queries/my-jobs-id",
+                    query_params={
+                        # We should still pass through the page size to the
+                        # subsequent calls to jobs.getQueryResults.
+                        #
+                        # See internal issue 344008814.
+                        "maxResults": 11,
+                        "formatOptions.useInt64Timestamp": True,
+                        "location": "my-jobs-location",
+                    },
+                    timeout=None,
+                ),
+            ]
+        )
+
     def test_query_and_wait_w_project(self):
         query = "select count(*) from `bigquery-public-data.usa_names.usa_1910_2013`"
         jobs_query_response = {
diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py
index 46eb1d6b..298ab9a5 100644
--- a/tests/unit/test_job_retry.py
+++ b/tests/unit/test_job_retry.py
@@ -247,7 +247,7 @@ def test_retry_connection_error_with_default_retries_and_successful_first_job(
                 method="GET",
                 path="/projects/PROJECT/queries/1",
                 query_params={"maxResults": 0, "location": "test-loc"},
-                timeout=120,
+                timeout=None,
             ),
         ],
     )