Refactor: Move engine adapter instantiation after loading on-demand (#4578)

Author: themisvaltinos

sqlmesh/core/config/scheduler.py

@@ -47,8 +47,8 @@ def create_state_sync(self, context: GenericContext) -> StateSync:
         """
 
     @abc.abstractmethod
-    def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
-        """Returns the default catalog for the Scheduler.
+    def get_default_catalog_per_gateway(self, context: GenericContext) -> t.Dict[str, str]:
+        """Returns the default catalog for each gateway.
 
         Args:
             context: The SQLMesh Context.
@@ -66,7 +66,7 @@ def state_sync_fingerprint(self, context: GenericContext) -> str:
 class _EngineAdapterStateSyncSchedulerConfig(SchedulerConfig):
     def create_state_sync(self, context: GenericContext) -> StateSync:
         state_connection = (
-            context.config.get_state_connection(context.gateway) or context._connection_config
+            context.config.get_state_connection(context.gateway) or context.connection_config
         )
 
         warehouse_connection = context.config.get_connection(context.gateway)
@@ -110,7 +110,7 @@ def create_state_sync(self, context: GenericContext) -> StateSync:
 
     def state_sync_fingerprint(self, context: GenericContext) -> str:
         state_connection = (
-            context.config.get_state_connection(context.gateway) or context._connection_config
+            context.config.get_state_connection(context.gateway) or context.connection_config
         )
         return md5(
             [
@@ -132,12 +132,16 @@ def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
             state_sync=context.state_sync,
             snapshot_evaluator=context.snapshot_evaluator,
             create_scheduler=context.create_scheduler,
-            default_catalog=self.get_default_catalog(context),
+            default_catalog=context.default_catalog,
             console=context.console,
         )
 
-    def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
-        return context.engine_adapter.default_catalog
+    def get_default_catalog_per_gateway(self, context: GenericContext) -> t.Dict[str, str]:
+        default_catalogs_per_gateway: t.Dict[str, str] = {}
+        for gateway, adapter in context.engine_adapters.items():
+            if catalog := adapter.default_catalog:
+                default_catalogs_per_gateway[gateway] = catalog
+        return default_catalogs_per_gateway
 
 
 SCHEDULER_CONFIG_TO_TYPE = {

sqlmesh/core/context.py

@@ -62,7 +62,9 @@
     Config,
     load_configs,
 )
+from sqlmesh.core.config.connection import ConnectionConfig
 from sqlmesh.core.config.loader import C
+from sqlmesh.core.config.root import RegexKeyDict
 from sqlmesh.core.console import get_console
 from sqlmesh.core.context_diff import ContextDiff
 from sqlmesh.core.dialect import (
@@ -91,6 +93,7 @@
 from sqlmesh.core.plan.definition import UserProvidedFlags
 from sqlmesh.core.reference import ReferenceGraph
 from sqlmesh.core.scheduler import Scheduler, CompletionStatus
+from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.core.schema_loader import create_external_models_file
 from sqlmesh.core.selector import Selector
 from sqlmesh.core.snapshot import (
@@ -367,6 +370,9 @@ def __init__(
         self._excluded_requirements: t.Set[str] = set()
         self._default_catalog: t.Optional[str] = None
         self._default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None
+        self._engine_adapter: t.Optional[EngineAdapter] = None
+        self._connection_config: t.Optional[ConnectionConfig] = None
+        self._test_connection_config: t.Optional[ConnectionConfig] = None
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
 
@@ -407,24 +413,15 @@ def __init__(
             for path, config in self.configs.items()
         ]
 
-        self._connection_config = self.config.get_connection(self.gateway)
+        self._concurrent_tasks = concurrent_tasks
         self._state_connection_config = (
-            self.config.get_state_connection(self.gateway) or self._connection_config
+            self.config.get_state_connection(self.gateway) or self.connection_config
         )
-        self.concurrent_tasks = concurrent_tasks or self._connection_config.concurrent_tasks
-
-        self._engine_adapters: t.Dict[str, EngineAdapter] = {
-            self.selected_gateway: self._connection_config.create_engine_adapter()
-        }
 
         self._snapshot_evaluator: t.Optional[SnapshotEvaluator] = None
 
         self.console = get_console()
-        setattr(self.console, "dialect", self.engine_adapter.dialect)
-
-        self._test_connection_config = self.config.get_test_connection(
-            self.gateway, self.default_catalog, default_catalog_dialect=self.engine_adapter.DIALECT
-        )
+        setattr(self.console, "dialect", self.config.dialect)
 
         self._provided_state_sync: t.Optional[StateSync] = state_sync
         self._state_sync: t.Optional[StateSync] = None
@@ -435,14 +432,6 @@ def __init__(
         self.users = list({user.username: user for user in self.users}.values())
         self._register_notification_targets()
 
-        if (
-            self.config.environment_catalog_mapping
-            and not self.engine_adapter.catalog_support.is_multi_catalog_supported
-        ):
-            raise SQLMeshError(
-                "Environment catalog mapping is only supported for engine adapters that support multiple catalogs"
-            )
-
         if load:
             self.load()
 
@@ -453,7 +442,9 @@ def default_dialect(self) -> t.Optional[str]:
     @property
     def engine_adapter(self) -> EngineAdapter:
         """Returns the default engine adapter."""
-        return self._engine_adapters[self.selected_gateway]
+        if self._engine_adapter is None:
+            self._engine_adapter = self.connection_config.create_engine_adapter()
+        return self._engine_adapter
 
     @property
     def snapshot_evaluator(self) -> SnapshotEvaluator:
@@ -980,8 +971,8 @@ def requirements(self) -> t.Dict[str, str]:
 
     @property
     def default_catalog(self) -> t.Optional[str]:
-        if self._default_catalog is None:
-            self._default_catalog = self._scheduler.get_default_catalog(self)
+        if self._default_catalog is None and self.default_catalog_per_gateway:
+            self._default_catalog = self.default_catalog_per_gateway[self.selected_gateway]
         return self._default_catalog
 
     @python_api_analytics
@@ -1538,7 +1529,7 @@ def plan_builder(
             allow_destructive_models=expanded_destructive_models,
             environment_ttl=environment_ttl,
             environment_suffix_target=self.config.environment_suffix_target,
-            environment_catalog_mapping=self.config.environment_catalog_mapping,
+            environment_catalog_mapping=self.environment_catalog_mapping,
             categorizer_config=categorizer_config or self.auto_categorize_changes,
             auto_categorization_enabled=not no_auto_categorization,
             effective_from=effective_from,
@@ -1550,7 +1541,7 @@ def plan_builder(
             ),
             end_bounded=not run,
             ensure_finalized_snapshots=self.config.plan.use_finalized_state,
-            engine_schema_differ=self.engine_adapter.SCHEMA_DIFFER,
+            engine_schema_differ=SchemaDiffer(),  # TODO: fix to properly handle it
             interval_end_per_model=max_interval_end_per_model,
             console=self.console,
             user_provided_flags=user_provided_flags,
@@ -1639,7 +1630,7 @@ def diff(self, environment: t.Optional[str] = None, detailed: bool = False) -> b
             self.console.show_model_difference_summary(
                 context_diff,
                 EnvironmentNamingInfo.from_environment_catalog_mapping(
-                    self.config.environment_catalog_mapping,
+                    self.environment_catalog_mapping,
                     name=environment,
                     suffix_target=self.config.environment_suffix_target,
                     normalize_name=context_diff.normalize_environment_name,
@@ -1993,7 +1984,7 @@ def create_test(
 
         try:
             model_to_test = self.get_model(model, raise_if_missing=True)
-            test_adapter = self._test_connection_config.create_engine_adapter(
+            test_adapter = self.test_connection_config.create_engine_adapter(
                 register_comments_override=False
             )
 
@@ -2039,7 +2030,7 @@ def test(
             preserve_fixtures=preserve_fixtures,
             stream=stream,
             default_catalog=self.default_catalog,
-            default_catalog_dialect=self.engine_adapter.DIALECT,
+            default_catalog_dialect=self.config.dialect or "",
         )
 
     @python_api_analytics
@@ -2478,7 +2469,7 @@ def _run_plan_tests(
                 self.console.log_test_results(
                     result,
                     test_output,
-                    self._test_connection_config._engine_adapter.DIALECT,
+                    self.test_connection_config._engine_adapter.DIALECT,
                 )
             if not result.wasSuccessful():
                 raise PlanError(
@@ -2499,7 +2490,7 @@ def _model_tables(self) -> t.Dict[str, str]:
                 if snapshot.version
                 else snapshot.qualified_view_name.for_environment(
                     EnvironmentNamingInfo.from_environment_catalog_mapping(
-                        self.config.environment_catalog_mapping,
+                        self.environment_catalog_mapping,
                         name=c.PROD,
                         suffix_target=self.config.environment_suffix_target,
                     )
@@ -2511,24 +2502,63 @@ def _model_tables(self) -> t.Dict[str, str]:
     @cached_property
     def engine_adapters(self) -> t.Dict[str, EngineAdapter]:
         """Returns all the engine adapters for the gateways defined in the configuration."""
+        adapters: t.Dict[str, EngineAdapter] = {self.selected_gateway: self.engine_adapter}
         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(concurrent_tasks=self.concurrent_tasks)
-                self._engine_adapters[gateway_name] = adapter
-        return self._engine_adapters
+                adapters[gateway_name] = adapter
+        return adapters
 
     @cached_property
     def default_catalog_per_gateway(self) -> t.Dict[str, str]:
         """Returns the default catalogs for each engine adapter."""
         if self._default_catalog_per_gateway is None:
-            self._default_catalog_per_gateway = {
-                name: adapter.default_catalog
-                for name, adapter in self.engine_adapters.items()
-                if adapter.default_catalog
-            }
+            self._default_catalog_per_gateway = self._scheduler.get_default_catalog_per_gateway(
+                self
+            )
         return self._default_catalog_per_gateway
 
+    @cached_property
+    def concurrent_tasks(self) -> int:
+        if self._concurrent_tasks is None:
+            self._concurrent_tasks = self.connection_config.concurrent_tasks
+        return self._concurrent_tasks
+
+    @cached_property
+    def connection_config(self) -> ConnectionConfig:
+        if self._connection_config is None:
+            self._connection_config = self.config.get_connection(self.selected_gateway)
+        return self._connection_config
+
+    @cached_property
+    def test_connection_config(self) -> ConnectionConfig:
+        if self._test_connection_config is None:
+            self._test_connection_config = self.config.get_test_connection(
+                self.gateway,
+                self.default_catalog,
+                default_catalog_dialect=self.engine_adapter.DIALECT,
+            )
+        return self._test_connection_config
+
+    @cached_property
+    def environment_catalog_mapping(self) -> RegexKeyDict:
+        engine_adapter = None
+        try:
+            engine_adapter = self.engine_adapter
+        except Exception:
+            pass
+
+        if (
+            self.config.environment_catalog_mapping
+            and engine_adapter
+            and not self.engine_adapter.catalog_support.is_multi_catalog_supported
+        ):
+            raise SQLMeshError(
+                "Environment catalog mapping is only supported for engine adapters that support multiple catalogs"
+            )
+        return self.config.environment_catalog_mapping
+
     def _get_engine_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
         if gateway:
             if adapter := self.engine_adapters.get(gateway):

sqlmesh/integrations/github/cicd/controller.py

@@ -726,7 +726,7 @@ def conclusion_handler(
                 self._console.log_test_results(
                     result,
                     output,
-                    self._context._test_connection_config._engine_adapter.DIALECT,
+                    self._context.test_connection_config._engine_adapter.DIALECT,
                 )
                 test_summary = self._console.consume_captured_output()
                 test_title = "Tests Passed" if result.wasSuccessful() else "Tests Failed"

tests/core/test_config.py

@@ -742,7 +742,7 @@ def test_multi_gateway_config(tmp_path, mocker: MockerFixture):
 
     ctx = Context(paths=tmp_path, config=config)
 
-    assert isinstance(ctx._connection_config, RedshiftConnectionConfig)
+    assert isinstance(ctx.connection_config, RedshiftConnectionConfig)
     assert len(ctx.engine_adapters) == 3
     assert isinstance(ctx.engine_adapters["athena"], AthenaEngineAdapter)
     assert isinstance(ctx.engine_adapters["redshift"], RedshiftEngineAdapter)
@@ -782,7 +782,7 @@ def test_multi_gateway_single_threaded_config(tmp_path):
     )
 
     ctx = Context(paths=tmp_path, config=config)
-    assert isinstance(ctx._connection_config, DuckDBConnectionConfig)
+    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)

tests/core/test_context.py

@@ -332,18 +332,18 @@ def test_evaluate_limit():
 def test_gateway_specific_adapters(copy_to_temp_path, mocker):
     path = copy_to_temp_path("examples/sushi")
     ctx = Context(paths=path, config="isolated_systems_config", gateway="prod")
-    assert len(ctx._engine_adapters) == 3
-    assert ctx.engine_adapter == ctx._engine_adapters["prod"]
-    assert ctx._get_engine_adapter("dev") == ctx._engine_adapters["dev"]
+    assert len(ctx.engine_adapters) == 3
+    assert ctx.engine_adapter == ctx.engine_adapters["prod"]
+    assert ctx._get_engine_adapter("dev") == ctx.engine_adapters["dev"]
 
     ctx = Context(paths=path, config="isolated_systems_config")
-    assert len(ctx._engine_adapters) == 3
-    assert ctx.engine_adapter == ctx._engine_adapters["dev"]
+    assert len(ctx.engine_adapters) == 3
+    assert ctx.engine_adapter == ctx.engine_adapters["dev"]
 
     ctx = Context(paths=path, config="isolated_systems_config")
     assert len(ctx.engine_adapters) == 3
     assert ctx.engine_adapter == ctx._get_engine_adapter()
-    assert ctx._get_engine_adapter("test") == ctx._engine_adapters["test"]
+    assert ctx._get_engine_adapter("test") == ctx.engine_adapters["test"]
 
 
 def test_multiple_gateways(tmp_path: Path):
@@ -800,7 +800,8 @@ def test_janitor(sushi_context, mocker: MockerFixture) -> None:
         ),
     ]
 
-    sushi_context._engine_adapters = {sushi_context.config.default_gateway: adapter_mock}
+    sushi_context._engine_adapter = adapter_mock
+    sushi_context.engine_adapters = {sushi_context.config.default_gateway: adapter_mock}
     sushi_context._state_sync = state_sync_mock
     state_sync_mock.get_expired_snapshots.return_value = []
 

tests/core/test_integration.py

@@ -4872,13 +4872,11 @@ def test_multi(mocker):
     assert context.fetchdf("select * from after_1").to_dict()["repo_1"][0] == "repo_1"
     assert context.fetchdf("select * from after_2").to_dict()["repo_2"][0] == "repo_2"
 
-    adapter = context.engine_adapter
     context = Context(
         paths=["examples/multi/repo_1"],
         state_sync=context.state_sync,
         gateway="memory",
     )
-    context._engine_adapters["memory"] = adapter
 
     model = context.get_model("bronze.a")
     assert model.project == "repo_1"
@@ -4935,6 +4933,8 @@ def test_multi_virtual_layer(copy_to_temp_path):
     )
 
     context = Context(paths=paths, config=config)
+    assert context.default_catalog_per_gateway == {"first": "db_1", "second": "db_2"}
+    assert len(context.engine_adapters) == 2
 
     # For the model without gateway the default should be used and the gateway variable should overide the global
     assert (

tests/core/test_model.py

@@ -4518,7 +4518,7 @@ def test_model_session_properties(sushi_context):
                 name test_schema.test_model,
                 session_properties (
                     'query_label' = (
-                        'some value', 
+                        'some value',
                         'another value',
                         'yet another value',
                     )
@@ -8350,7 +8350,7 @@ def test_gateway_specific_render(assert_exp_eq) -> None:
         default_gateway="main",
     )
     context = Context(config=config)
-    assert context.engine_adapter == context._engine_adapters["main"]
+    assert context.engine_adapter == context.engine_adapters["main"]
 
     @model(
         name="dummy_model",
@@ -8376,7 +8376,7 @@ def dummy_model_entry(evaluator: MacroEvaluator) -> exp.Select:
         """,
     )
     assert isinstance(context._get_engine_adapter("duckdb"), DuckDBEngineAdapter)
-    assert len(context._engine_adapters) == 2
+    assert len(context.engine_adapters) == 2
 
 
 def test_model_on_virtual_update(make_snapshot: t.Callable):