first attempt at publisher

Author: tswast

bigframes/core/events.py

@@ -24,6 +24,8 @@
 import google.cloud.bigquery.job.query
 import google.cloud.bigquery.table
 
+import bigframes.formatting_helpers
+
 
 @dataclasses.dataclass(frozen=True)
 class Subscriber:
@@ -66,6 +68,7 @@ def send(self, event: Event):
 
 
 publisher = Publisher()
+publisher.subscribe(bigframes.formatting_helpers.progress_callback)
 
 
 class Event:
@@ -85,7 +88,7 @@ class ExecutionStopped(Event):
 
 
 @dataclasses.dataclass(frozen=True)
-class BigQuerySentEvent(ExecutionStarted):
+class BigQuerySentEvent(ExecutionRunning):
     """Query sent to BigQuery."""
 
     query: str
@@ -158,7 +161,7 @@ def from_bqclient(
 
 
 @dataclasses.dataclass(frozen=True)
-class BigQueryFinishedEvent(ExecutionStopped):
+class BigQueryFinishedEvent(ExecutionRunning):
     """Query finished successfully."""
 
     billing_project: Optional[str] = None

bigframes/formatting_helpers.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 """Shared helper functions for formatting jobs related info."""
-# TODO(orrbradford): cleanup up typings and documenttion in this file
+
+from __future__ import annotations
 
 import datetime
 import random
-from typing import Any, Optional, Type, Union
+from typing import Any, Optional, Type, TYPE_CHECKING, Union
 
 import bigframes_vendored.constants as constants
 import google.api_core.exceptions as api_core_exceptions
@@ -27,6 +28,9 @@
 import IPython.display as display
 import ipywidgets as widgets
 
+if TYPE_CHECKING:
+    import bigframes.core.events
+
 GenericJob = Union[
     bigquery.LoadJob, bigquery.ExtractJob, bigquery.QueryJob, bigquery.CopyJob
 ]
@@ -119,71 +123,51 @@ def repr_query_job(query_job: Optional[bigquery.QueryJob]):
     return res
 
 
-def wait_for_query_job(
-    query_job: bigquery.QueryJob,
-    max_results: Optional[int] = None,
-    page_size: Optional[int] = None,
-    progress_bar: Optional[str] = None,
-) -> bigquery.table.RowIterator:
-    """Return query results. Displays a progress bar while the query is running
-    Args:
-        query_job (bigquery.QueryJob, Optional):
-            The job representing the execution of the query on the server.
-        max_results (int, Optional):
-            The maximum number of rows the row iterator should return.
-        page_size (int, Optional):
-            The number of results to return on each results page.
-        progress_bar (str, Optional):
-            Which progress bar to show.
-    Returns:
-        A row iterator over the query results.
-    """
+current_display: Optional[display.HTML] = None
+current_display_id: Optional[str] = None
+
+
+def progress_callback(
+    event: bigframes.core.events.Event,
+):
+    """Displays a progress bar while the query is running"""
+    global current_display, current_display_id
+
+    import bigframes._config
+    import bigframes.core.events
+
+    progress_bar = bigframes._config.options.display.progress_bar
+
     if progress_bar == "auto":
         progress_bar = "notebook" if in_ipython() else "terminal"
 
-    try:
-        if progress_bar == "notebook":
-            display_id = str(random.random())
-            loading_bar = display.HTML(get_query_job_loading_html(query_job))
-            display.display(loading_bar, display_id=display_id)
-            query_result = query_job.result(
-                max_results=max_results, page_size=page_size
-            )
-            query_job.reload()
+    if progress_bar == "notebook":
+        if (
+            isinstance(event, bigframes.core.events.ExecutionStarted)
+            or current_display is None
+            or current_display_id is None
+        ):
+            current_display_id = str(random.random())
+            current_display = display.HTML("Starting execution.")
+            display.display(current_display)
+
+        if isinstance(event, bigframes.core.events.ExecutionRunning):
             display.update_display(
-                display.HTML(get_query_job_loading_html(query_job)),
-                display_id=display_id,
+                display.HTML("Execution happening."),
+                display_id=current_display_id,
             )
-        elif progress_bar == "terminal":
-            initial_loading_bar = get_query_job_loading_string(query_job)
-            print(initial_loading_bar)
-            query_result = query_job.result(
-                max_results=max_results, page_size=page_size
-            )
-            query_job.reload()
-            if initial_loading_bar != get_query_job_loading_string(query_job):
-                print(get_query_job_loading_string(query_job))
-        else:
-            # No progress bar.
-            query_result = query_job.result(
-                max_results=max_results, page_size=page_size
+        elif isinstance(event, bigframes.core.events.ExecutionStopped):
+            display.update_display(
+                display.HTML("Execution done."),
+                display_id=current_display_id,
             )
-            query_job.reload()
-        return query_result
-    except api_core_exceptions.RetryError as exc:
-        add_feedback_link(exc)
-        raise
-    except api_core_exceptions.GoogleAPICallError as exc:
-        add_feedback_link(exc)
-        raise
-    except KeyboardInterrupt:
-        query_job.cancel()
-        print(
-            f"Requested cancellation for {query_job.job_type.capitalize()}"
-            f" job {query_job.job_id} in location {query_job.location}..."
-        )
-        # begin the cancel request before immediately rethrowing
-        raise
+    elif progress_bar == "terminal":
+        if isinstance(event, bigframes.core.events.ExecutionStarted):
+            print("Starting execution.")
+        elif isinstance(event, bigframes.core.events.ExecutionRunning):
+            print("Execution happening.")
+        elif isinstance(event, bigframes.core.events.ExecutionStopped):
+            print("Execution done.")
 
 
 def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None):

bigframes/session/_io/bigquery/__init__.py

@@ -29,11 +29,13 @@
 import google.api_core.exceptions
 import google.api_core.retry
 import google.cloud.bigquery as bigquery
+import google.cloud.bigquery._job_helpers
+import google.cloud.bigquery.table
 
 from bigframes.core import log_adapter
 import bigframes.core.compile.googlesql as googlesql
+import bigframes.core.events
 import bigframes.core.sql
-import bigframes.formatting_helpers as formatting_helpers
 import bigframes.session.metrics
 
 CHECK_DRIVE_PERMISSIONS = "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions."
@@ -238,6 +240,15 @@ def add_and_trim_labels(job_config):
     )
 
 
