Introduce a thread local pool with a shared connection

Author: izeigerman

sqlmesh/core/config/connection.py

@@ -8,7 +8,7 @@
 import re
 import typing as t
 from enum import Enum
-from functools import partial, lru_cache
+from functools import partial
 
 import pydantic
 from pydantic import Field
@@ -97,21 +97,13 @@ def is_forbidden_for_state_sync(self) -> bool:
     @property
     def _connection_factory_with_kwargs(self) -> t.Callable[[], t.Any]:
         """A function that is called to return a connection object for the given Engine Adapter"""
-        factory = partial(
+        return partial(
             self._connection_factory,
             **{
                 **self._static_connection_kwargs,
                 **{k: v for k, v in self.dict().items() if k in self._connection_kwargs_keys},
             },
         )
-        if self.shared_connection:
-            # Make sure that a single connection is created and returned
-            @lru_cache
-            def _cached_connection() -> t.Any:
-                return factory()
-
-            return _cached_connection
-        return factory
 
     def connection_validator(self) -> t.Callable[[], None]:
         """A function that validates the connection configuration"""

sqlmesh/utils/connection_pool.py

@@ -111,22 +111,18 @@ def _do_rollback(self) -> None:
             self.get().rollback()
 
 
-class ThreadLocalConnectionPool(_TransactionManagementMixin):
+class _ThreadLocalBase(_TransactionManagementMixin):
     def __init__(
         self,
         connection_factory: t.Callable[[], t.Any],
-        shared_connection: bool = False,
         cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
     ):
         self._connection_factory = connection_factory
-        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()
-        self._shared_connection = shared_connection
         self._cursor_init = cursor_init
 
     def get_cursor(self) -> t.Any:
@@ -138,13 +134,6 @@ def get_cursor(self) -> t.Any:
                     self._cursor_init(self._thread_cursors[thread_id])
             return self._thread_cursors[thread_id]
 
-    def get(self) -> t.Any:
-        thread_id = get_ident()
-        with self._thread_connections_lock:
-            if thread_id not in self._thread_connections:
-                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)
@@ -178,6 +167,28 @@ def close_cursor(self) -> None:
                 _try_close(self._thread_cursors[thread_id], "cursor")
                 self._thread_cursors.pop(thread_id)
 
+    def _discard_transaction(self, thread_id: t.Hashable) -> None:
+        with self._thread_transactions_lock:
+            self._thread_transactions.discard(thread_id)
+
+
+class ThreadLocalConnectionPool(_ThreadLocalBase):
+    def __init__(
+        self,
+        connection_factory: t.Callable[[], t.Any],
+        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
+    ):
+        super().__init__(connection_factory, cursor_init)
+        self._thread_connections: t.Dict[t.Hashable, t.Any] = {}
+        self._thread_connections_lock = Lock()
+
+    def get(self) -> t.Any:
+        thread_id = get_ident()
+        with self._thread_connections_lock:
+            if thread_id not in self._thread_connections:
+                self._thread_connections[thread_id] = self._connection_factory()
+            return self._thread_connections[thread_id]
+
     def close(self) -> None:
         thread_id = get_ident()
         with self._thread_cursors_lock, self._thread_connections_lock:
@@ -189,23 +200,55 @@ def close(self) -> None:
             self._thread_attributes.pop(thread_id, None)
 
     def close_all(self, exclude_calling_thread: bool = False) -> None:
-        if exclude_calling_thread and self._shared_connection:
-            return
-
         calling_thread_id = get_ident()
         with self._thread_cursors_lock, self._thread_connections_lock:
             for thread_id, connection in self._thread_connections.copy().items():
                 if not exclude_calling_thread or thread_id != calling_thread_id:
-                    # NOTE: the access to the connection instance itself is not thread-safe here.
                     _try_close(connection, "connection")
                     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:
-            self._thread_transactions.discard(thread_id)
+
+class ThreadLocalSharedConnectionPool(_ThreadLocalBase):
+    def __init__(
+        self,
+        connection_factory: t.Callable[[], t.Any],
+        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
+    ):
+        super().__init__(connection_factory, cursor_init)
+        self._connection: t.Optional[t.Any] = None
+        self._connection_lock = Lock()
+
+    def get(self) -> t.Any:
+        with self._connection_lock:
+            if self._connection is None:
+                self._connection = self._connection_factory()
+            return self._connection
+
+    def close(self) -> None:
+        thread_id = get_ident()
+        with self._thread_cursors_lock, self._connection_lock:
+            if thread_id in self._thread_cursors:
+                _try_close(self._thread_cursors[thread_id], "cursor")
+                self._thread_cursors.pop(thread_id)
+            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()
+        with self._thread_cursors_lock, self._connection_lock:
+            for thread_id, cursor in self._thread_cursors.copy().items():
+                if not exclude_calling_thread or thread_id != calling_thread_id:
+                    _try_close(cursor, "cursor")
+                    self._thread_cursors.pop(thread_id)
+                    self._discard_transaction(thread_id)
+                self._thread_attributes.pop(thread_id, None)
+
+            if not exclude_calling_thread:
+                _try_close(self._connection, "connection")
+                self._connection = None
 
 
 class SingletonConnectionPool(_TransactionManagementMixin):
@@ -277,13 +320,14 @@ def create_connection_pool(
     shared_connection: bool = False,
     cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
 ) -> ConnectionPool:
-    return (
-        ThreadLocalConnectionPool(
-            connection_factory, shared_connection=shared_connection, cursor_init=cursor_init
-        )
+    pool_class = (
+        ThreadLocalSharedConnectionPool
+        if multithreaded and shared_connection
+        else ThreadLocalConnectionPool
         if multithreaded
-        else SingletonConnectionPool(connection_factory, cursor_init=cursor_init)
+        else SingletonConnectionPool
     )
+    return pool_class(connection_factory, cursor_init=cursor_init)
 
 
 def _try_close(closeable: t.Any, kind: str) -> None:

tests/core/test_connection_config.py

@@ -607,7 +607,7 @@ def test_duckdb_attach_options():
 
 def test_duckdb_multithreaded_connection_factory(make_config):
     from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
-    from sqlmesh.utils.connection_pool import ThreadLocalConnectionPool
+    from sqlmesh.utils.connection_pool import ThreadLocalSharedConnectionPool
     from threading import Thread
 
     config = make_config(type="duckdb")
@@ -620,7 +620,7 @@ def test_duckdb_multithreaded_connection_factory(make_config):
     config = make_config(type="duckdb", concurrent_tasks=8)
     adapter = config.create_engine_adapter()
     assert isinstance(adapter, DuckDBEngineAdapter)
-    assert isinstance(adapter._connection_pool, ThreadLocalConnectionPool)
+    assert isinstance(adapter._connection_pool, ThreadLocalSharedConnectionPool)
 
     threads = []
     connection_objects = []

tests/utils/test_connection_pool.py

@@ -6,6 +6,7 @@
 from sqlmesh.utils.connection_pool import (
     SingletonConnectionPool,
     ThreadLocalConnectionPool,
+    ThreadLocalSharedConnectionPool,
 )
 
 
@@ -207,3 +208,95 @@ def thread():
     assert cursor_mock_thread_one.rollback.call_count == 1
 
     assert cursor_mock_thread_two.begin.call_count == 1
+
+
+def test_thread_local_shared_connection_pool(mocker: MockerFixture):
+    cursor_mock_thread_one = mocker.Mock()
+    cursor_mock_thread_two = mocker.Mock()
+    connection_mock = mocker.Mock()
+    connection_mock.cursor.side_effect = [
+        cursor_mock_thread_one,
+        cursor_mock_thread_two,
+        cursor_mock_thread_one,
+    ]
+
+    test_thread_id = get_ident()
+
+    connection_factory_mock = mocker.Mock(return_value=connection_mock)
+    pool = ThreadLocalSharedConnectionPool(connection_factory_mock)
+
+    assert pool.get_cursor() == cursor_mock_thread_one
+    assert pool.get_cursor() == cursor_mock_thread_one
+    assert pool.get() == connection_mock
+    assert pool.get() == connection_mock
+
+    def thread():
+        assert pool.get_cursor() == cursor_mock_thread_two
+        assert pool.get_cursor() == cursor_mock_thread_two
+        assert pool.get() == connection_mock
+        assert pool.get() == connection_mock
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(thread).result()
+
+    assert pool._connection is not None
+    assert len(pool._thread_cursors) == 2
+
+    pool.close_all(exclude_calling_thread=True)
+
+    assert pool._connection is not None
+    assert len(pool._thread_cursors) == 1
+    assert test_thread_id in pool._thread_cursors
+
+    pool.close_cursor()
+    pool.close()
+
+    assert pool.get_cursor() == cursor_mock_thread_one
+
+    pool.close_all()
+
+    assert connection_factory_mock.call_count == 1
+
+    assert cursor_mock_thread_one.close.call_count == 2
+    assert connection_mock.cursor.call_count == 3
+    assert connection_mock.close.call_count == 1
+
+
+def test_thread_local_shared_connection_pool_close(mocker: MockerFixture):
+    connection_mock = mocker.Mock()
+    cursor_mock = mocker.Mock()
+    connection_mock.cursor.return_value = cursor_mock
+
+    connection_factory_mock = mocker.Mock(return_value=connection_mock)
+    pool = ThreadLocalSharedConnectionPool(connection_factory_mock)
+
+    # First time we get a connection
+    pool.get()
+    pool.get()
+    pool.get_cursor()
+    pool.get_cursor()
+
+    # This shouldn't close the connection, only the cursor
+    pool.close()
+    pool.get()
+    pool.get()
+    pool.get_cursor()
+
+    pool.get_cursor()
+    # This shouldn't close the connection either
+    pool.close_all(exclude_calling_thread=True)
+
+    pool.get()
+    pool.get()
+    # Now this should close the connection
+    pool.close_all()
+
+    # Re-open the connection
+    pool.get()
+    pool.get()
+    # Close it again
+    pool.close_all()
+
+    assert cursor_mock.close.call_count == 2
+    assert connection_factory_mock.call_count == 2
+    assert connection_mock.close.call_count == 2