fix remaining mypy errors

Author: tswast

bigframes/pandas/__init__.py

@@ -291,6 +291,7 @@ def clean_up_by_session_id(
             session.bqclient,
             location=location,
             project=project,
+            publisher=session._publisher,
         )
 
     bigframes.session._io.bigquery.delete_tables_matching_session_id(

bigframes/session/__init__.py

@@ -233,6 +233,7 @@ def __init__(
             location=self._location,
             session_id=self._session_id,
             kms_key=self._bq_kms_key_name,
+            publisher=self._publisher,
         )
         # Session temp tables don't support specifying kms key, so use anon dataset if kms key specified
         self._session_resource_manager = (

bigframes/session/anonymous_dataset.py

@@ -20,6 +20,7 @@
 import google.cloud.bigquery as bigquery
 
 from bigframes import constants
+import bigframes.core.events
 from bigframes.session import temporary_storage
 import bigframes.session._io.bigquery as bf_io_bigquery
 
@@ -37,10 +38,12 @@ def __init__(
         location: str,
         session_id: str,
         *,
-        kms_key: Optional[str] = None
+        kms_key: Optional[str] = None,
+        publisher: bigframes.core.events.Publisher,
     ):
         self.bqclient = bqclient
         self._location = location
+        self._publisher = publisher
 
         self.session_id = session_id
         self._table_ids: List[bigquery.TableReference] = []
@@ -62,6 +65,7 @@ def dataset(self) -> bigquery.DatasetReference:
                 self._datset_ref = bf_io_bigquery.create_bq_dataset_reference(
                     self.bqclient,
                     location=self._location,
+                    publisher=self._publisher,
                 )
         return self._datset_ref
 

bigframes/session/loader.py

@@ -647,6 +647,7 @@ def read_gbq_table(
             bq_time=self._clock.get_time(),
             cache=self._df_snapshot,
             use_cache=use_cache,
+            publisher=self._publisher,
         )
 
         if table.location.casefold() != self._storage_manager.location.casefold():
@@ -767,6 +768,7 @@ def read_gbq_table(
             filter_str,
             should_warn=True,
             should_dry_run=True,
+            publisher=self._publisher,
         )
 
         # ----------------------------
@@ -793,6 +795,7 @@ def read_gbq_table(
                 self._bqclient,
                 table=table,
                 index_cols=index_cols,
+                publisher=self._publisher,
             )
             if publish_execution:
                 self._publisher.send(

tests/system/small/engines/conftest.py

@@ -19,7 +19,7 @@
 import pytest
 
 import bigframes
-from bigframes.core import ArrayValue, local_data
+from bigframes.core import ArrayValue, events, local_data
 from bigframes.session import (
     direct_gbq_execution,
     local_scan_executor,
@@ -50,11 +50,14 @@ def engine(request, bigquery_client: bigquery.Client) -> semi_executor.SemiExecu
         return local_scan_executor.LocalScanExecutor()
     if request.param == "polars":
         return polars_executor.PolarsExecutor()
+    publisher = events.Publisher()
     if request.param == "bq":
-        return direct_gbq_execution.DirectGbqExecutor(bigquery_client)
+        return direct_gbq_execution.DirectGbqExecutor(
+            bigquery_client, publisher=publisher
+        )
     if request.param == "bq-sqlglot":
         return direct_gbq_execution.DirectGbqExecutor(
-            bigquery_client, compiler="sqlglot"
+            bigquery_client, compiler="sqlglot", publisher=publisher
         )
     raise ValueError(f"Unrecognized param: {request.param}")
 

tests/system/small/engines/test_aggregation.py

@@ -15,7 +15,14 @@
 from google.cloud import bigquery
 import pytest
 
-from bigframes.core import agg_expressions, array_value, expression, identifiers, nodes
+from bigframes.core import (
+    agg_expressions,
+    array_value,
+    events,
+    expression,
+    identifiers,
+    nodes,
+)
 import bigframes.operations.aggregations as agg_ops
 from bigframes.session import direct_gbq_execution, polars_executor
 from bigframes.testing.engine_utils import assert_equivalence_execution
@@ -93,9 +100,12 @@ def test_sql_engines_median_op_aggregates(
         scalars_array_value,
         agg_ops.MedianOp(),
     ).node
-    left_engine = direct_gbq_execution.DirectGbqExecutor(bigquery_client)
+    publisher = events.Publisher()
+    left_engine = direct_gbq_execution.DirectGbqExecutor(
+        bigquery_client, publisher=publisher
+    )
     right_engine = direct_gbq_execution.DirectGbqExecutor(
-        bigquery_client, compiler="sqlglot"
+        bigquery_client, compiler="sqlglot", publisher=publisher
     )
     assert_equivalence_execution(node, left_engine, right_engine)
 

tests/system/small/engines/test_windowing.py

@@ -18,6 +18,7 @@
 from bigframes.core import (
     agg_expressions,
     array_value,
+    events,
     expression,
     identifiers,
     nodes,
@@ -64,8 +65,11 @@ def test_engines_with_rows_window(
         skip_reproject_unsafe=False,
     )
 
-    bq_executor = direct_gbq_execution.DirectGbqExecutor(bigquery_client)
+    publisher = events.Publisher()
+    bq_executor = direct_gbq_execution.DirectGbqExecutor(
+        bigquery_client, publisher=publisher
+    )
     bq_sqlgot_executor = direct_gbq_execution.DirectGbqExecutor(
-        bigquery_client, compiler="sqlglot"
+        bigquery_client, compiler="sqlglot", publisher=publisher
     )
     assert_equivalence_execution(window_node, bq_executor, bq_sqlgot_executor)

tests/unit/session/test_read_gbq_table.py

@@ -139,6 +139,7 @@ def test_check_if_index_columns_are_unique(index_cols, values_distinct, expected
         bqclient=bqclient,
         table=table,
         index_cols=index_cols,
+        publisher=session._publisher,
     )
 
     assert result == expected