fix: Respect display.progress_bar=None in background threads

packages/bigframes/bigframes/core/events.py

@@ -18,17 +18,19 @@
 import datetime
 import threading
 import uuid
-from typing import Any, Callable, Optional, Set
+from typing import Any, Callable, Literal, Set
 
 import google.cloud.bigquery._job_helpers
 import google.cloud.bigquery.job.query
 import google.cloud.bigquery.table
 
 import bigframes.session.executor
 
+_DEFAULT: Literal["default"] = "default"
+
 
 class Subscriber:
-    def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):
+    def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):  # noqa: E501
         self._publisher = publisher
         self._callback = callback
         self._subscriber_id = uuid.uuid4()
@@ -106,7 +108,7 @@ class ExecutionRunning(Event):
 
 @dataclasses.dataclass(frozen=True)
 class ExecutionFinished(Event):
-    result: Optional[bigframes.session.executor.ExecuteResult] = None
+    result: bigframes.session.executor.ExecuteResult | None = None
 
 
 @dataclasses.dataclass(frozen=True)
@@ -121,19 +123,26 @@ 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
+    billing_project: str | None = None
+    location: str | None = None
+    job_id: str | None = None
+    request_id: str | None = None
+    progress_bar: Literal["default", "auto", "notebook", "terminal"] | None = _DEFAULT
 
     @classmethod
-    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QuerySentEvent):
+    def from_bqclient(
+        cls,
+        event: google.cloud.bigquery._job_helpers.QuerySentEvent,
+        progress_bar: Literal["default", "auto", "notebook", "terminal"]
+        | None = _DEFAULT,
+    ):
         return cls(
             query=event.query,
             billing_project=event.billing_project,
             location=event.location,
             job_id=event.job_id,
             request_id=event.request_id,
+            progress_bar=progress_bar,
         )
 
 
@@ -142,39 +151,50 @@ 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
+    billing_project: str | None = None
+    location: str | None = None
+    job_id: str | None = None
+    request_id: str | None = None
+    progress_bar: Literal["default", "auto", "notebook", "terminal"] | None = _DEFAULT
 
     @classmethod
-    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QueryRetryEvent):
+    def from_bqclient(
+        cls,
+        event: google.cloud.bigquery._job_helpers.QueryRetryEvent,
+        progress_bar: Literal["default", "auto", "notebook", "terminal"]
+        | None = _DEFAULT,
+    ):
         return cls(
             query=event.query,
             billing_project=event.billing_project,
             location=event.location,
             job_id=event.job_id,
             request_id=event.request_id,
+            progress_bar=progress_bar,
         )
 
 
 @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