+def publish_bq_event(event):
+    if isinstance(event, google.cloud.bigquery._job_helpers.QuerySentEvent):
+        bf_event = bigframes.core.events.BigQuerySentEvent.from_bqclient(event)
+    else:
+        bf_event = bigframes.core.events.BigQueryUnknownEvent(event)
+
+    bigframes.core.events.publisher.send(bf_event)
+
+
 @overload
 def start_query_with_client(
     bq_client: bigquery.Client,
@@ -249,7 +260,7 @@ def start_query_with_client(
     timeout: Optional[float],
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[True],
-) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
+) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]:
     ...
 
 
@@ -264,7 +275,7 @@ def start_query_with_client(
     timeout: Optional[float],
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[False],
-) -> Tuple[bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
+) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     ...
 
 
@@ -280,7 +291,7 @@ def start_query_with_client(
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[True],
     job_retry: google.api_core.retry.Retry,
-) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
+) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]:
     ...
 
 
@@ -296,7 +307,7 @@ def start_query_with_client(
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[False],
     job_retry: google.api_core.retry.Retry,
-) -> Tuple[bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
+) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     ...
 
 
@@ -315,23 +326,25 @@ def start_query_with_client(
     # https://github.com/googleapis/python-bigquery/pull/2256 merged, likely
     # version 3.36.0 or later.
     job_retry: google.api_core.retry.Retry = third_party_gcb_retry.DEFAULT_JOB_RETRY,
-) -> Tuple[bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
+) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     """
     Starts query job and waits for results.
     """
+    # Note: Ensure no additional labels are added to job_config after this
+    # point, as `add_and_trim_labels` ensures the label count does not
+    # exceed MAX_LABELS_COUNT.
+    add_and_trim_labels(job_config)
+
     try:
-        # Note: Ensure no additional labels are added to job_config after this
-        # point, as `add_and_trim_labels` ensures the label count does not
-        # exceed MAX_LABELS_COUNT.
-        add_and_trim_labels(job_config)
         if not query_with_job:
-            results_iterator = bq_client.query_and_wait(
+            results_iterator = bq_client._query_and_wait_bigframes(
                 sql,
                 job_config=job_config,
                 location=location,
                 project=project,
                 api_timeout=timeout,
                 job_retry=job_retry,
+                callback=publish_bq_event,
             )
             if metrics is not None:
                 metrics.count_job_stats(row_iterator=results_iterator)
@@ -350,14 +363,32 @@ def start_query_with_client(
             ex.message += CHECK_DRIVE_PERMISSIONS
         raise
 
-    opts = bigframes.options.display
-    if opts.progress_bar is not None and not query_job.configuration.dry_run:
-        results_iterator = formatting_helpers.wait_for_query_job(
-            query_job,
-            progress_bar=opts.progress_bar,
+    if not query_job.configuration.dry_run:
+        bigframes.core.events.publisher.send(
+            bigframes.core.events.BigQuerySentEvent(
+                sql,
+                billing_project=query_job.project,
+                location=query_job.location,
+                job_id=query_job.job_id,
+                request_id=None,
+            )
+        )
+    results_iterator = query_job.result()
+    if not query_job.configuration.dry_run:
+        bigframes.core.events.publisher.send(
+            bigframes.core.events.BigQueryFinishedEvent(
+                billing_project=query_job.project,
+                location=query_job.location,
+                job_id=query_job.job_id,
+                destination=query_job.destination,
+                total_rows=results_iterator.total_rows,
+                total_bytes_processed=query_job.total_bytes_processed,
+                slot_millis=query_job.slot_millis,
+                created=query_job.created,
+                started=query_job.started,
+                ended=query_job.ended,
+            )
         )
-    else:
-        results_iterator = query_job.result()
 
     if metrics is not None:
         metrics.count_job_stats(query_job=query_job)