work-in-progress: callbacks for query finished

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -506,6 +506,7 @@ def query_and_wait(
             retry=retry,
             page_size=page_size,
             max_results=max_results,
+            callback=callback,
         )
 
     path = _to_query_path(project)
@@ -516,7 +517,7 @@ 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():
@@ -576,6 +577,7 @@ def do_query():
                 retry=retry,
                 page_size=page_size,
                 max_results=max_results,
+                callback=callback,
             )
 
         callback(
@@ -657,19 +659,33 @@ def _wait_or_cancel(
     retry: Optional[retries.Retry],
     page_size: Optional[int],
     max_results: Optional[int],
+    callback: Callable,
 ) -> table.RowIterator:
     """Wait for a job to complete and return the results.
 
     If we can't return the results within the ``wait_timeout``, try to cancel
     the job.
     """
     try:
-        return job.result(
+        query_results = job.result(
             page_size=page_size,
             max_results=max_results,
             retry=retry,
             timeout=wait_timeout,
         )
+        callback(
+            QueryCompleteEvent(
+                billing_project=job.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=job.destination,
+            )
+        )
+        return query_results
     except Exception:
         # Attempt to cancel the job since we can't return the results.
         try:
@@ -683,6 +699,7 @@ def _wait_or_cancel(
 @dataclasses.dataclass(frozen=True)
 class QueryCompleteEvent:
     """Query finished successfully."""
+
     billing_project: str
     location: Optional[str]
     query_id: Optional[str]
@@ -691,11 +708,12 @@ class QueryCompleteEvent:
     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]
@@ -712,7 +730,7 @@ class QuerySentEventFactory:
 
     def __init__(self):
         self._event_constructor = QuerySentEvent
-    
+
     def __call__(self, **kwargs):
         result = self._event_constructor(**kwargs)
         self._event_constructor = QueryRetryEvent

tests/unit/test_client_bigframes.py

@@ -22,6 +22,7 @@
 
 import google.auth.credentials
 from google.api_core import exceptions
+from google.cloud import bigquery
 import google.cloud.bigquery.client
 from google.cloud.bigquery import _job_helpers
 
@@ -40,10 +41,15 @@ def make_response(body, *, status_code: int = 200):
 @pytest.fixture
 def client():
     """A real client object with mocked API requests."""
-    credentials = mock.create_autospec(google.auth.credentials.Credentials, instance=True)
+    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,
+        project=PROJECT,
+        credentials=credentials,
+        _http=http_session,
+        location=LOCATION,
     )
 
 
@@ -66,39 +72,89 @@ def test_query_and_wait_bigframes_callback(client):
     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,
-               )
+                _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_callback(client):
+    client._http.request.side_effect = [
+        make_response(
+            {
+                # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/insert
+                # https://cloud.google.com/bigquery/docs/reference/rest/v2/Job
+                "jobReference": {},
+            }
+        ),
+        make_response(
+            {
+                # https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/insert
+                # https://cloud.google.com/bigquery/docs/reference/rest/v2/Job
+                "jobReference": {},
+                "status": {"state": "DONE"},
+            }
+        ),
+    ]
+    callback = mock.Mock()
+    config = bigquery.QueryJobConfig()
+    config.destination = "proj.dset.table"
+    client._query_and_wait_bigframes(
+        query="SELECT 1", job_config=config, 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.QueryCompleteEvent(
+                _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,  
+                    job_id=None,
                     destination=None,
-               ),
+                ),
             ),
         ]
     )
 
 
-def test_query_and_wait_bigframes_with_job_retry_callbacks(client):
+def test_query_and_wait_bigframes_with_query_retry_callbacks(client):
     client._http.request.side_effect = [
         exceptions.InternalServerError(
-            "first try",
-            errors=(
-                {"reason": "jobInternalError"},
-            )
+            "first try", errors=({"reason": "jobInternalError"},)
         ),
         make_response(
             {
@@ -117,37 +173,37 @@ def test_query_and_wait_bigframes_with_job_retry_callbacks(client):
     callback.assert_has_calls(
         [
             mock.call(
-               _job_helpers.QuerySentEvent(
+                _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,  
+                    job_id=None,
                     request_id=mock.ANY,
-               )
+                )
             ),
             mock.call(
-               _job_helpers.QueryRetryEvent(
+                _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,  
+                    job_id=None,
                     request_id=mock.ANY,
-               )
+                )
             ),
             mock.call(
-               _job_helpers.QueryCompleteEvent(
+                _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,  
+                    job_id=None,
                     destination=None,
-               ),
+                ),
             ),
         ]
     )