fix: init catalogs at cursor level (#1630)

Author: eakmanrq

examples/sushi/config.py

@@ -113,3 +113,15 @@
     model_defaults=ModelDefaultsConfig(dialect="duckdb"),
     environment_suffix_target=EnvironmentSuffixTarget.TABLE,
 )
+
+
+CATALOGS = {
+    "in_memory": ":memory:",
+    "other_catalog": f":memory:",
+}
+
+local_catalogs = Config(
+    default_connection=DuckDBConnectionConfig(catalogs=CATALOGS),
+    default_test_connection=DuckDBConnectionConfig(catalogs=CATALOGS),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)

sqlmesh/core/config/connection.py

@@ -62,6 +62,11 @@ def _cursor_kwargs(self) -> t.Optional[t.Dict[str, t.Any]]:
         """Key-value arguments that will be passed during cursor construction."""
         return None
 
+    @property
+    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
+        """A function that is called to initialize the cursor"""
+        return None
+
     def create_engine_adapter(self) -> EngineAdapter:
         """Returns a new instance of the Engine Adapter."""
         return self._engine_adapter(
@@ -73,6 +78,7 @@ def create_engine_adapter(self) -> EngineAdapter:
             ),
             multithreaded=self.concurrent_tasks > 1,
             cursor_kwargs=self._cursor_kwargs,
+            cursor_init=self._cursor_init,
             **self._extra_engine_config,
         )
 
@@ -118,25 +124,29 @@ def _connection_factory(self) -> t.Callable:
 
         return duckdb.connect
 
-    def create_engine_adapter(self) -> EngineAdapter:
-        """Returns a new instance of the Engine Adapter."""
+    @property
+    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
+        """A function that is called to initialize the cursor"""
+        import duckdb
         from duckdb import BinderException
 
-        engine_adapter = super().create_engine_adapter()
-        for i, (alias, path) in enumerate((self.catalogs or {}).items()):
-            try:
-                engine_adapter.execute(f"ATTACH '{path}' AS {alias}")
-            except BinderException as e:
-                # If a user tries to create a catalog pointing at `:memory:` and with the name `memory`
-                # then we don't want to raise since this happens by default. They are just doing this to
-                # set it as the default catalog.
-                if not (
-                    'database with name "memory" already exists' in str(e) and path == ":memory:"
-                ):
-                    raise e
-            if i == 0 and not self.database:
-                engine_adapter.set_current_catalog(alias)
-        return engine_adapter
+        def init_catalogs(cursor: duckdb.DuckDBPyConnection) -> None:
+            for i, (alias, path) in enumerate((self.catalogs or {}).items()):
+                try:
+                    cursor.execute(f"ATTACH '{path}' AS {alias}")
+                except BinderException as e:
+                    # If a user tries to create a catalog pointing at `:memory:` and with the name `memory`
+                    # then we don't want to raise since this happens by default. They are just doing this to
+                    # set it as the default catalog.
+                    if not (
+                        'database with name "memory" already exists' in str(e)
+                        and path == ":memory:"
+                    ):
+                        raise e
+                if i == 0 and not self.database:
+                    cursor.execute(f"USE {alias}")
+
+        return init_catalogs
 
 
 class SnowflakeConnectionConfig(ConnectionConfig):

sqlmesh/core/engine_adapter/base.py

