feat: Add component to manage temporary tables (#1559)

TrevorBergeron · web-flow · commit 0a4e245670e6 · 2025-03-28T13:26:59.000-07:00
diff --git a/bigframes/session/bigquery_session.py b/bigframes/session/bigquery_session.py
@@ -0,0 +1,171 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import logging
+import threading
+from typing import Callable, Optional, Sequence
+import uuid
+
+# TODO: Non-ibis implementation
+import bigframes_vendored.ibis.backends.bigquery.datatypes as ibis_bq
+import google.cloud.bigquery as bigquery
+
+from bigframes.core.compile import googlesql
+
+KEEPALIVE_QUERY_TIMEOUT_SECONDS = 5.0
+
+KEEPALIVE_FREQUENCY = datetime.timedelta(hours=6)
+
+
+logger = logging.getLogger(__name__)
+
+
+class SessionResourceManager:
+    """
+    Responsible for allocating and cleaning up temporary gbq tables used by a BigFrames session.
+    """
+
+    def __init__(
+        self, bqclient: bigquery.Client, location: str, *, kms_key: Optional[str] = None
+    ):
+        self.bqclient = bqclient
+        self.location = location
+        self._kms_key = kms_key
+        self._session_id: Optional[str] = None
+        self._sessiondaemon: Optional[RecurringTaskDaemon] = None
+        self._session_lock = threading.RLock()
+
+    def create_temp_table(
+        self, schema: Sequence[bigquery.SchemaField], cluster_cols: Sequence[str] = []
+    ) -> bigquery.TableReference:
+        """Create a temporary session table. Session is an exclusive resource, so throughput is limited"""
+        # Can't set a table in _SESSION as destination via query job API, so we
+        # run DDL, instead.
+        with self._session_lock:
+            table_ref = bigquery.TableReference(
+                bigquery.DatasetReference(self.bqclient.project, "_SESSION"),
+                uuid.uuid4().hex,
+            )
+            job_config = bigquery.QueryJobConfig(
+                connection_properties=[
+                    bigquery.ConnectionProperty("session_id", self._get_session_id())
+                ]
+            )
+            if self._kms_key:
+                job_config.destination_encryption_configuration = (
+                    bigquery.EncryptionConfiguration(kms_key_name=self._kms_key)
+                )
+
+            ibis_schema = ibis_bq.BigQuerySchema.to_ibis(list(schema))
+
+            fields = [
+                f"{googlesql.identifier(name)} {ibis_bq.BigQueryType.from_ibis(ibis_type)}"
+                for name, ibis_type in ibis_schema.fields.items()
+            ]
+            fields_string = ",".join(fields)
+
+            cluster_string = ""
+            if cluster_cols:
+                cluster_cols_sql = ", ".join(
+                    f"{googlesql.identifier(cluster_col)}"
+                    for cluster_col in cluster_cols
+                )
+                cluster_string = f"\nCLUSTER BY {cluster_cols_sql}"
+
+            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)
+
+            job.result()
+            # return the fully qualified table, so it can be used outside of the session
+            return job.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(f"CALL BQ.ABORT_SESSION('{self._session_id}')")
+
+    def _get_session_id(self) -> str:
+        if self._session_id:
+            return self._session_id
+        with self._session_lock:
+            if self._session_id is None:
+                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.result()  # blocks until finished
+                assert query_job.session_info is not None
+                assert query_job.session_info.session_id is not None
+                self._session_id = query_job.session_info.session_id
+                self._sessiondaemon = RecurringTaskDaemon(
+                    task=self._keep_session_alive, frequency=KEEPALIVE_FREQUENCY
+                )
+                self._sessiondaemon.start()
+                return query_job.session_info.session_id
+            else:
+                return self._session_id
+
+    def _keep_session_alive(self):
+        # bq sessions will default expire after 24 hours of disuse, but if queried, this is renewed to a maximum of 7 days
+        with self._session_lock:
+            job_config = bigquery.QueryJobConfig(
+                connection_properties=[
+                    bigquery.ConnectionProperty("session_id", self._get_session_id())
+                ]
+            )
+            try:
+                self.bqclient.query_and_wait(
+                    "SELECT 1",
+                    location=self.location,
+                    job_config=job_config,
+                    wait_timeout=KEEPALIVE_QUERY_TIMEOUT_SECONDS,
+                )
+            except Exception as e:
+                logging.warning("BigQuery session keep-alive query errored : %s", e)
+
+
+class RecurringTaskDaemon:
+    def __init__(self, task: Callable[[], None], frequency: datetime.timedelta):
+        self._stop_event = threading.Event()
+        self._frequency = frequency
+        self._thread = threading.Thread(target=self._run_loop, daemon=True)
+        self._task = task
+
+    def start(self):
+        """Start the daemon. Cannot be restarted once stopped."""
+        if self._stop_event.is_set():
+            raise RuntimeError("Cannot restart daemon thread.")
+        self._thread.start()
+
+    def _run_loop(self):
+        while True:
+            self._stop_event.wait(self._frequency.total_seconds())
+            if self._stop_event.is_set():
+                return
+            try:
+                self._task()
+            except Exception as e:
+                logging.warning("RecurringTaskDaemon task errorred: %s", e)
+
+    def stop(self, timeout_seconds: Optional[float] = None):
+        """Stop and cleanup the daemon."""
+        if self._thread.is_alive():
+            self._stop_event.set()
+            self._thread.join(timeout=timeout_seconds)
diff --git a/tests/system/small/test_bq_sessions.py b/tests/system/small/test_bq_sessions.py
@@ -0,0 +1,80 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from concurrent.futures import ThreadPoolExecutor
+
+import google
+import google.api_core.exceptions
+import google.cloud
+from google.cloud import bigquery
+import pytest
+
+from bigframes.session import bigquery_session
+
+TEST_SCHEMA = [
+    bigquery.SchemaField("bool field", "BOOLEAN"),
+    bigquery.SchemaField("string field", "STRING"),
+    bigquery.SchemaField("float array_field", "FLOAT", mode="REPEATED"),
+    bigquery.SchemaField(
+        "struct field",
+        "RECORD",
+        fields=(bigquery.SchemaField("int subfield", "INTEGER"),),
+    ),
+]
+
+
+@pytest.fixture
+def session_resource_manager(
+    bigquery_client,
+) -> bigquery_session.SessionResourceManager:
+    return bigquery_session.SessionResourceManager(bigquery_client, "US")
+
+
+def test_bq_session_create_temp_table_clustered(bigquery_client: bigquery.Client):
+    session_resource_manager = bigquery_session.SessionResourceManager(
+        bigquery_client, "US"
+    )
+    cluster_cols = ["string field", "bool field"]
+
+    session_table_ref = session_resource_manager.create_temp_table(
+        TEST_SCHEMA, cluster_cols=cluster_cols
+    )
+    session_resource_manager._keep_session_alive()
+
+    result_table = bigquery_client.get_table(session_table_ref)
+    assert result_table.schema == TEST_SCHEMA
+    assert result_table.clustering_fields == cluster_cols
+
+    session_resource_manager.close()
+    with pytest.raises(google.api_core.exceptions.NotFound):
+        bigquery_client.get_table(session_table_ref)
+
+
+def test_bq_session_create_multi_temp_tables(bigquery_client: bigquery.Client):
+    session_resource_manager = bigquery_session.SessionResourceManager(
+        bigquery_client, "US"
+    )
+
+    def create_table():
+        return session_resource_manager.create_temp_table(TEST_SCHEMA)
+
+    with ThreadPoolExecutor() as executor:
+        results = [executor.submit(create_table) for i in range(10)]
+
+    for future in results:
+        table = future.result()
+        result_table = bigquery_client.get_table(table)
+        assert result_table.schema == TEST_SCHEMA
+
+    session_resource_manager.close()
diff --git a/tests/unit/test_daemon.py b/tests/unit/test_daemon.py
@@ -0,0 +1,42 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import time
+from unittest.mock import MagicMock
+
+from bigframes.session.bigquery_session import RecurringTaskDaemon
+
+
+def test_recurring_task_daemon_calls():
+    mock_task = MagicMock()
+    daemon = RecurringTaskDaemon(
+        task=mock_task, frequency=datetime.timedelta(seconds=0.1)
+    )
+    daemon.start()
+    time.sleep(1.0)
+    daemon.stop()
+    time.sleep(0.5)
+    # be lenient, but number of calls should be in this ballpark regardless of scheduling hiccups
+    assert mock_task.call_count > 6
+    assert mock_task.call_count < 12
+
+
+def test_recurring_task_daemon_never_started():
+    mock_task = MagicMock()
+    _ = RecurringTaskDaemon(
+        task=mock_task, frequency=datetime.timedelta(seconds=0.0001)
+    )
+    time.sleep(0.1)
+    assert mock_task.call_count == 0
