Refactor: Move engine adapter instantiation after loading on-demand

Author: themisvaltinos

sqlmesh/core/config/scheduler.py

@@ -54,6 +54,14 @@ def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
             context: The SQLMesh Context.
         """
 
+    @abc.abstractmethod
+    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.
+        """
+
     @abc.abstractmethod
     def state_sync_fingerprint(self, context: GenericContext) -> str:
         """Returns the fingerprint of the State Sync configuration.
@@ -139,6 +147,13 @@ def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
     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]:
+        return {
+            name: adapter.default_catalog
+            for name, adapter in context.engine_adapters.items()
+            if adapter.default_catalog
+        }
+
 
 SCHEDULER_CONFIG_TO_TYPE = {
     tpe.all_field_infos()["type_"].default: tpe

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 (
@@ -366,7 +369,7 @@ def __init__(
         self._environment_statements: t.List[EnvironmentStatements] = []
         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._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
 
@@ -407,24 +410,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.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 +429,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 +439,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:
@@ -1535,7 +1523,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,
@@ -1547,7 +1535,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,
@@ -1636,7 +1624,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,
@@ -2036,7 +2024,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
@@ -2496,7 +2484,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,
                     )
@@ -2508,23 +2496,47 @@ 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
-            }
-        return self._default_catalog_per_gateway
+        if self.gateway_managed_virtual_layer:
+            return self._scheduler.get_default_catalog_per_gateway(self)
+        return {}
+
+    @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:
+        return self.config.get_connection(self.gateway)
+
+    @cached_property
+    def _test_connection_config(self) -> ConnectionConfig:
+        return self.config.get_test_connection(
+            self.gateway, self.default_catalog, default_catalog_dialect=self.engine_adapter.DIALECT
+        )
+
+    @cached_property
+    def environment_catalog_mapping(self) -> RegexKeyDict:
+        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"
+            )
+        return self.config.environment_catalog_mapping
 
     def _get_engine_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
         if gateway:

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

@@ -4878,7 +4878,7 @@ def test_multi(mocker):
         state_sync=context.state_sync,
         gateway="memory",
     )
-    context._engine_adapters["memory"] = adapter
+    context._engine_adapter = adapter
 
     model = context.get_model("bronze.a")
     assert model.project == "repo_1"
@@ -5064,6 +5064,39 @@ def test_multi_virtual_layer(copy_to_temp_path):
         context.apply(plan)
 
 
+def test_multi_virtual_layer_catalogs(copy_to_temp_path):
+    paths = copy_to_temp_path("tests/fixtures/multi_virtual_layer")
+    path = Path(paths[0])
+    first_db_path = str(path / "db_1.db")
+    second_db_path = str(path / "db_2.db")
+
+    config = Config(
+        gateways={
+            "first": GatewayConfig(
+                connection=DuckDBConnectionConfig(database=first_db_path),
+                variables={"overriden_var": "gateway_1"},
+            ),
+            "second": GatewayConfig(
+                connection=DuckDBConnectionConfig(database=second_db_path),
+                variables={"overriden_var": "gateway_2"},
+            ),
+        },
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        model_naming=NameInferenceConfig(infer_names=True),
+        default_gateway="first",
+        variables={"overriden_var": "global", "global_one": 88},
+    )
+
+    # With gateway_managed_virtual_layer to False the catalogs won't be retrieved
+    context = Context(paths=paths, config=config)
+    assert context.default_catalog_per_gateway == {}
+
+    config.gateway_managed_virtual_layer = True
+    context = Context(paths=paths, config=config)
+    assert context.default_catalog_per_gateway == {"first": "db_1", "second": "db_2"}
+    assert len(context.engine_adapters) == 2
+
+
 def test_multi_dbt(mocker):
     context = Context(paths=["examples/multi_dbt/bronze", "examples/multi_dbt/silver"])
     context._new_state_sync().reset(default_catalog=context.default_catalog)

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):

tests/core/test_test.py

@@ -2128,16 +2128,16 @@ def test_test_with_gateway_specific_model(tmp_path: Path, mocker: MockerFixture)
         return_value=pd.DataFrame({"c": [5]}),
     )
 
-    assert context.engine_adapter == context._engine_adapters["main"]
+    assert context.engine_adapter == context.engine_adapters["main"]
     with pytest.raises(
         SQLMeshError, match=r"Gateway 'wrong' not found in the available engine adapters."
     ):
         context._get_engine_adapter("wrong")
 
     # Create test should use the gateway specific engine adapter
     context.create_test("sqlmesh_example.gw_model", input_queries=input_queries, overwrite=True)
-    assert context._get_engine_adapter("second") == context._engine_adapters["second"]
-    assert len(context._engine_adapters) == 2
+    assert context._get_engine_adapter("second") == context.engine_adapters["second"]
+    assert len(context.engine_adapters) == 2
 
     test = load_yaml(context.path / c.TESTS / "test_gw_model.yaml")