add lock for publisher and publish temp table creations

Author: tswast

bigframes/core/events.py

@@ -16,6 +16,7 @@
 
 import dataclasses
 import datetime
+import threading
 from typing import Any, Callable, Optional, Set
 import uuid
 
@@ -30,7 +31,7 @@ class Subscriber:
     def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):
         self._publisher = publisher
         self._callback = callback
-        self._subscriber_id = str(uuid.uuid4())
+        self._subscriber_id = uuid.uuid4()
 
     def __call__(self, *args, **kwargs):
         return self._callback(*args, **kwargs)
@@ -65,21 +66,25 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 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)
-        self._subscribers.add(subscriber)
+        with self._subscribers_lock:
+            self._subscribers.add(subscriber)
         return subscriber
 
     def unsubscribe(self, subscriber: Subscriber):
-        self._subscribers.remove(subscriber)
+        with self._subscribers_lock:
+            self._subscribers.remove(subscriber)
 
     def publish(self, event: Event):
-        for subscriber in self._subscribers:
-            subscriber(event)
+        with self._subscribers_lock:
+            for subscriber in self._subscribers:
+                subscriber(event)
 
 
 class Event:

bigframes/session/__init__.py

@@ -240,6 +240,7 @@ def __init__(
             bigquery_session.SessionResourceManager(
                 self.bqclient,
                 self._location,
+                publisher=self._publisher,
             )
             if (self._bq_kms_key_name is None)
             else None

bigframes/session/bigquery_session.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import datetime
 import logging
 import threading
@@ -23,7 +25,9 @@
 import google.cloud.bigquery as bigquery
 
 from bigframes.core.compile import googlesql
+import bigframes.core.events
 from bigframes.session import temporary_storage
+import bigframes.session._io.bigquery as bfbqio
 
 KEEPALIVE_QUERY_TIMEOUT_SECONDS = 5.0
 
@@ -38,12 +42,19 @@ class SessionResourceManager(temporary_storage.TemporaryStorageManager):
     Responsible for allocating and cleaning up temporary gbq tables used by a BigFrames session.
     """
 
-    def __init__(self, bqclient: bigquery.Client, location: str):
+    def __init__(
+        self,
+        bqclient: bigquery.Client,
+        location: str,
+        *,
+        publisher: bigframes.core.events.Publisher,
+    ):
         self.bqclient = bqclient
         self._location = location
         self._session_id: Optional[str] = None
         self._sessiondaemon: Optional[RecurringTaskDaemon] = None
         self._session_lock = threading.RLock()
+        self._publisher = publisher
 
     @property
     def location(self):
@@ -84,21 +95,38 @@ def create_temp_table(
 
             ddl = f"CREATE TEMP TABLE `_SESSION`.{googlesql.identifier(table_ref.table_id)} ({fields_string}){cluster_string}"
 
-            job = self.bqclient.query(
-                ddl, job_config=job_config, location=self.location
+            _, job = bfbqio.start_query_with_client(
+                self.bqclient,
+                ddl,
+                job_config=job_config,
+                location=self.location,
+                project=None,
+                timeout=None,
+                metrics=None,
+                query_with_job=True,
+                publisher=self._publisher,
             )
             job.result()
             # return the fully qualified table, so it can be used outside of the session
-            return job.destination
+            destination = job.destination
+            assert destination is not None, "Failure to create temp table."
+            return destination
 
     def close(self):
         if self._sessiondaemon is not None:
             self._sessiondaemon.stop()
 
         if self._session_id is not None and self.bqclient is not None:
-            self.bqclient.query_and_wait(
+            bfbqio.start_query_with_client(
+                self.bqclient,
                 f"CALL BQ.ABORT_SESSION('{self._session_id}')",
+                job_config=bigquery.QueryJobConfig(),
                 location=self.location,
+                project=None,
+                timeout=None,
+                metrics=None,
+                query_with_job=False,
+                publisher=self._publisher,
             )
 
     def _get_session_id(self) -> str:
@@ -109,8 +137,16 @@ def _get_session_id(self) -> str:
                 job_config = bigquery.QueryJobConfig(create_session=True)
                 # Make sure the session is a new one, not one associated with another query.
                 job_config.use_query_cache = False
-                query_job = self.bqclient.query(
-                    "SELECT 1", job_config=job_config, location=self.location
+                _, query_job = bfbqio.start_query_with_client(
+                    self.bqclient,
+                    "SELECT 1",
+                    job_config=job_config,
+                    location=self.location,
+                    project=None,
+                    timeout=None,
+                    metrics=None,
+                    query_with_job=True,
+                    publisher=self._publisher,
                 )
                 query_job.result()  # blocks until finished
                 assert query_job.session_info is not None
@@ -133,11 +169,16 @@ def _keep_session_alive(self):
                 ]
             )
             try:
-                self.bqclient.query_and_wait(
+                bfbqio.start_query_with_client(
+                    self.bqclient,
                     "SELECT 1",
-                    location=self.location,
                     job_config=job_config,
-                    wait_timeout=KEEPALIVE_QUERY_TIMEOUT_SECONDS,
+                    location=self.location,
+                    project=None,
+                    timeout=KEEPALIVE_QUERY_TIMEOUT_SECONDS,
+                    metrics=None,
+                    query_with_job=False,
+                    publisher=self._publisher,
                 )
             except Exception as e:
                 logging.warning("BigQuery session keep-alive query errored : %s", e)

tests/system/small/test_bq_sessions.py

@@ -17,10 +17,10 @@
 
 import google
 import google.api_core.exceptions
-import google.cloud
 from google.cloud import bigquery
 import pytest
 
+import bigframes.core.events
 from bigframes.session import bigquery_session
 
 TEST_SCHEMA = [
@@ -39,12 +39,14 @@
 def session_resource_manager(
     bigquery_client,
 ) -> bigquery_session.SessionResourceManager:
-    return bigquery_session.SessionResourceManager(bigquery_client, "US")
+    return bigquery_session.SessionResourceManager(
+        bigquery_client, "US", publisher=bigframes.core.events.Publisher()
+    )
 
 
 def test_bq_session_create_temp_table_clustered(bigquery_client: bigquery.Client):
     session_resource_manager = bigquery_session.SessionResourceManager(
-        bigquery_client, "US"
+        bigquery_client, "US", publisher=bigframes.core.events.Publisher()
     )
     cluster_cols = ["string field", "bool field"]
 
@@ -68,7 +70,7 @@ def test_bq_session_create_temp_table_clustered(bigquery_client: bigquery.Client
 
 def test_bq_session_create_multi_temp_tables(bigquery_client: bigquery.Client):
     session_resource_manager = bigquery_session.SessionResourceManager(
-        bigquery_client, "US"
+        bigquery_client, "US", publisher=bigframes.core.events.Publisher()
     )
 
     def create_table():