Refactor to pass as argument to create _engine_adapter

Author: themisvaltinos

sqlmesh/core/config/connection.py

@@ -109,8 +109,12 @@ 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."""
+        if concurrent_tasks:
+            self.concurrent_tasks = concurrent_tasks
         return self._engine_adapter(
             self._connection_factory_with_kwargs,
             multithreaded=self.concurrent_tasks > 1,
@@ -288,7 +292,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."""
@@ -319,7 +325,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/config/scheduler.py

@@ -90,10 +90,10 @@ def create_state_sync(self, context: GenericContext) -> StateSync:
                     + f"multi threaded mode with {warehouse_connection.concurrent_tasks} concurrent tasks."
                     + " This can cause SQLMesh to hang. Overriding the duckdb state connection config to use multi threaded mode."
                 )
-                # this triggers multithreaded mode and has to happen before the engine adapter is created below
-                state_connection.concurrent_tasks = warehouse_connection.concurrent_tasks
 
-        engine_adapter = state_connection.create_engine_adapter()
+        engine_adapter = state_connection.create_engine_adapter(
+            concurrent_tasks=warehouse_connection.concurrent_tasks
+        )
         if state_connection.is_forbidden_for_state_sync:
             raise ConfigError(
                 f"The {engine_adapter.DIALECT.upper()} engine cannot be used to store SQLMesh state - please specify a different `state_connection` engine."

sqlmesh/core/context.py

@@ -2197,8 +2197,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)
-                connection.concurrent_tasks = self.concurrent_tasks
-                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