feat: add Client.query_and_wait which directly returns a RowIterator of results

Set the `QUERY_PREVIEW_ENABLED=TRUE` environment variable to use this with the
new JOB_CREATION_OPTIONAL mode (currently in preview).

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -12,9 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Helpers for interacting with the job REST APIs from the client."""
+"""Helpers for interacting with the job REST APIs from the client.
+
+For queries, there are three cases to consider:
+
+1. jobs.insert: This always returns a job resource.
+2. jobs.query, jobCreationMode=JOB_CREATION_REQUIRED:
+   This sometimes can return the results inline, but always includes a job ID.
+3. jobs.query, jobCreationMode=JOB_CREATION_OPTIONAL:
+   This sometimes doesn't create a job at all, instead returning the results.
+   For better debugging, a query ID is included in the response (not always a
+   job ID).
+
+Client.query() calls either (1) or (2), depending on what the user provides
+for the api_method parameter. query() always returns a QueryJob object, which
+can retry the query when the query job fails for a retriable reason.
+
+Client.query_and_wait() calls (3). This returns a RowIterator that may wrap
+local results from the response or may wrap a query job containing multiple
+pages of results. Even though query_and_wait() waits for the job to complete,
+we still need a separate job_retry object because there are different
+predicates where it is safe to generate a new query ID.
+"""
 
 import copy
+import os
 import uuid
 from typing import Any, Dict, TYPE_CHECKING, Optional
 
@@ -23,6 +45,7 @@
 
 from google.cloud.bigquery import job
 import google.cloud.bigquery.query
+from google.cloud.bigquery import table
 
 # Avoid circular imports
 if TYPE_CHECKING:  # pragma: NO COVER
@@ -123,7 +146,12 @@ def do_query():
     return future
 
 
-def _to_query_request(job_config: Optional[job.QueryJobConfig]) -> Dict[str, Any]:
+def _to_query_request(
+    query: str,
+    job_config: Optional[job.QueryJobConfig],
+    location: Optional[str],
+    timeout: Optional[float],
+) -> Dict[str, Any]:
     """Transform from Job resource to QueryRequest resource.
 
     Most of the keys in job.configuration.query are in common with
