move publisher to session

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

@@ -68,10 +68,6 @@ def send(self, event: Event):
             callback(event)
 
 
-publisher = Publisher()
-publisher.subscribe(bigframes.formatting_helpers.progress_callback)
-
-
 class Event:
     pass
 

bigframes/functions/_function_client.py

@@ -145,6 +145,7 @@ def _create_bq_function(self, create_function_ddl: str) -> None:
             timeout=None,
             metrics=None,
             query_with_job=True,
+            publisher=self._session._publisher,
         )
         logger.info(f"Created bigframes function {query_job.ddl_target_routine}")
 

bigframes/functions/function.py

@@ -219,7 +219,13 @@ def __call__(self, *args, **kwargs):
 
         args_string = ", ".join(map(bf_sql.simple_literal, args))
         sql = f"SELECT `{str(self._udf_def.routine_ref)}`({args_string})"
-        iter, job = bf_io_bigquery.start_query_with_client(self._session.bqclient, sql=sql, query_with_job=True, job_config=bigquery.QueryJobConfig())  # type: ignore
+        iter, job = bf_io_bigquery.start_query_with_client(
+            self._session.bqclient,
+            sql=sql,
+            query_with_job=True,
+            job_config=bigquery.QueryJobConfig(),
+            publisher=self._session._publisher,
+        )  # type: ignore
         return list(iter.to_arrow().to_pydict().values())[0][0]
 
     @property
@@ -297,7 +303,13 @@ def __call__(self, *args, **kwargs):
 
         args_string = ", ".join(map(bf_sql.simple_literal, args))
         sql = f"SELECT `{str(self._udf_def.routine_ref)}`({args_string})"
-        iter, job = bf_io_bigquery.start_query_with_client(self._session.bqclient, sql=sql, query_with_job=True, job_config=bigquery.QueryJobConfig())  # type: ignore
+        iter, job = bf_io_bigquery.start_query_with_client(
+            self._session.bqclient,
+            sql=sql,
+            query_with_job=True,
+            job_config=bigquery.QueryJobConfig(),
+            publisher=self._session._publisher,
+        )  # type: ignore
         return list(iter.to_arrow().to_pydict().values())[0][0]
 
     @property

bigframes/session/__init__.py

@@ -67,10 +67,9 @@
 import bigframes.constants
 import bigframes.core
 from bigframes.core import blocks, log_adapter, utils
+import bigframes.core.events
 import bigframes.core.pyformat
-
-# Even though the ibis.backends.bigquery import is unused, it's needed
-# to register new and replacement ops with the Ibis BigQuery backend.
+import bigframes.formatting_helpers
 import bigframes.functions._function_session as bff_session
 import bigframes.functions.function as bff
 from bigframes.session import bigquery_session, bq_caching_executor, executor
@@ -137,6 +136,11 @@ def __init__(
 
         _warn_if_bf_version_is_obsolete()
 
+        # Publisher needs to be created before the other objects, especially
+        # the executors, because they access it.
+        self._publisher = bigframes.core.events.Publisher()
+        self._publisher.subscribe(bigframes.formatting_helpers.progress_callback)
+
         if context is None:
             context = bigquery_options.BigQueryOptions()
 
@@ -251,6 +255,7 @@ def __init__(
             scan_index_uniqueness=self._strictly_ordered,
             force_total_order=self._strictly_ordered,
             metrics=self._metrics,
+            publisher=self._publisher,
         )
         self._executor: executor.Executor = bq_caching_executor.BigQueryCachingExecutor(
             bqclient=self._clients_provider.bqclient,
@@ -260,6 +265,7 @@ def __init__(
             strictly_ordered=self._strictly_ordered,
             metrics=self._metrics,
             enable_polars_execution=context.enable_polars_execution,
+            publisher=self._publisher,
         )
 
     def __del__(self):
@@ -2150,6 +2156,7 @@ def _start_query_ml_ddl(
             timeout=None,
             query_with_job=True,
             job_retry=third_party_gcb_retry.DEFAULT_ML_JOB_RETRY,
+            publisher=self._publisher,
         )
         return iterator, query_job
 
@@ -2177,6 +2184,7 @@ def _create_object_table(self, path: str, connection: str) -> str:
             project=None,
             timeout=None,
             query_with_job=True,
+            publisher=self._publisher,
         )
 
         return table

bigframes/session/_io/bigquery/__init__.py

@@ -240,19 +240,22 @@ def add_and_trim_labels(job_config):
     )
 
 
