Fix: When initialising multiple connections pass concurrent tasks (#4176)

Author: themisvaltinos

sqlmesh/core/config/connection.py

@@ -109,11 +109,15 @@ def connection_validator(self) -> t.Callable[[], None]:
         """A function that validates the connection configuration"""
         return self.create_engine_adapter().ping
 
-    def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter:
+    def create_engine_adapter(
+        self, register_comments_override: bool = False, concurrent_tasks: t.Optional[int] = None
+    ) -> EngineAdapter:
         """Returns a new instance of the Engine Adapter."""
+
+        concurrent_tasks = concurrent_tasks or self.concurrent_tasks
         return self._engine_adapter(
             self._connection_factory_with_kwargs,
-            multithreaded=self.concurrent_tasks > 1,
+            multithreaded=concurrent_tasks > 1,
             default_catalog=self.get_catalog(),
             cursor_init=self._cursor_init,
             register_comments=register_comments_override or self.register_comments,
@@ -284,7 +288,9 @@ def init(cursor: duckdb.DuckDBPyConnection) -> None:
 
         return init
 
-    def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter:
+    def create_engine_adapter(
+        self, register_comments_override: bool = False, concurrent_tasks: t.Optional[int] = None
+    ) -> EngineAdapter:
         """Checks if another engine adapter has already been created that shares a catalog that points to the same data
         file. If so, it uses that same adapter instead of creating a new one. As a result, any additional configuration
         associated with the new adapter will be ignored."""
@@ -315,7 +321,9 @@ def create_engine_adapter(self, register_comments_override: bool = False) -> Eng
             logger.info(f"Creating new DuckDB adapter for data files: {masked_files}")
         else:
             logger.info("Creating new DuckDB adapter for in-memory database")
-        adapter = super().create_engine_adapter(register_comments_override)
+        adapter = super().create_engine_adapter(
+            register_comments_override, concurrent_tasks=concurrent_tasks
+        )
         for data_file in data_files:
             key = data_file if isinstance(data_file, str) else data_file.path
             BaseDuckDBConnectionConfig._data_file_to_adapter[key] = adapter

sqlmesh/core/context.py

@@ -2254,7 +2254,7 @@ def engine_adapters(self) -> t.Dict[str, EngineAdapter]:
         for gateway_name in self.config.gateways:
             if gateway_name != self.selected_gateway:
                 connection = self.config.get_connection(gateway_name)
-                adapter = connection.create_engine_adapter()
+                adapter = connection.create_engine_adapter(concurrent_tasks=self.concurrent_tasks)
                 self._engine_adapters[gateway_name] = adapter
         return self._engine_adapters
 

tests/core/test_config.py

@@ -26,6 +26,7 @@
 )
 from sqlmesh.core.context import Context
 from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter
+from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
 from sqlmesh.core.engine_adapter.redshift import RedshiftEngineAdapter
 from sqlmesh.core.notification_target import ConsoleNotificationTarget
 from sqlmesh.core.user import User
@@ -709,6 +710,10 @@ def test_multi_gateway_config(tmp_path, mocker: MockerFixture):
             aws_secret_access_key: accesskey
             work_group: group
             s3_warehouse_location: s3://location
+    duckdb:
+        connection:
+            type: duckdb
+            database: db.db
 
 default_gateway: redshift
 
@@ -725,11 +730,53 @@ def test_multi_gateway_config(tmp_path, mocker: MockerFixture):
     ctx = Context(paths=tmp_path, config=config)
 
     assert isinstance(ctx._connection_config, RedshiftConnectionConfig)
-    assert len(ctx.engine_adapters) == 2
+    assert len(ctx.engine_adapters) == 3
     assert isinstance(ctx.engine_adapters["athena"], AthenaEngineAdapter)
     assert isinstance(ctx.engine_adapters["redshift"], RedshiftEngineAdapter)
+    assert isinstance(ctx.engine_adapters["duckdb"], DuckDBEngineAdapter)
     assert ctx.engine_adapter == ctx._get_engine_adapter("redshift")
 
+    # The duckdb engine adapter should be have been set as multithreaded as well
+    assert ctx.engine_adapters["duckdb"]._multithreaded
+
+
+def test_multi_gateway_single_threaded_config(tmp_path):
+    config_path = tmp_path / "config_duck_athena.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+gateways:
+    duckdb:
+        connection:
+            type: duckdb
+            database: db.db
+    athena:
+        connection:
+            type: athena
+            aws_access_key_id: '1234'
+            aws_secret_access_key: accesskey
+            work_group: group
+            s3_warehouse_location: s3://location
+default_gateway: duckdb
+model_defaults:
+    dialect: duckdb
+        """
+        )
+
+    config = load_config_from_paths(
+        Config,
+        project_paths=[config_path],
+    )
+
+    ctx = Context(paths=tmp_path, config=config)
+    assert isinstance(ctx._connection_config, DuckDBConnectionConfig)
+    assert len(ctx.engine_adapters) == 2
+    assert ctx.engine_adapter == ctx._get_engine_adapter("duckdb")
+    assert isinstance(ctx.engine_adapters["athena"], AthenaEngineAdapter)
+
+    # In this case athena should use 1 concurrent task as the default gateway is duckdb
+    assert not ctx.engine_adapters["athena"]._multithreaded
+
 
 def test_trino_schema_location_mapping_syntax(tmp_path):
     config_path = tmp_path / "config_trino.yaml"