add explicit unsubscribe

Author: tswast

bigframes/core/events.py

@@ -16,56 +16,70 @@
 
 import dataclasses
 import datetime
-import threading
-from typing import List, Optional
-import weakref
+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.formatting_helpers
 import bigframes.session.executor
 
 
-@dataclasses.dataclass(frozen=True)
 class Subscriber:
-    callback_ref: weakref.ref
-    # TODO(tswast): Add block_id to allow filter in context managers.
+    def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):
+        self._publisher = publisher
+        self._callback = callback
+        self._subscriber_id = str(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: List[Subscriber] = []
-        self._subscribers_lock = threading.Lock()
-
-    def subscribe(self, callback):
-        subscriber = Subscriber(callback_ref=weakref.ref(callback))
-
-        with self._subscribers_lock:
-            # TODO(tswast): Add block_id to allow filter in context managers.
-            self._subscribers.append(subscriber)
-
-    def send(self, event: Event):
-        to_delete = []
-        to_call = []
+        self._subscribers: Set[Subscriber] = set()
 
-        with self._subscribers_lock:
-            for sid, subscriber in enumerate(self._subscribers):
-                callback = subscriber.callback_ref()
+    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)
+        self._subscribers.add(subscriber)
+        return subscriber
 
-                if callback is None:
-                    to_delete.append(sid)
-                else:
-                    # TODO(tswast): Add if statement for block_id to allow filter
-                    # in context managers.
-                    to_call.append(callback)
+    def unsubscribe(self, subscriber: Subscriber):
+        self._subscribers.remove(subscriber)
 
-            for sid in reversed(to_delete):
-                del self._subscribers[sid]
-
-        for callback in to_call:
-            callback(event)
+    def publish(self, event: Event):
+        for subscriber in self._subscribers:
+            subscriber(event)
 
 
 class Event:
@@ -90,6 +104,13 @@ 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."""

bigframes/session/__init__.py

@@ -383,7 +383,9 @@ def close(self):
 
         publisher_session = getattr(self, "_publisher", None)
         if publisher_session:
-            publisher_session.send(self.session_id)
+            publisher_session.publish(
+                bigframes.core.events.SessionClosed(self.session_id)
+            )
 
     @overload
     def read_gbq(  # type: ignore[overload-overlap]

bigframes/session/_io/bigquery/__init__.py

@@ -253,7 +253,7 @@ def publish_bq_event(event):
         else:
             bf_event = bigframes.core.events.BigQueryUnknownEvent(event)
 
-        publisher.send(bf_event)
+        publisher.publish(bf_event)
 
     return publish_bq_event
 
@@ -378,7 +378,7 @@ def start_query_with_client(
         raise
 
     if not query_job.configuration.dry_run:
-        publisher.send(
+        publisher.publish(
             bigframes.core.events.BigQuerySentEvent(
                 sql,
                 billing_project=query_job.project,
@@ -389,7 +389,7 @@ def start_query_with_client(
         )
     results_iterator = query_job.result()
     if not query_job.configuration.dry_run:
-        publisher.send(
+        publisher.publish(
             bigframes.core.events.BigQueryFinishedEvent(
                 billing_project=query_job.project,
                 location=query_job.location,

bigframes/session/bq_caching_executor.py

@@ -192,7 +192,7 @@ def execute(
         array_value: bigframes.core.ArrayValue,
         execution_spec: ex_spec.ExecutionSpec,
     ) -> executor.ExecuteResult:
-        self._publisher.send(bigframes.core.events.ExecutionStarted())
+        self._publisher.publish(bigframes.core.events.ExecutionStarted())
 
         # TODO: Support export jobs in combination with semi executors
         if execution_spec.destination_spec is None:
@@ -202,7 +202,7 @@ def execute(
                     plan, ordered=execution_spec.ordered, peek=execution_spec.peek
                 )
                 if maybe_result:
-                    self._publisher.send(
+                    self._publisher.publish(
                         bigframes.core.events.ExecutionFinished(
                             result=maybe_result,
                         )
@@ -216,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)
-            self._publisher.send(
+            self._publisher.publish(
                 bigframes.core.events.ExecutionFinished(
                     result=result,
                 )
@@ -236,7 +236,7 @@ def execute(
         if isinstance(execution_spec.destination_spec, ex_spec.GcsOutputSpec):
             self._export_result_gcs(result, execution_spec.destination_spec)
 
-        self._publisher.send(
+        self._publisher.publish(
             bigframes.core.events.ExecutionFinished(
                 result=result,
             )

bigframes/session/loader.py

@@ -788,7 +788,7 @@ def read_gbq_table(
         # If non in strict ordering mode, don't go through overhead of scanning index column(s) to determine if unique
         if not primary_key and self._scan_index_uniqueness and index_cols:
             if publish_execution:
-                self._publisher.send(
+                self._publisher.publish(
                     bigframes.core.events.ExecutionStarted(),
                 )
             primary_key = bf_read_gbq_table.check_if_index_columns_are_unique(
@@ -798,7 +798,7 @@ def read_gbq_table(
                 publisher=self._publisher,
             )
             if publish_execution:
-                self._publisher.send(
+                self._publisher.publish(
                     bigframes.core.events.ExecutionFinished(),
                 )
 
@@ -1021,7 +1021,7 @@ def read_gbq_query(
 
         # We want to make sure we show progress when we actually do execute a
         # query. Since we have got this far, we know it's not a dry run.
-        self._publisher.send(
+        self._publisher.publish(
             bigframes.core.events.ExecutionStarted(),
         )
 
@@ -1086,7 +1086,7 @@ def read_gbq_query(
                 index_col=index_col,
                 columns=columns,
             )
-            self._publisher.send(
+            self._publisher.publish(
                 bigframes.core.events.ExecutionFinished(),
             )
             return df
@@ -1098,7 +1098,7 @@ def read_gbq_query(
                 query_job_for_metrics,
                 session=self._session,
             )
-            self._publisher.send(
+            self._publisher.publish(
                 bigframes.core.events.ExecutionFinished(),
             )
             return df
@@ -1116,7 +1116,7 @@ def read_gbq_query(
                 query_job_for_metrics,
                 session=self._session,
             )
-            self._publisher.send(
+            self._publisher.publish(
                 bigframes.core.events.ExecutionFinished(),
             )
             return df
@@ -1140,7 +1140,7 @@ def read_gbq_query(
             # max_results and filters are omitted because they are already
             # handled by to_query(), above.
         )
-        self._publisher.send(
+        self._publisher.publish(
             bigframes.core.events.ExecutionFinished(),
         )
         return df