Feat: Support sessions in the BigQuery adapter (#1185)

Author: izeigerman

sqlmesh/core/engine_adapter/base.py

@@ -840,7 +840,7 @@ def fetch_pyspark_df(
     @contextlib.contextmanager
     def transaction(
         self, transaction_type: TransactionType = TransactionType.DML
-    ) -> t.Generator[None, None, None]:
+    ) -> t.Iterator[None]:
         """A transaction context manager."""
         if self._connection_pool.is_transaction_active or not self.supports_transactions(
             transaction_type
@@ -860,6 +860,29 @@ def supports_transactions(self, transaction_type: TransactionType) -> bool:
         """Whether or not the engine adapter supports transactions for the given transaction type."""
         return True
 
+    @contextlib.contextmanager
+    def session(self) -> t.Iterator[None]:
+        """A session context manager."""
+        if self._is_session_active():
+            yield
+            return
+
+        self._begin_session()
+        try:
+            yield
+        finally:
+            self._end_session()
+
+    def _begin_session(self) -> None:
+        """Begin a new session."""
+
+    def _end_session(self) -> None:
+        """End the existing session."""
+
+    def _is_session_active(self) -> bool:
+        """Indicates whether or not a session is active."""
+        return False
+
     def execute(
         self,
         expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],

sqlmesh/core/engine_adapter/bigquery.py

@@ -87,6 +87,21 @@ def _job_params(self) -> t.Dict[str, t.Any]:
             params["maximum_bytes_billed"] = self._extra_config.get("maximum_bytes_billed")
         return params
 
+    def _begin_session(self) -> None:
+        from google.cloud.bigquery import QueryJobConfig
+
+        job = self.client.query("SELECT 1;", job_config=QueryJobConfig(create_session=True))
+        session_info = job.session_info
+        session_id = session_info.session_id if session_info else None
+        self._session_id = session_id
+        job.result()
+
+    def _end_session(self) -> None:
+        self._session_id = None
+
+    def _is_session_active(self) -> bool:
+        return self._session_id is not None
+
     def create_schema(self, schema_name: str, ignore_if_exists: bool = True) -> None:
         """Create a schema from a name or qualified table name."""
         from google.api_core.exceptions import Conflict
@@ -153,7 +168,7 @@ def fetchone(
             quote_identifiers=quote_identifiers,
         )
         try:
-            return next(self.cursor._query_data)
+            return next(self._query_data)
         except StopIteration:
             return ()
 
@@ -172,7 +187,7 @@ def fetchall(
             ignore_unsupported_errors=ignore_unsupported_errors,
             quote_identifiers=quote_identifiers,
         )
-        return list(self.cursor._query_data)
+        return list(self._query_data)
 
     def _create_table_from_df(
         self,
@@ -410,7 +425,7 @@ def _fetch_native_df(
         self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
     ) -> DF:
         self.execute(query, quote_identifiers=quote_identifiers)
-        return self.cursor._query_job.to_dataframe()
+        return self._query_job.to_dataframe()
 
     def _create_table_properties(
         self,
@@ -487,6 +502,7 @@ def execute(
     ) -> None:
         """Execute a sql query."""
         from google.cloud.bigquery import QueryJobConfig
+        from google.cloud.bigquery.query import ConnectionProperty
 
         to_sql_kwargs = (
             {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {}
@@ -503,19 +519,30 @@ def execute(
             # BigQuery's Python DB API implementation does not support retries, so we have to implement them ourselves.
             # So we update the cursor's query job and query data with the results of the new query job. This makes sure
             # that other cursor based operations execute correctly.
-            job_config = QueryJobConfig(**self._job_params)
-            self.cursor._query_job = self._db_call(
+            session_id = self._session_id
+            connection_properties = (
+                [
+                    ConnectionProperty(key="session_id", value=session_id),
+                ]
+                if session_id
+                else []
+            )
+
+            job_config = QueryJobConfig(
+                **self._job_params, connection_properties=connection_properties
+            )
+            self._query_job = self._db_call(
                 self.client.query,
                 query=sql,
                 job_config=job_config,
                 timeout=self._extra_config.get("job_creation_timeout_seconds"),
             )
             results = self._db_call(
-                self.cursor._query_job.result,
+                self._query_job.result,
                 timeout=self._extra_config.get("job_execution_timeout_seconds"),  # type: ignore
             )
-            self.cursor._query_data = iter(results) if results.total_rows else iter([])
-            query_results = self.cursor._query_job._query_results
+            self._query_data = iter(results) if results.total_rows else iter([])
+            query_results = self._query_job._query_results
             self.cursor._set_rowcount(query_results)
             self.cursor._set_description(query_results.schema)
 
@@ -541,6 +568,30 @@ def _get_data_objects(
             for table in all_tables
         ]
 
+    @property
+    def _query_data(self) -> t.Any:
+        return self._connection_pool.get_attribute("query_data")
+
+    @_query_data.setter
+    def _query_data(self, value: t.Any) -> None:
+        return self._connection_pool.set_attribute("query_data", value)
+
+    @property
+    def _query_job(self) -> t.Any:
+        return self._connection_pool.get_attribute("query_job")
+
+    @_query_job.setter
+    def _query_job(self, value: t.Any) -> None:
+        return self._connection_pool.set_attribute("query_job", value)
+
+    @property
+    def _session_id(self) -> t.Any:
+        return self._connection_pool.get_attribute("session_id")
+
+    @_session_id.setter
+    def _session_id(self, value: t.Any) -> None:
+        return self._connection_pool.set_attribute("session_id", value)
+
 
 class _ErrorCounter:
     """

sqlmesh/core/snapshot/evaluator.py

@@ -145,7 +145,7 @@ def apply(query_or_df: QueryOrDF, index: int = 0) -> None:
             transaction_type=TransactionType.DDL
             if model.kind.is_view or model.kind.is_full
             else TransactionType.DML
-        ):
+        ), self.adapter.session():
             if not limit:
                 self.adapter.execute(model.render_pre_statements(**render_statements_kwargs))
 

sqlmesh/utils/connection_pool.py

@@ -1,6 +1,7 @@
 import abc
 import logging
 import typing as t
+from collections import defaultdict
 from threading import Lock, get_ident
 
 logger = logging.getLogger(__name__)
@@ -27,6 +28,26 @@ def get(self) -> t.Any:
             A connection instance.
         """
 
+    @abc.abstractmethod
+    def get_attribute(self, key: str) -> t.Optional[t.Any]:
+        """Returns an attribute associated with the connection.
+
+        Args:
+            key: Attribute key.
+
+        Returns:
+            Attribute value or None if not found.
+        """
+
+    @abc.abstractmethod
+    def set_attribute(self, key: str, value: t.Any) -> None:
+        """Sets an attribute associated with the connection.
+
+        Args:
+            key: Attribute key.
+            value: Attribute value.
+        """
+
     @abc.abstractmethod
     def begin(self) -> None:
         """Starts a new transaction."""
@@ -96,6 +117,7 @@ def __init__(self, connection_factory: t.Callable[[], t.Any]):
         self._thread_connections: t.Dict[t.Hashable, t.Any] = {}
         self._thread_cursors: t.Dict[t.Hashable, t.Any] = {}
         self._thread_transactions: t.Set[t.Hashable] = set()
+        self._thread_attributes: t.Dict[t.Hashable, t.Dict[str, t.Any]] = defaultdict(dict)
         self._thread_connections_lock = Lock()
         self._thread_cursors_lock = Lock()
         self._thread_transactions_lock = Lock()
@@ -114,6 +136,14 @@ def get(self) -> t.Any:
                 self._thread_connections[thread_id] = self._connection_factory()
             return self._thread_connections[thread_id]
 
+    def get_attribute(self, key: str) -> t.Optional[t.Any]:
+        thread_id = get_ident()
+        return self._thread_attributes[thread_id].get(key)
+
+    def set_attribute(self, key: str, value: t.Any) -> None:
+        thread_id = get_ident()
+        self._thread_attributes[thread_id][key] = value
+
     def begin(self) -> None:
         self._do_begin()
         with self._thread_transactions_lock:
@@ -147,6 +177,7 @@ def close(self) -> None:
                 self._thread_connections.pop(thread_id)
                 self._thread_cursors.pop(thread_id, None)
                 self._discard_transaction(thread_id)
+            self._thread_attributes.pop(thread_id, None)
 
     def close_all(self, exclude_calling_thread: bool = False) -> None:
         calling_thread_id = get_ident()
@@ -158,6 +189,7 @@ def close_all(self, exclude_calling_thread: bool = False) -> None:
                     self._thread_connections.pop(thread_id)
                     self._thread_cursors.pop(thread_id, None)
                     self._discard_transaction(thread_id)
+                self._thread_attributes.pop(thread_id, None)
 
     def _discard_transaction(self, thread_id: t.Hashable) -> None:
         with self._thread_transactions_lock:
@@ -169,6 +201,7 @@ def __init__(self, connection_factory: t.Callable[[], t.Any]):
         self._connection_factory = connection_factory
         self._connection: t.Optional[t.Any] = None
         self._cursor: t.Optional[t.Any] = None
+        self._attributes: t.Dict[str, t.Any] = {}
         self._is_transaction_active: bool = False
 
     def get_cursor(self) -> t.Any:
@@ -181,6 +214,12 @@ def get(self) -> t.Any:
             self._connection = self._connection_factory()
         return self._connection
 
+    def get_attribute(self, key: str) -> t.Optional[t.Any]:
+        return self._attributes.get(key)
+
+    def set_attribute(self, key: str, value: t.Any) -> None:
+        self._attributes[key] = value
+
     def begin(self) -> None:
         self._do_begin()
         self._is_transaction_active = True
@@ -206,6 +245,7 @@ def close(self) -> None:
         self._connection = None
         self._cursor = None
         self._is_transaction_active = False
+        self._attributes.clear()
 
     def close_all(self, exclude_calling_thread: bool = False) -> None:
         if not exclude_calling_thread:

tests/core/engine_adapter/test_bigquery.py

@@ -364,3 +364,38 @@ def test_merge(mocker: MockerFixture):
         "job_config",
         "table",
     ]
+
+
+def test_begin_end_session(mocker: MockerFixture):
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    cursor_mock.connection = connection_mock
+    connection_mock.cursor.return_value = cursor_mock
+
+    query_result_mock = mocker.Mock()
+    query_result_mock.total_rows = 0
+    job_mock = mocker.Mock()
+    job_mock.result.return_value = query_result_mock
+    connection_mock._client.query.return_value = job_mock
+
+    adapter = BigQueryEngineAdapter(lambda: connection_mock, job_retries=0)
+
+    with adapter.session():
+        assert adapter._connection_pool.get_attribute("session_id") is not None
+        adapter.execute("SELECT 2;")
+
+    assert adapter._connection_pool.get_attribute("session_id") is None
+    adapter.execute("SELECT 3;")
+
+    begin_session_call = connection_mock._client.query.call_args_list[0]
+    assert begin_session_call[0][0] == "SELECT 1;"
+
+    execute_a_call = connection_mock._client.query.call_args_list[1]
+    assert execute_a_call[1]["query"] == "SELECT 2;"
+    assert len(execute_a_call[1]["job_config"].connection_properties) == 1
+    assert execute_a_call[1]["job_config"].connection_properties[0].key == "session_id"
+    assert execute_a_call[1]["job_config"].connection_properties[0].value
+
+    execute_b_call = connection_mock._client.query.call_args_list[2]
+    assert execute_b_call[1]["query"] == "SELECT 3;"
+    assert not execute_b_call[1]["job_config"].connection_properties

tests/core/test_snapshot_evaluator.py

@@ -62,8 +62,13 @@ def adapter_mock(mocker: MockerFixture):
     transaction_mock.__enter__ = mocker.Mock()
     transaction_mock.__exit__ = mocker.Mock()
 
+    session_mock = mocker.Mock()
+    session_mock.__enter__ = mocker.Mock()
+    session_mock.__exit__ = mocker.Mock()
+
     adapter_mock = mocker.Mock()
     adapter_mock.transaction.return_value = transaction_mock
+    adapter_mock.session.return_value = session_mock
     adapter_mock.dialect = "duckdb"
     return adapter_mock