Fix: Properly support engines that can share a single connection instance across threads

Author: izeigerman

sqlmesh/core/config/connection.py

@@ -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]:
@@ -94,13 +97,21 @@ 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"""
-        return partial(
+        factory = 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"""
@@ -116,6 +127,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 +194,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 +226,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

@@ -115,6 +115,7 @@ class ThreadLocalConnectionPool(_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
@@ -125,6 +126,7 @@ def __init__(
         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:
@@ -187,6 +189,9 @@ 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():
@@ -269,10 +274,13 @@ 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)
+        ThreadLocalConnectionPool(
+            connection_factory, shared_connection=shared_connection, cursor_init=cursor_init
+        )
         if multithreaded
         else SingletonConnectionPool(connection_factory, cursor_init=cursor_init)
     )