+    billing_project: str | None = None
+    location: str | None = None
+    job_id: str | None = None
+    statement_type: str | None = None
+    state: str | None = None
+    query_plan: list[google.cloud.bigquery.job.query.QueryPlanEntry] | None = None
+    created: datetime.datetime | None = None
+    started: datetime.datetime | None = None
+    ended: datetime.datetime | None = None
+    progress_bar: Literal["default", "auto", "notebook", "terminal"] | None = _DEFAULT
 
     @classmethod
     def from_bqclient(
-        cls, event: google.cloud.bigquery._job_helpers.QueryReceivedEvent
+        cls,
+        event: google.cloud.bigquery._job_helpers.QueryReceivedEvent,
+        progress_bar: Literal["default", "auto", "notebook", "terminal"]
+        | None = _DEFAULT,
     ):
         return cls(
             billing_project=event.billing_project,
@@ -186,28 +206,33 @@ def from_bqclient(
             created=event.created,
             started=event.started,
             ended=event.ended,
+            progress_bar=progress_bar,
         )
 
 
 @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
+    billing_project: str | None = None
+    location: str | None = None
+    query_id: str | None = None
+    job_id: str | None = None
+    destination: google.cloud.bigquery.table.TableReference | None = None
+    total_rows: int | None = None
+    total_bytes_processed: int | None = None
+    slot_millis: int | None = None
+    created: datetime.datetime | None = None
+    started: datetime.datetime | None = None
+    ended: datetime.datetime | None = None
+    progress_bar: Literal["default", "auto", "notebook", "terminal"] | None = _DEFAULT
 
     @classmethod
     def from_bqclient(
-        cls, event: google.cloud.bigquery._job_helpers.QueryFinishedEvent
+        cls,
+        event: google.cloud.bigquery._job_helpers.QueryFinishedEvent,
+        progress_bar: Literal["default", "auto", "notebook", "terminal"]
+        | None = _DEFAULT,
     ):
         return cls(
             billing_project=event.billing_project,
@@ -221,6 +246,7 @@ def from_bqclient(
             created=event.created,
             started=event.started,
             ended=event.ended,
+            progress_bar=progress_bar,
         )
 
 

packages/bigframes/bigframes/formatting_helpers.py

@@ -152,7 +152,10 @@ def progress_callback(
         # This will allow cleanup to continue.
         return
 
-    progress_bar = bigframes._config.options.display.progress_bar
+    # Prioritize progress_bar set on the event, falling back to thread-local option.
+    progress_bar = getattr(event, "progress_bar", bigframes.core.events._DEFAULT)
+    if progress_bar == bigframes.core.events._DEFAULT:
+        progress_bar = bigframes._config.options.display.progress_bar
 
     if progress_bar == "auto":
         progress_bar = "notebook" if in_ipython() else "terminal"

packages/bigframes/bigframes/pandas/io/api.py

@@ -654,8 +654,8 @@ def from_glob_path(
 def _get_bqclient_and_project() -> Tuple[bigquery.Client, str]:
     # Address circular imports in doctest due to bigframes/session/__init__.py
     # containing a lot of logic and samples.
-    from bigframes.session import clients
     import bigframes._config.auth
+    from bigframes.session import clients
 
     credentials, project = bigframes._config.auth.resolve_credentials_and_project(
         config.options.bigquery

packages/bigframes/bigframes/session/_io/bigquery/__init__.py

@@ -245,18 +245,23 @@ def add_and_trim_labels(job_config, session=None):
 
 
 def create_bq_event_callback(publisher):
-    def publish_bq_event(event):
-        if isinstance(event, google.cloud.bigquery._job_helpers.QueryFinishedEvent):
-            bf_event = bigframes.core.events.BigQueryFinishedEvent.from_bqclient(event)
-        elif isinstance(event, google.cloud.bigquery._job_helpers.QueryReceivedEvent):
-            bf_event = bigframes.core.events.BigQueryReceivedEvent.from_bqclient(event)
-        elif isinstance(event, google.cloud.bigquery._job_helpers.QueryRetryEvent):
-            bf_event = bigframes.core.events.BigQueryRetryEvent.from_bqclient(event)
-        elif 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)
+    import bigframes._config
+
+    progress_bar = bigframes._config.options.display.progress_bar
 
+    event_map = {
+        google.cloud.bigquery._job_helpers.QueryFinishedEvent: bigframes.core.events.BigQueryFinishedEvent,
+        google.cloud.bigquery._job_helpers.QueryReceivedEvent: bigframes.core.events.BigQueryReceivedEvent,
+        google.cloud.bigquery._job_helpers.QueryRetryEvent: bigframes.core.events.BigQueryRetryEvent,
+        google.cloud.bigquery._job_helpers.QuerySentEvent: bigframes.core.events.BigQuerySentEvent,
+    }
+
+    def publish_bq_event(event):
+        bf_event = bigframes.core.events.BigQueryUnknownEvent(event)
+        for bq_type, bf_type in event_map.items():
+            if isinstance(event, bq_type):
+                bf_event = bf_type.from_bqclient(event, progress_bar=progress_bar)  # type: ignore
+                break
         publisher.publish(bf_event)
 
     return publish_bq_event

packages/bigframes/tests/system/small/test_progress_bar.py

@@ -104,6 +104,23 @@ def test_progress_bar_load_jobs(
     assert_loading_msg_exist(capsys.readouterr().out, pattern="Load")
 
 
+def test_progress_bar_uniqueness_check(session: bf.Session, capsys):
+    # Ensure strictly_ordered is True (default) to trigger uniqueness check
+    assert session._strictly_ordered
+
+    capsys.readouterr()  # clear output
+
+    with bf.option_context("display.progress_bar", "terminal"):
+        # Read a table and specify a non-unique index_col to trigger the check.
+        # We use a public table to make it a "real" test.
+        session.read_gbq_table(
+            "bigquery-public-data.ml_datasets.penguins",
+            index_col="island",
+        )
+
+    assert_loading_msg_exist(capsys.readouterr().out)
+
+
 def assert_loading_msg_exist(capstdout: str, pattern=job_load_message_regex):
     num_loading_msg = 0
     lines = capstdout.split("\n")

packages/bigframes/tests/unit/test_formatting_helpers.py

@@ -212,3 +212,28 @@ def test_get_job_url():
         job_id=job_id, location=location, project_id=project_id
     )
     assert actual_url == expected_url
+
+
+def test_progress_callback_respects_event_progress_bar():
+    event = bfevents.BigQuerySentEvent(
+        query="SELECT * FROM my_table",
+        progress_bar=None,
+    )
+
+    with mock.patch("bigframes._config.options.display.progress_bar", "terminal"):
+        with mock.patch("bigframes.formatting_helpers.in_ipython", return_value=False):
+            with mock.patch("builtins.print") as mock_print:
+                formatting_helpers.progress_callback(event)
+                mock_print.assert_not_called()
+
+
+def test_progress_callback_falls_back_to_global():
+    event = bfevents.BigQuerySentEvent(
+        query="SELECT * FROM my_table",
+    )
+
+    with mock.patch("bigframes._config.options.display.progress_bar", "terminal"):
+        with mock.patch("bigframes.formatting_helpers.in_ipython", return_value=False):
+            with mock.patch("builtins.print") as mock_print:
+                formatting_helpers.progress_callback(event)
+                mock_print.assert_called_once()