@@ -164,11 +164,12 @@ def __init__(
         sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None,
         multithreaded: bool = False,
         cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
         self._connection_pool = create_connection_pool(
-            connection_factory, multithreaded, cursor_kwargs=cursor_kwargs
+            connection_factory, multithreaded, cursor_kwargs=cursor_kwargs, cursor_init=cursor_init
         )
         self.sql_gen_kwargs = sql_gen_kwargs or {}
         self._extra_config = kwargs

sqlmesh/core/engine_adapter/databricks.py

@@ -41,6 +41,7 @@ def __init__(
         sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None,
         multithreaded: bool = False,
         cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
         **kwargs: t.Any,
     ):
         super().__init__(
@@ -49,6 +50,7 @@ def __init__(
             sql_gen_kwargs,
             multithreaded,
             cursor_kwargs,
+            cursor_init,
             **kwargs,
         )
         self._spark: t.Optional[PySparkSession] = None

sqlmesh/utils/connection_pool.py

@@ -116,6 +116,7 @@ def __init__(
         self,
         connection_factory: t.Callable[[], t.Any],
         cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
     ):
         self._connection_factory = connection_factory
         self._thread_connections: t.Dict[t.Hashable, t.Any] = {}
@@ -126,12 +127,15 @@ def __init__(
         self._thread_cursors_lock = Lock()
         self._thread_transactions_lock = Lock()
         self._cursor_kwargs = cursor_kwargs or {}
+        self._cursor_init = cursor_init
 
     def get_cursor(self) -> t.Any:
         thread_id = get_ident()
         with self._thread_cursors_lock:
             if thread_id not in self._thread_cursors:
                 self._thread_cursors[thread_id] = self.get().cursor(**self._cursor_kwargs)
+                if self._cursor_init:
+                    self._cursor_init(self._thread_cursors[thread_id])
             return self._thread_cursors[thread_id]
 
     def get(self) -> t.Any:
@@ -206,17 +210,21 @@ def __init__(
         self,
         connection_factory: t.Callable[[], t.Any],
         cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
     ):
         self._connection_factory = connection_factory
         self._connection: t.Optional[t.Any] = None
         self._cursor: t.Optional[t.Any] = None
         self._cursor_kwargs = cursor_kwargs or {}
         self._attributes: t.Dict[str, t.Any] = {}
         self._is_transaction_active: bool = False
+        self._cursor_init = cursor_init
 
     def get_cursor(self) -> t.Any:
         if not self._cursor:
             self._cursor = self.get().cursor(**self._cursor_kwargs)
+            if self._cursor_init:
+                self._cursor_init(self._cursor)
         return self._cursor
 
     def get(self) -> t.Any:
@@ -266,11 +274,16 @@ def create_connection_pool(
     connection_factory: t.Callable[[], t.Any],
     multithreaded: bool,
     cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+    cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
 ) -> ConnectionPool:
     return (
-        ThreadLocalConnectionPool(connection_factory, cursor_kwargs=cursor_kwargs)
+        ThreadLocalConnectionPool(
+            connection_factory, cursor_kwargs=cursor_kwargs, cursor_init=cursor_init
+        )
         if multithreaded
-        else SingletonConnectionPool(connection_factory, cursor_kwargs=cursor_kwargs)
+        else SingletonConnectionPool(
+            connection_factory, cursor_kwargs=cursor_kwargs, cursor_init=cursor_init
+        )
     )
 
 

tests/conftest.py

@@ -178,6 +178,13 @@ def sushi_test_dbt_context(mocker: MockerFixture) -> Context:
     return context
 
 
+@pytest.fixture()
+def sushi_default_catalog(mocker: MockerFixture) -> Context:
+    context, plan = init_and_plan_context("examples/sushi", mocker, "local_catalogs")
+    context.apply(plan)
+    return context
+
+
 def init_and_plan_context(
     paths: str | t.List[str],
     mocker: MockerFixture,

tests/core/test_integration.py

@@ -418,7 +418,8 @@ def test_plan_set_choice_is_reflected_in_missing_intervals(mocker: MockerFixture
 @pytest.mark.integration
 @pytest.mark.core_integration
 @pytest.mark.parametrize(
-    "context_fixture", ["sushi_context", "sushi_dbt_context", "sushi_test_dbt_context"]
+    "context_fixture",
+    ["sushi_context", "sushi_dbt_context", "sushi_test_dbt_context", "sushi_default_catalog"],
 )
 def test_model_add(context_fixture: Context, request):
     initial_add(request.getfixturevalue(context_fixture), "dev")