fix: show progress even in job optional queries (#2119)

* fix: show progress even in job optional queries

* first attempt at publisher

* report execution started/stopped in read_gbq_query

* render bigquery sent events

* Feat render more events (#2121)

* feat: Render more BigQuery events in progress bar

This change updates bigframes/formatting_helpers.py to render more event types from bigframes/core/events.py.

Specifically, it adds rendering support for:
- BigQueryRetryEvent
- BigQueryReceivedEvent
- BigQueryFinishedEvent
- BigQueryUnknownEvent

This provides users with more detailed feedback during query execution in both notebook (HTML) and terminal (plaintext) environments.

* feat: Render more BigQuery events in progress bar

This change updates bigframes/formatting_helpers.py to render more event types from bigframes/core/events.py.

Specifically, it adds rendering support for:
- BigQueryRetryEvent
- BigQueryReceivedEvent
- BigQueryFinishedEvent
- BigQueryUnknownEvent

This provides users with more detailed feedback during query execution in both notebook (HTML) and terminal (plaintext) environments.

Unit tests have been added to verify the rendering of each new event type.

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

* fix job links

* fix system tests

* fix mypy

* fix unit tests

* support more event types

* move publisher to session

* fix remaining mypy errors

* update text

* add explicit unsubscribe

* fix presubmits

* add lock for publisher and publish temp table creations

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: tswast

bigframes/blob/_functions.py

@@ -99,6 +99,7 @@ def _create_udf(self):
             project=None,
             timeout=None,
             query_with_job=True,
+            publisher=self._session._publisher,
         )
 
         return udf_name

bigframes/core/events.py

@@ -0,0 +1,237 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import dataclasses
+import datetime
+import threading
+from typing import Any, Callable, Optional, Set
+import uuid
+
+import google.cloud.bigquery._job_helpers
+import google.cloud.bigquery.job.query
+import google.cloud.bigquery.table
+
+import bigframes.session.executor
+
+
+class Subscriber:
+    def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):
+        self._publisher = publisher
+        self._callback = callback
+        self._subscriber_id = uuid.uuid4()
+
+    def __call__(self, *args, **kwargs):
+        return self._callback(*args, **kwargs)
+
+    def __hash__(self) -> int:
+        return hash(self._subscriber_id)
+
+    def __eq__(self, value: object):
+        if not isinstance(value, Subscriber):
+            return NotImplemented
+        return value._subscriber_id == self._subscriber_id
+
+    def close(self):
+        self._publisher.unsubscribe(self)
+        del self._publisher
+        del self._callback
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_value is not None:
+            self(
+                UnknownErrorEvent(
+                    exc_type=exc_type,
+                    exc_value=exc_value,
+                    traceback=traceback,
+                )
+            )
+        self.close()
+
+
+class Publisher:
+    def __init__(self):
+        self._subscribers_lock = threading.Lock()
+        self._subscribers: Set[Subscriber] = set()
+
+    def subscribe(self, callback: Callable[[Event], None]) -> Subscriber:
+        # TODO(b/448176657): figure out how to handle subscribers/publishers in
+        # a background thread. Maybe subscribers should be thread-local?
+        subscriber = Subscriber(callback, publisher=self)
+        with self._subscribers_lock:
+            self._subscribers.add(subscriber)
+        return subscriber
+
+    def unsubscribe(self, subscriber: Subscriber):
+        with self._subscribers_lock:
+            self._subscribers.remove(subscriber)
+
+    def publish(self, event: Event):
+        with self._subscribers_lock:
+            for subscriber in self._subscribers:
+                subscriber(event)
+
+
+class Event:
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class SessionClosed(Event):
+    session_id: str
+
+
+class ExecutionStarted(Event):
+    pass
+
+
+class ExecutionRunning(Event):
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class ExecutionFinished(Event):
+    result: Optional[bigframes.session.executor.ExecuteResult] = None
+
+
+@dataclasses.dataclass(frozen=True)
+class UnknownErrorEvent(Event):
+    exc_type: Any
+    exc_value: Any
+    traceback: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQuerySentEvent(ExecutionRunning):
+    """Query sent to BigQuery."""
+
+    query: str
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    job_id: Optional[str] = None
+    request_id: Optional[str] = None
+
+    @classmethod
+    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QuerySentEvent):
+        return cls(
+            query=event.query,
+            billing_project=event.billing_project,
+            location=event.location,
+            job_id=event.job_id,
+            request_id=event.request_id,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryRetryEvent(ExecutionRunning):
+    """Query sent another time because the previous attempt failed."""
+
+    query: str
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    job_id: Optional[str] = None
+    request_id: Optional[str] = None
+
+    @classmethod
+    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QueryRetryEvent):
+        return cls(
+            query=event.query,
+            billing_project=event.billing_project,
+            location=event.location,
+            job_id=event.job_id,
+            request_id=event.request_id,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryReceivedEvent(ExecutionRunning):
+    """Query received and acknowledged by the BigQuery API."""
+
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    job_id: Optional[str] = None
+    statement_type: Optional[str] = None
+    state: Optional[str] = None
+    query_plan: Optional[list[google.cloud.bigquery.job.query.QueryPlanEntry]] = None
+    created: Optional[datetime.datetime] = None
+    started: Optional[datetime.datetime] = None
+    ended: Optional[datetime.datetime] = None
+
+    @classmethod
+    def from_bqclient(
+        cls, event: google.cloud.bigquery._job_helpers.QueryReceivedEvent
+    ):
+        return cls(
+            billing_project=event.billing_project,
+            location=event.location,
+            job_id=event.job_id,
+            statement_type=event.statement_type,
+            state=event.state,
+            query_plan=event.query_plan,
+            created=event.created,
+            started=event.started,
+            ended=event.ended,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryFinishedEvent(ExecutionRunning):
+    """Query finished successfully."""
+
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    query_id: Optional[str] = None
+    job_id: Optional[str] = None
+    destination: Optional[google.cloud.bigquery.table.TableReference] = None
+    total_rows: Optional[int] = None
+    total_bytes_processed: Optional[int] = None
+    slot_millis: Optional[int] = None
+    created: Optional[datetime.datetime] = None
+    started: Optional[datetime.datetime] = None
+    ended: Optional[datetime.datetime] = None
+
+    @classmethod
+    def from_bqclient(
+        cls, event: google.cloud.bigquery._job_helpers.QueryFinishedEvent
+    ):
+        return cls(
+            billing_project=event.billing_project,
+            location=event.location,
+            query_id=event.query_id,
+            job_id=event.job_id,
+            destination=event.destination,
+            total_rows=event.total_rows,
+            total_bytes_processed=event.total_bytes_processed,
+            slot_millis=event.slot_millis,
+            created=event.created,
+            started=event.started,
+            ended=event.ended,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryUnknownEvent(ExecutionRunning):
+    """Got unknown event from the BigQuery client library."""
+
+    # TODO: should we just skip sending unknown events?
+
+    event: object
+
+    @classmethod
+    def from_bqclient(cls, event):
+        return cls(event)

bigframes/dataframe.py

@@ -4670,24 +4670,24 @@ def to_string(
     ) -> str | None:
         return self.to_pandas(allow_large_results=allow_large_results).to_string(
             buf,
-            columns,  # type: ignore
-            col_space,
-            header,  # type: ignore
-            index,
-            na_rep,
-            formatters,
-            float_format,
-            sparsify,
-            index_names,
-            justify,
-            max_rows,
-            max_cols,
-            show_dimensions,
-            decimal,
-            line_width,
-            min_rows,
-            max_colwidth,
-            encoding,
+            columns=columns,  # type: ignore
+            col_space=col_space,
+            header=header,  # type: ignore
+            index=index,
+            na_rep=na_rep,
+            formatters=formatters,
+            float_format=float_format,
+            sparsify=sparsify,
+            index_names=index_names,
+            justify=justify,
+            max_rows=max_rows,
+            max_cols=max_cols,
+            show_dimensions=show_dimensions,
+            decimal=decimal,
+            line_width=line_width,
+            min_rows=min_rows,
+            max_colwidth=max_colwidth,
+            encoding=encoding,
         )
 
     def to_html(
@@ -4720,28 +4720,28 @@ def to_html(
     ) -> str:
         return self.to_pandas(allow_large_results=allow_large_results).to_html(
             buf,
-            columns,  # type: ignore
-            col_space,
-            header,
-            index,
-            na_rep,
-            formatters,
-            float_format,
-            sparsify,
-            index_names,
-            justify,  # type: ignore
-            max_rows,
-            max_cols,
-            show_dimensions,
-            decimal,
-            bold_rows,
-            classes,
-            escape,
-            notebook,
-            border,
-            table_id,
-            render_links,
-            encoding,
+            columns=columns,  # type: ignore
+            col_space=col_space,
+            header=header,
+            index=index,
+            na_rep=na_rep,
+            formatters=formatters,
+            float_format=float_format,
+            sparsify=sparsify,
+            index_names=index_names,
+            justify=justify,  # type: ignore
+            max_rows=max_rows,
+            max_cols=max_cols,
+            show_dimensions=show_dimensions,
+            decimal=decimal,
+            bold_rows=bold_rows,
+            classes=classes,
+            escape=escape,
+            notebook=notebook,
+            border=border,
+            table_id=table_id,
+            render_links=render_links,
+            encoding=encoding,
         )
 
     def to_markdown(
@@ -4753,7 +4753,7 @@ def to_markdown(
         allow_large_results: Optional[bool] = None,
         **kwargs,
     ) -> str | None:
-        return self.to_pandas(allow_large_results=allow_large_results).to_markdown(buf, mode, index, **kwargs)  # type: ignore
+        return self.to_pandas(allow_large_results=allow_large_results).to_markdown(buf, mode=mode, index=index, **kwargs)  # type: ignore
 
     def to_pickle(self, path, *, allow_large_results=None, **kwargs) -> None:
         return self.to_pandas(allow_large_results=allow_large_results).to_pickle(