@@ -150,6 +178,12 @@ def _to_query_request(job_config: Optional[job.QueryJobConfig]) -> Dict[str, Any
     request_body.setdefault("formatOptions", {})
     request_body["formatOptions"]["useInt64Timestamp"] = True  # type: ignore
 
+    if timeout is not None:
+        # Subtract a buffer for context switching, network latency, etc.
+        request_body["timeoutMs"] = max(0, int(1000 * timeout) - _TIMEOUT_BUFFER_MILLIS)
+    request_body["location"] = location
+    request_body["query"] = query
+
     return request_body
 
 
@@ -207,6 +241,10 @@ def _to_query_job(
     return query_job
 
 
+def _to_query_path(project: str) -> str:
+    return f"/projects/{project}/queries"
+
+
 def query_jobs_query(
     client: "Client",
     query: str,
@@ -217,18 +255,12 @@ def query_jobs_query(
     timeout: Optional[float],
     job_retry: retries.Retry,
 ) -> job.QueryJob:
-    """Initiate a query using jobs.query.
+    """Initiate a query using jobs.query with jobCreationMode=JOB_CREATION_REQUIRED.
 
     See: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
     """
-    path = f"/projects/{project}/queries"
-    request_body = _to_query_request(job_config)
-
-    if timeout is not None:
-        # Subtract a buffer for context switching, network latency, etc.
-        request_body["timeoutMs"] = max(0, int(1000 * timeout) - _TIMEOUT_BUFFER_MILLIS)
-    request_body["location"] = location
-    request_body["query"] = query
+    path = _to_query_path(project)
+    request_body = _to_query_request(query, job_config, location, timeout)
 
     def do_query():
         request_body["requestId"] = make_job_id()
@@ -253,3 +285,84 @@ def do_query():
     future._job_retry = job_retry
 
     return future
+
+
+def query_and_wait(
+    client: "Client",
+    query: str,
+    job_config: Optional[job.QueryJobConfig],
+    location: Optional[str],
+    project: str,
+    retry: retries.Retry,
+    timeout: Optional[float],
+    job_retry: retries.Retry,
+) -> table.RowIterator:
+    """Initiate a query using jobs.query and waits for results.
+     
+    While ``jobCreationMode=JOB_CREATION_OPTIONAL`` is in preview, use the
+    default ``jobCreationMode`` unless the environment variable
+    ``QUERY_PREVIEW_ENABLED=true``. After ``jobCreationMode`` is GA, this
+    method will always use ``jobCreationMode=JOB_CREATION_OPTIONAL``.
+
+    See: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
+    """
+    path = _to_query_path(project)
+    request_body = _to_query_request(query, job_config, location, timeout)
+
+    if os.getenv("QUERY_PREVIEW_ENABLED").casefold() == "true":
+        request_body["jobCreationMode"] = "JOB_CREATION_OPTIONAL"
+
+    @job_retry
+    def do_query():
+        request_body["requestId"] = make_job_id()
+        span_attributes = {"path": path}
+
+        response = client._call_api(
+            retry,
+            span_name="BigQuery.query",
+            span_attributes=span_attributes,
+            method="POST",
+            path=path,
+            data=request_body,
+            timeout=timeout,
+        )
+
+        # The query hasn't finished, so we expect there to be a job ID now.
+        # Wait until the query finishes.
+        if not response.get("jobComplete", False):
+            return _to_query_job(client, query, job_config, response).result()
+
+        # Even if we run with JOB_CREATION_OPTIONAL, if there are more pages
+        # to fetch, there will be a job ID for jobs.getQueryResults.
+        query_results = google.cloud.bigquery.query._QueryResults.from_api_repr(response)
+        job_id = query_results.job_id
+        location = query_results.location
+        rows = query_results.rows
+        total_rows = query_results.total_rows
+        more_pages = (
+            job_id is not None
+            and location is not None
+            and len(rows) < total_rows
+        )
+        
+        if more_pages:
+            # TODO(swast): Call client._list_rows_from_query_results directly
+            # after updating RowIterator to fetch destination only if needed.
+            return _to_query_job(client, query, job_config, response).result()
+        
+        return table.RowIterator(
+            client=client,
+            api_request=client._call_api,
+            path=None,
+            schema=query_results.schema,
+            # TODO(swast): Support max_results
+            max_results=None,
+            total_rows=total_rows,
+            first_page_response=response,
+            location=location,
+            job_id=job_id,
+            query_id=query_results.query_id,
+            project=project,            
+        )
+
+    return do_query()

google/cloud/bigquery/client.py

@@ -3405,6 +3405,20 @@ def query(
         else:
             raise ValueError(f"Got unexpected value for api_method: {repr(api_method)}")
 
+    def query_and_wait(self) -> RowIterator:
+        """[Preview] Run the query and return the results."""
+        maybe_job = _job_helpers.query_jobs_query(
+            self,
+            query,
+            job_config,
+            location,
+            project,
+            retry,
+            timeout,
+            job_retry,
+            # TODO: add no job mode
+        )
+
     def insert_rows(
         self,
         table: Union[Table, TableReference, str],
@@ -3853,7 +3867,7 @@ def _list_rows_from_query_results(
         job_id: str,
         location: str,
         project: str,
-        schema: SchemaField,
+        schema: Sequence[SchemaField],
         total_rows: Optional[int] = None,
         destination: Optional[Union[Table, TableReference, TableListItem, str]] = None,
         max_results: Optional[int] = None,

google/cloud/bigquery/query.py

@@ -911,6 +911,18 @@ def job_id(self):
         """
         return self._properties.get("jobReference", {}).get("jobId")
 
+    @property
+    def location(self):
+        """Location of the query job these results are from.
+
+        See:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#body.QueryResponse.FIELDS.job_reference
+
+        Returns:
+            str: Job ID of the query job.
+        """
+        return self._properties.get("jobReference", {}).get("location")
+
     @property
     def query_id(self) -> Optional[str]:
         """[Preview] ID of a completed query.

google/cloud/bigquery/table.py

@@ -1591,7 +1591,36 @@ def __init__(
         self._location = location
         self._job_id = job_id
         self._query_id = query_id
-        self._project = project
+
+        if project:
+            self._project = project
+        elif client is not None:
+            self._project = client.project
+        else:
+            self._project = None
+
+    def project(self) -> Optional[str]:
+        """GCP Project ID where these rows are read from."""
+        return self._project
+
+    def location(self) -> Optional[str]:
+        """Location where the query executed (if applicable).
+
+        See: https://cloud.google.com/bigquery/docs/locations
+        """
+        self._location
+
+    def job_id(self) -> Optional[str]:
+        """ID of the query job (if applicable).
+
+        To get the job metadata, call
+        ``job = client.get_job(rows.job_id, location=rows.location)``.
+        """
+        return self._job_id
+
+    def query_id(self) -> Optional[str]:
+        """ID of the stateless query (if applicable)."""
+        return self._query_id
 
     @property
     def _billing_project(self) -> Optional[str]: