Fix: Properly support engines that can share a single connection instance across threads (#4124)

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
@@ -51,6 +51,9 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     pre_ping: bool
     pretty_sql: bool = False
 
+    # Whether to share a  single connection across threads or create a new connection per thread.
+    shared_connection: t.ClassVar[bool] = False
+
     @property
     @abc.abstractmethod
     def _connection_kwargs_keys(self) -> t.Set[str]:
@@ -116,6 +119,7 @@ def create_engine_adapter(self, register_comments_override: bool = False) -> Eng
             register_comments=register_comments_override or self.register_comments,
             pre_ping=self.pre_ping,
             pretty_sql=self.pretty_sql,
+            shared_connection=self.shared_connection,
             **self._extra_engine_config,
         )
 
@@ -182,6 +186,8 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
 
     token: t.Optional[str] = None
 
+    shared_connection: t.ClassVar[bool] = True
+
     _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
 
     @model_validator(mode="before")
@@ -212,43 +218,6 @@ def _connection_kwargs_keys(self) -> t.Set[str]:
     def _connection_factory(self) -> t.Callable:
         import duckdb
 
-        if self.concurrent_tasks > 1:
-            # ensures a single connection instance is used across threads rather than a new connection being established per thread
-            # this is in line with https://duckdb.org/docs/guides/python/multiple_threads.html
-            # the important thing is that the *cursor*'s are per thread, but the connection should be shared
-            @lru_cache
-            def _factory(*args: t.Any, **kwargs: t.Any) -> t.Any:
-                class ConnWrapper:
-                    def __init__(self, conn: duckdb.DuckDBPyConnection):
-                        self.conn = conn
-
-                    def __getattr__(self, attr: str) -> t.Any:
-                        return getattr(self.conn, attr)
-
-                    def close(self) -> None:
-                        # This overrides conn.close() to be a no-op to work with ThreadLocalConnectionPool which assumes that a new connection should
-                        # be created per thread. However, DuckDB expects the same connection instance to be shared across threads. There is a pattern
-                        # in the SQLMesh codebase that `EngineAdapter.recycle()` is called after doing things like merging intervals. This in turn causes
-                        # `ThreadLocalConnectionPool.close_all(exclude_calling_thread=True)` to be called.
-                        #
-                        # The problem with sharing a connection across threads and then allowing it to be closed for every thread except the current one
-                        # is that it gets closed for the current one too because its shared. This causes any ":memory:" databases to be discarded.
-                        # ":memory:" databases are convienient and are used heavily in our test suite amongst other things.
-                        #
-                        # Ok, so why not have a connection per thread as is the default for ThreadLocalConnectionPool? Two reasons:
-                        # - It makes any ":memory:" databases unique to that thread. So if one thread creates tables, another thread cant see them
-                        # - If you use local files instead (eg point each connection to the same db file) then all the connection instances
-                        #   fight over locks to the same file and performance tanks heavily
-                        #
-                        # From what I can tell, DuckDB expects the single process reading / writing the database from multiple
-                        # threads to /share the same connection/ and just use thread-local cursors. In order to support ":memory:" databases
-                        # and remove lock contention, the connection needs to live for the life of the application and not be closed
-                        pass
-
-                return ConnWrapper(duckdb.connect(*args, **kwargs))
-
-            return _factory
-
         return duckdb.connect
 
     @property

sqlmesh/core/engine_adapter/base.py

@@ -119,14 +119,18 @@ def __init__(
         register_comments: bool = True,
         pre_ping: bool = False,
         pretty_sql: bool = False,
+        shared_connection: bool = False,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
         self._connection_pool = (
             connection_factory_or_pool
             if isinstance(connection_factory_or_pool, ConnectionPool)
             else create_connection_pool(
-                connection_factory_or_pool, multithreaded, cursor_init=cursor_init
+                connection_factory_or_pool,
+                multithreaded,
+                shared_connection=shared_connection,
+                cursor_init=cursor_init,
             )
         )
         self._sql_gen_kwargs = sql_gen_kwargs or {}

sqlmesh/utils/connection_pool.py

@@ -111,18 +111,16 @@ def _do_rollback(self) -> None:
             self.get().rollback()
 
 
-class ThreadLocalConnectionPool(_TransactionManagementMixin):
+class _ThreadLocalBase(_TransactionManagementMixin):
     def __init__(
         self,
         connection_factory: t.Callable[[], t.Any],
         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._cursor_init = cursor_init
@@ -136,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)
@@ -176,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:
@@ -191,16 +204,51 @@ def close_all(self, exclude_calling_thread: bool = False) -> None:
         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):
@@ -269,13 +317,17 @@ def close_all(self, exclude_calling_thread: bool = False) -> None:
 def create_connection_pool(
     connection_factory: t.Callable[[], t.Any],
     multithreaded: bool,
+    shared_connection: bool = False,
     cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
 ) -> ConnectionPool:
-    return (
-        ThreadLocalConnectionPool(connection_factory, 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/engine_adapter/integration/test_integration_duckdb.py

@@ -5,7 +5,7 @@
 from sqlglot import exp
 
 from sqlmesh.core.config.connection import DuckDBConnectionConfig
-from sqlmesh.utils.connection_pool import ThreadLocalConnectionPool
+from sqlmesh.utils.connection_pool import ThreadLocalSharedConnectionPool
 
 pytestmark = [pytest.mark.duckdb, pytest.mark.engine, pytest.mark.slow]
 
@@ -21,7 +21,7 @@ def test_multithread_concurrency(tmp_path, database: t.Optional[str]):
 
     adapter = config.create_engine_adapter()
 
-    assert isinstance(adapter._connection_pool, ThreadLocalConnectionPool)
+    assert isinstance(adapter._connection_pool, ThreadLocalSharedConnectionPool)
 
     # this test loosely follows this example: https://duckdb.org/docs/guides/python/multiple_threads.html
     adapter.execute(

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/core/test_context.py

@@ -39,7 +39,7 @@
 from sqlmesh.core.plan import BuiltInPlanEvaluator, PlanBuilder
 from sqlmesh.core.state_sync.cache import CachingStateSync
 from sqlmesh.core.state_sync.db import EngineAdapterStateSync
-from sqlmesh.utils.connection_pool import SingletonConnectionPool, ThreadLocalConnectionPool
+from sqlmesh.utils.connection_pool import SingletonConnectionPool, ThreadLocalSharedConnectionPool
 from sqlmesh.utils.date import (
     make_inclusive_end,
     now,
@@ -1209,7 +1209,7 @@ def test_duckdb_state_connection_automatic_multithreaded_mode(tmp_path):
     state_sync = context.state_sync.state_sync
     assert isinstance(state_sync, EngineAdapterStateSync)
     assert isinstance(state_sync.engine_adapter, DuckDBEngineAdapter)
-    assert isinstance(state_sync.engine_adapter._connection_pool, ThreadLocalConnectionPool)
+    assert isinstance(state_sync.engine_adapter._connection_pool, ThreadLocalSharedConnectionPool)
 
 
 def test_requirements(copy_to_temp_path: t.Callable):

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