fix unit tests

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -39,6 +39,7 @@
 
 import copy
 import dataclasses
+import datetime
 import functools
 import uuid
 import textwrap
@@ -120,7 +121,8 @@ def query_jobs_insert(
     retry: Optional[retries.Retry],
     timeout: Optional[float],
     job_retry: Optional[retries.Retry],
-    callback: Callable,
+    *,
+    callback: Callable = lambda _: None,
 ) -> job.QueryJob:
     """Initiate a query using jobs.insert.
 
@@ -146,15 +148,16 @@ def do_query():
 
         try:
             query_job._begin(retry=retry, timeout=timeout)
-            callback(
-                query_sent_factory(
-                    query=query,
-                    billing_project=query_job.project,
-                    location=query_job.location,
-                    job_id=query_job.job_id,
-                    request_id=None,
+            if job_config is not None and not job_config.dry_run:
+                callback(
+                    query_sent_factory(
+                        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
@@ -526,15 +529,16 @@ def do_query():
         request_body["requestId"] = request_id
         span_attributes = {"path": path}
 
-        callback(
-            query_sent_factory(
-                query=query,
-                billing_project=project,
-                location=location,
-                job_id=None,
-                request_id=request_id,
+        if "dryRun" not in request_body:
+            callback(
+                query_sent_factory(
+                    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:
@@ -581,18 +585,25 @@ def do_query():
                 callback=callback,
             )
 
-        callback(
-            QueryFinishedEvent(
-                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,
+        if "dryRun" not in request_body:
+            callback(
+                QueryFinishedEvent(
+                    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,
+                    # TODO(tswast): After
+                    # https://github.com/googleapis/python-bigquery/pull/2260 goes in, add
+                    # created, started, ended properties here.
+                    created=None,
+                    started=None,
+                    ended=None,
+                )
             )
-        )
         return table.RowIterator(
             client=client,
             api_request=functools.partial(client._call_api, retry, timeout=api_timeout),
@@ -660,42 +671,51 @@ def _wait_or_cancel(
     retry: Optional[retries.Retry],
     page_size: Optional[int],
     max_results: Optional[int],
-    callback: Callable,
+    *,
+    callback: Callable = lambda _: None,
 ) -> 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:
-        callback(
-            QueryReceivedEvent(
-                billing_project=job.project,
-                location=job.location,
-                job_id=job.job_id,
-                statement_type=job.statement_type,
-                state=job.state,
-                query_plan=job.query_plan,
+        if not job.dry_run:
+            callback(
+                QueryReceivedEvent(
+                    billing_project=job.project,
+                    location=job.location,
+                    job_id=job.job_id,
+                    statement_type=job.statement_type,
+                    state=job.state,
+                    query_plan=job.query_plan,
+                    created=job.created,
+                    started=job.started,
+                    ended=job.ended,
+                )
             )
-        )
         query_results = job.result(
             page_size=page_size,
             max_results=max_results,
             retry=retry,
             timeout=wait_timeout,
         )
-        callback(
-            QueryFinishedEvent(
-                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,
+        if not job.dry_run:
+            callback(
+                QueryFinishedEvent(
+                    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,
+                    created=job.created,
+                    started=job.started,
+                    ended=job.ended,
+                )
             )
-        )
         return query_results
     except Exception:
         # Attempt to cancel the job since we can't return the results.
@@ -719,6 +739,9 @@ class QueryFinishedEvent:
     total_rows: Optional[int]
     total_bytes_processed: Optional[int]
     slot_millis: Optional[int]
+    created: Optional[datetime.datetime]
+    started: Optional[datetime.datetime]
+    ended: Optional[datetime.datetime]
 
 
 @dataclasses.dataclass(frozen=True)
@@ -731,6 +754,9 @@ class QueryReceivedEvent:
     statement_type: Optional[str]
     state: Optional[str]
     query_plan: Optional[list[google.cloud.bigquery.job.query.QueryPlanEntry]]
+    created: Optional[datetime.datetime]
+    started: Optional[datetime.datetime]
+    ended: Optional[datetime.datetime]
 
 
 @dataclasses.dataclass(frozen=True)

google/cloud/bigquery/job/query.py

@@ -1550,6 +1550,8 @@ def result(  # type: ignore  # (incompatible with supertype)
             return _EmptyRowIterator(
                 project=self.project,
                 location=self.location,
+                schema=self.schema,
+                total_bytes_processed=self.total_bytes_processed,
                 # Intentionally omit job_id and query_id since this doesn't
                 # actually correspond to a finished query job.
             )
@@ -1737,7 +1739,11 @@ def is_job_done():
                 project=self.project,
                 job_id=self.job_id,
                 query_id=self.query_id,
+                schema=self.schema,
                 num_dml_affected_rows=self._query_results.num_dml_affected_rows,
+                query=self.query,
+                total_bytes_processed=self.total_bytes_processed,
+                slot_millis=self.slot_millis,
             )
 
         # We know that there's at least 1 row, so only treat the response from

google/cloud/bigquery/table.py

@@ -1826,7 +1826,7 @@ def __init__(
             page_start=_rows_page_start,
             next_token="pageToken",
         )
-        schema = _to_schema_fields(schema)
+        schema = _to_schema_fields(schema) if schema else ()
         self._field_to_index = _helpers._field_to_index_mapping(schema)
         self._page_size = page_size
         self._preserve_order = False
@@ -2888,7 +2888,6 @@ class _EmptyRowIterator(RowIterator):
     statements.
     """
 
-    schema = ()
     pages = ()
     total_rows = 0