add QueryCompleteEvent

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -132,6 +132,7 @@ def query_jobs_insert(
     job_id_given = job_id is not None
     job_id_save = job_id
     job_config_save = job_config
+    query_sent_factory = QuerySentEventFactory()
 
     def do_query():
         # Make a copy now, so that original doesn't get changed by the process
@@ -145,7 +146,7 @@ def do_query():
         try:
             query_job._begin(retry=retry, timeout=timeout)
             callback(
-                QuerySentEvent(
+                query_sent_factory(
                     query=query,
                     billing_project=query_job.project,
                     location=query_job.location,
@@ -515,14 +516,16 @@ def query_and_wait(
         request_body["maxResults"] = page_size or max_results
     if client.default_job_creation_mode:
         request_body["jobCreationMode"] = client.default_job_creation_mode
+    
+    query_sent_factory = QuerySentEventFactory()
 
     def do_query():
         request_id = make_job_id()
         request_body["requestId"] = request_id
         span_attributes = {"path": path}
 
         callback(
-            QuerySentEvent(
+            query_sent_factory(
                 query=query,
                 billing_project=project,
                 location=location,
@@ -575,6 +578,18 @@ def do_query():
                 max_results=max_results,
             )
 
+        callback(
+            QueryCompleteEvent(
+                billing_project=project,
+                location=query_results.location,
+                query_id=query_results.query_id,
+                job_id=query_results.job_id,
+                total_rows=query_results.total_rows,
+                total_bytes_processed=query_results.total_bytes_processed,
+                slot_millis=query_results.slot_millis,
+                destination=None,
+            )
+        )
         return table.RowIterator(
             client=client,
             api_request=functools.partial(client._call_api, retry, timeout=api_timeout),
@@ -665,10 +680,40 @@ def _wait_or_cancel(
         raise
 
 
+@dataclasses.dataclass(frozen=True)
+class QueryCompleteEvent:
+    """Query finished successfully."""
+    billing_project: str
+    location: Optional[str]
+    query_id: Optional[str]
+    job_id: Optional[str]
+    destination: Optional[table.TableReference]
+    total_rows: Optional[int]
+    total_bytes_processed: Optional[int]
+    slot_millis: Optional[int]
+    
+
 @dataclasses.dataclass(frozen=True)
 class QuerySentEvent:
+    """Query sent to BigQuery."""
     query: str
     billing_project: str
     location: Optional[str]
     job_id: Optional[str]
     request_id: Optional[str]
+
+
+class QueryRetryEvent(QuerySentEvent):
+    """Query sent another time because the previous failed."""
+
+
+class QuerySentEventFactory:
+    """Creates a QuerySentEvent first, then QueryRetryEvent after that."""
+
+    def __init__(self):
+        self._event_constructor = QuerySentEvent
+    
+    def __call__(self, **kwargs):
+        result = self._event_constructor(**kwargs)
+        self._event_constructor = QueryRetryEvent
+        return result

google/cloud/bigquery/query.py

@@ -1228,11 +1228,18 @@ def location(self):
 
         See:
         https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#body.QueryResponse.FIELDS.job_reference
+        or https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#body.QueryResponse.FIELDS.location
 
         Returns:
             str: Job ID of the query job.
         """
-        return self._properties.get("jobReference", {}).get("location")
+        location = self._properties.get("jobReference", {}).get("location")
+
+        # Sometimes there's no job, but we still want to get the location
+        # information. Prefer the value from job for backwards compatibilitity.
+        if not location:
+            location = self._properties.get("location")
+        return location
 
     @property
     def query_id(self) -> Optional[str]:

tests/unit/test_client_bigframes.py

@@ -21,6 +21,7 @@
 import pytest
 
 import google.auth.credentials
+from google.api_core import exceptions
 import google.cloud.bigquery.client
 from google.cloud.bigquery import _job_helpers
 
@@ -32,7 +33,7 @@
 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 = {}
+    response.json.return_value = body
     return response
 
 
@@ -49,15 +50,104 @@ def client():
 def test_query_and_wait_bigframes_callback(client):
     client._http.request.side_effect = [
         make_response(
-            {"jobComplete": True}
+            {
+                # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
+                "location": LOCATION,
+                "queryId": "abcdefg",
+                "totalRows": "100",
+                "totalBytesProcessed": "123",
+                "totalSlotMs": "987",
+                "jobComplete": True,
+            }
         ),
     ]
     callback = mock.Mock()
     client._query_and_wait_bigframes(query="SELECT 1", callback=callback)
     callback.assert_has_calls(
-        mock.call(
-           _job_helpers.QuerySentEvent(
-               query="SELECT 1",
+        [
+            mock.call(
+               _job_helpers.QuerySentEvent(
+                   query="SELECT 1",
+                   billing_project=PROJECT,
+                   location=LOCATION,
+                   # No job ID, because a basic query is elegible for jobs.query.
+                   job_id=None,  
+                   request_id=mock.ANY,
+               )
             ),
-        )
+            mock.call(
+               _job_helpers.QueryCompleteEvent(
+                    billing_project=PROJECT,
+                    location=LOCATION,
+                    query_id="abcdefg",
+                    total_rows=100,
+                    total_bytes_processed=123,
+                    slot_millis=987,
+                    # No job ID or destination, because a basic query is elegible for jobs.query.
+                    job_id=None,  
+                    destination=None,
+               ),
+            ),
+        ]
+    )
+
+
+def test_query_and_wait_bigframes_with_job_retry_callbacks(client):
+    client._http.request.side_effect = [
+        exceptions.InternalServerError(
+            "first try",
+            errors=(
+                {"reason": "jobInternalError"},
+            )
+        ),
+        make_response(
+            {
+                # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
+                "location": LOCATION,
+                "queryId": "abcdefg",
+                "totalRows": "100",
+                "totalBytesProcessed": "123",
+                "totalSlotMs": "987",
+                "jobComplete": True,
+            }
+        ),
+    ]
+    callback = mock.Mock()
+    client._query_and_wait_bigframes(query="SELECT 1", callback=callback)
+    callback.assert_has_calls(
+        [
+            mock.call(
+               _job_helpers.QuerySentEvent(
+                    query="SELECT 1",
+                    billing_project=PROJECT,
+                    location=LOCATION,
+                    # No job ID, because a basic query is elegible for jobs.query.
+                    job_id=None,  
+                    request_id=mock.ANY,
+               )
+            ),
+            mock.call(
+               _job_helpers.QueryRetryEvent(
+                    query="SELECT 1",
+                    billing_project=PROJECT,
+                    location=LOCATION,
+                    # No job ID, because a basic query is elegible for jobs.query.
+                    job_id=None,  
+                    request_id=mock.ANY,
+               )
+            ),
+            mock.call(
+               _job_helpers.QueryCompleteEvent(
+                    billing_project=PROJECT,
+                    location=LOCATION,
+                    query_id=mock.ANY,
+                    total_rows=100,
+                    total_bytes_processed=123,
+                    slot_millis=987,
+                    # No job ID or destination, because a basic query is elegible for jobs.query.
+                    job_id=None,  
+                    destination=None,
+               ),
+            ),
+        ]
     )