chore: add private _query_and_wait_bigframes method

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -35,11 +35,14 @@
 predicates where it is safe to generate a new query ID.
 """
 
+from __future__ import annotations
+
 import copy
+import dataclasses
 import functools
 import uuid
 import textwrap
-from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+from typing import Any, Callable, Dict, Optional, TYPE_CHECKING, Union
 import warnings
 
 import google.api_core.exceptions as core_exceptions
@@ -116,10 +119,15 @@ def query_jobs_insert(
     retry: Optional[retries.Retry],
     timeout: Optional[float],
     job_retry: Optional[retries.Retry],
+    callback: Callable,
 ) -> job.QueryJob:
     """Initiate a query using jobs.insert.
 
     See: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/insert
+
+    Args:
+        callback (Callable):
+            A callback function used by bigframes to report query progress.
     """
     job_id_given = job_id is not None
     job_id_save = job_id
@@ -136,6 +144,15 @@ def do_query():
 
         try:
             query_job._begin(retry=retry, timeout=timeout)
+            callback(
+                QueryRequestedEvent(
+                    query=query,
+                    billing_project=query_job.project,
+                    location=query_job.location,
+                    job_id=query_job.job_id,
+                    request_id=None,
+                )
+            )
         except core_exceptions.Conflict as create_exc:
             # The thought is if someone is providing their own job IDs and they get
             # their job ID generation wrong, this could end up returning results for
@@ -396,6 +413,7 @@ def query_and_wait(
     job_retry: Optional[retries.Retry],
     page_size: Optional[int] = None,
     max_results: Optional[int] = None,
+    callback: Callable = lambda _: None,
 ) -> table.RowIterator:
     """Run the query, wait for it to finish, and return the results.
 
@@ -415,9 +433,8 @@ def query_and_wait(
         location (Optional[str]):
             Location where to run the job. Must match the location of the
             table used in the query as well as the destination table.
-        project (Optional[str]):
-            Project ID of the project of where to run the job. Defaults
-            to the client's project.
+        project (str):
+            Project ID of the project of where to run the job.
         api_timeout (Optional[float]):
             The number of seconds to wait for the underlying HTTP transport
             before using ``retry``.
@@ -441,6 +458,8 @@ def query_and_wait(
             request. Non-positive values are ignored.
         max_results (Optional[int]):
             The maximum total number of rows from this request.
+        callback (Callable):
+            A callback function used by bigframes to report query progress.
 
     Returns:
         google.cloud.bigquery.table.RowIterator:
@@ -479,6 +498,7 @@ def query_and_wait(
                 retry=retry,
                 timeout=api_timeout,
                 job_retry=job_retry,
+                callback=callback,
             ),
             api_timeout=api_timeout,
             wait_timeout=wait_timeout,
@@ -497,9 +517,20 @@ def query_and_wait(
         request_body["jobCreationMode"] = client.default_job_creation_mode
 
     def do_query():
-        request_body["requestId"] = make_job_id()
+        request_id = make_job_id()
+        request_body["requestId"] = request_id
         span_attributes = {"path": path}
 
+        callback(
+            QueryRequestedEvent(
+                query=query,
+                billing_project=project,
+                location=location,
+                job_id=None,
+                request_id=request_id,
+            )
+        )
+
         # For easier testing, handle the retries ourselves.
         if retry is not None:
             response = retry(client._call_api)(
@@ -632,3 +663,12 @@ def _wait_or_cancel(
             # Don't eat the original exception if cancel fails.
             pass
         raise
+
+
+@dataclasses.dataclass(frozen=True)
+class QueryRequestedEvent:
+    query: str
+    billing_project: str
+    location: Optional[str]
+    job_id: Optional[str]
+    request_id: Optional[str]

google/cloud/bigquery/client.py

@@ -3633,8 +3633,8 @@ def query_and_wait(
                 rate-limit-exceeded errors.  Passing ``None`` disables
                 job retry. Not all jobs can be retried.
             page_size (Optional[int]):
-                The maximum number of rows in each page of results from this
-                request. Non-positive values are ignored.
+                The maximum number of rows in each page of results from the
+                initial jobs.query request. Non-positive values are ignored.
             max_results (Optional[int]):
                 The maximum total number of rows from this request.
 
@@ -3656,6 +3656,39 @@ def query_and_wait(
                 :class:`~google.cloud.bigquery.job.QueryJobConfig`
                 class.
         """
+        return self._query_and_wait_bigframes(
+            query,
+            job_config=job_config,
+            location=location,
+            project=project,
+            api_timeout=api_timeout,
+            wait_timeout=wait_timeout,
+            retry=retry,
+            job_retry=job_retry,
+            page_size=page_size,
+            max_results=max_results,
+        )
+
+    def _query_and_wait_bigframes(
+        self,
+        query,
+        *,
+        job_config: Optional[QueryJobConfig] = None,
+        location: Optional[str] = None,
+        project: Optional[str] = None,
+        api_timeout: TimeoutType = DEFAULT_TIMEOUT,
+        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,
+        max_results: Optional[int] = None,
+        callback = lambda _: None,
+    ) -> RowIterator:
+        """See query_and_wait.
+        
+        This method has an extra callback parameter, which is used by bigframes
+        to create better progress bars.
+        """
         if project is None:
             project = self.project
 
@@ -3681,6 +3714,7 @@ def query_and_wait(
             job_retry=job_retry,
             page_size=page_size,
             max_results=max_results,
+            callback=callback,
         )
 
     def insert_rows(

tests/unit/test_client_bigframes.py

@@ -0,0 +1,55 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+"""Tests for Client features enabling the bigframes integration."""
+
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+
+import google.auth.credentials
+import google.cloud.bigquery.client
+
+
+PROJECT = "test-project"
+LOCATION = "test-location"
+
+
+def make_response(body, *, status_code: int = 200):
+    response = mock.Mock()
+    type(response).status_code = mock.PropertyMock(return_value=status_code)
+    response.json.return_value = {}
+    return response
+
+
+@pytest.fixture
+def client():
+    """A real client object with mocked API requests."""
+    credentials = mock.create_autospec(google.auth.credentials.Credentials, instance=True)
+    http_session = mock.Mock()
+    return google.cloud.bigquery.client.Client(
+        project=PROJECT, credentials=credentials, _http=http_session, location=LOCATION,
+    )
+
+
+def test_query_and_wait_bigframes_callback(client):
+    client._http.request.side_effect = [
+        make_response(
+            {"jobComplete": True}
+        ),
+    ]
+    callback = lambda _: None
+    client._query_and_wait_bigframes(query="SELECT 1", callback=callback)