-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)
+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)
 
-    bigframes.core.events.publisher.send(bf_event)
+        publisher.send(bf_event)
+
+    return publish_bq_event
 
 
 @overload
@@ -266,6 +269,7 @@ def start_query_with_client(
     timeout: Optional[float],
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[True],
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]:
     ...
 
@@ -281,6 +285,7 @@ def start_query_with_client(
     timeout: Optional[float],
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[False],
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     ...
 
@@ -297,6 +302,7 @@ def start_query_with_client(
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[True],
     job_retry: google.api_core.retry.Retry,
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]:
     ...
 
@@ -313,6 +319,7 @@ def start_query_with_client(
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[False],
     job_retry: google.api_core.retry.Retry,
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     ...
 
@@ -332,6 +339,7 @@ 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,
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     """
     Starts query job and waits for results.
@@ -350,7 +358,7 @@ def start_query_with_client(
                 project=project,
                 api_timeout=timeout,
                 job_retry=job_retry,
-                callback=publish_bq_event,
+                callback=create_bq_event_callback(publisher),
             )
             if metrics is not None:
                 metrics.count_job_stats(row_iterator=results_iterator)
@@ -370,7 +378,7 @@ def start_query_with_client(
         raise
 
     if not query_job.configuration.dry_run:
-        bigframes.core.events.publisher.send(
+        publisher.send(
             bigframes.core.events.BigQuerySentEvent(
                 sql,
                 billing_project=query_job.project,
@@ -381,7 +389,7 @@ def start_query_with_client(
         )
     results_iterator = query_job.result()
     if not query_job.configuration.dry_run:
-        bigframes.core.events.publisher.send(
+        publisher.send(
             bigframes.core.events.BigQueryFinishedEvent(
                 billing_project=query_job.project,
                 location=query_job.location,
@@ -436,6 +444,8 @@ def create_bq_dataset_reference(
     bq_client: bigquery.Client,
     location: Optional[str] = None,
     project: Optional[str] = None,
+    *,
+    publisher: bigframes.core.events.Publisher,
 ) -> bigquery.DatasetReference:
     """Create and identify dataset(s) for temporary BQ resources.
 
@@ -467,6 +477,7 @@ def create_bq_dataset_reference(
         timeout=None,
         metrics=None,
         query_with_job=True,
+        publisher=publisher,
     )
 
     # The anonymous dataset is used by BigQuery to write query results and

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -28,6 +28,7 @@
 import google.cloud.bigquery as bigquery
 import google.cloud.bigquery.table
 
+import bigframes.core.events
 import bigframes.exceptions as bfe
 import bigframes.session._io.bigquery
 
@@ -43,6 +44,7 @@ def get_table_metadata(
     *,
     cache: Dict[bigquery.TableReference, Tuple[datetime.datetime, bigquery.Table]],
     use_cache: bool = True,
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[datetime.datetime, google.cloud.bigquery.table.Table]:
     """Get the table metadata, either from cache or via REST API."""
 
@@ -59,6 +61,7 @@ def get_table_metadata(
             # Don't warn, because that will already have been taken care of.
             should_warn=False,
             should_dry_run=False,
+            publisher=publisher,
         ):
             # This warning should only happen if the cached snapshot_time will
             # have any effect on bigframes (b/437090788). For example, with
@@ -108,6 +111,7 @@ def is_time_travel_eligible(
     *,
     should_warn: bool,
     should_dry_run: bool,
+    publisher: bigframes.core.events.Publisher,
 ):
     """Check if a table is eligible to use time-travel.
 
@@ -184,6 +188,7 @@ def is_time_travel_eligible(
                 timeout=None,
                 metrics=None,
                 query_with_job=False,
+                publisher=publisher,
             )
             return True
 
@@ -235,6 +240,8 @@ def check_if_index_columns_are_unique(
     bqclient: bigquery.Client,
     table: google.cloud.bigquery.table.Table,
     index_cols: List[str],
+    *,
+    publisher: bigframes.core.events.Publisher,
 ) -> Tuple[str, ...]:
     import bigframes.core.sql
     import bigframes.session._io.bigquery
@@ -252,6 +259,7 @@ def check_if_index_columns_are_unique(
         project=None,
         metrics=None,
         query_with_job=False,
+        publisher=publisher,
     )
     row = next(iter(results))
 

bigframes/session/bq_caching_executor.py

@@ -140,6 +140,7 @@ def __init__(
         strictly_ordered: bool = True,
         metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
         enable_polars_execution: bool = False,
+        publisher: bigframes.core.events.Publisher,
     ):
         self.bqclient = bqclient
         self.storage_manager = storage_manager
@@ -149,6 +150,9 @@ def __init__(
         self.loader = loader
         self.bqstoragereadclient = bqstoragereadclient
         self._enable_polars_execution = enable_polars_execution
+        self._publisher = publisher
+
+        # TODO(tswast): Send events from semi-executors, too.
         self._semi_executors: Sequence[semi_executor.SemiExecutor] = (
             read_api_execution.ReadApiSemiExecutor(
                 bqstoragereadclient=bqstoragereadclient,
@@ -188,7 +192,7 @@ def execute(
         array_value: bigframes.core.ArrayValue,
         execution_spec: ex_spec.ExecutionSpec,
     ) -> executor.ExecuteResult:
-        bigframes.core.events.publisher.send(bigframes.core.events.ExecutionStarted())
+        self._publisher.send(bigframes.core.events.ExecutionStarted())
 
         # TODO: Support export jobs in combination with semi executors
         if execution_spec.destination_spec is None:
@@ -198,7 +202,7 @@ def execute(
                     plan, ordered=execution_spec.ordered, peek=execution_spec.peek
                 )
                 if maybe_result:
-                    bigframes.core.events.publisher.send(
+                    self._publisher.send(
                         bigframes.core.events.ExecutionFinished(
                             result=maybe_result,
                         )
@@ -212,7 +216,7 @@ def execute(
                 )
             # separate path for export_gbq, as it has all sorts of annoying logic, such as possibly running as dml
             result = self._export_gbq(array_value, execution_spec.destination_spec)
-            bigframes.core.events.publisher.send(
+            self._publisher.send(
                 bigframes.core.events.ExecutionFinished(
                     result=result,
                 )
@@ -232,7 +236,7 @@ def execute(
         if isinstance(execution_spec.destination_spec, ex_spec.GcsOutputSpec):
             self._export_result_gcs(result, execution_spec.destination_spec)
 
-        bigframes.core.events.publisher.send(
+        self._publisher.send(
             bigframes.core.events.ExecutionFinished(
                 result=result,
             )
@@ -261,6 +265,7 @@ def _export_result_gcs(
             location=None,
             timeout=None,
             query_with_job=True,
+            publisher=self._publisher,
         )
 
     def _maybe_find_existing_table(
@@ -419,6 +424,7 @@ def _run_execute_query(
                     location=None,
                     timeout=None,
                     query_with_job=True,
+                    publisher=self._publisher,
                 )
             else:
                 return bq_io.start_query_with_client(
@@ -430,6 +436,7 @@ def _run_execute_query(
                     location=None,
                     timeout=None,
                     query_with_job=False,
+                    publisher=self._publisher,
                 )
 
         except google.api_core.exceptions.BadRequest as e: