Refactors; account for identical db names between warehouses when creating schemas

Author: themisvaltinos

sqlmesh/core/context.py

@@ -364,7 +364,7 @@ def __init__(
         self._environment_statements: t.List[EnvironmentStatements] = []
         self._excluded_requirements: t.Set[str] = set()
         self._default_catalog: t.Optional[str] = None
-        self._catalogs: t.Dict[str, str] = {}
+        self._default_catalog_per_gateway: t.Dict[str, str] = {}
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
 
@@ -2225,15 +2225,15 @@ def engine_adapters(self) -> t.Dict[str, EngineAdapter]:
         return self._engine_adapters
 
     @cached_property
-    def catalogs(self) -> t.Dict[str, str]:
+    def default_catalog_per_gateway(self) -> t.Dict[str, str]:
         """Returns the catalogs for each engine adapter in a multi virtual layer setup when the catalog isn't shared."""
         if self.gateway_managed_virtual_layer:
-            self._catalogs = {
+            self._default_catalog_per_gateway = {
                 name: adapter.default_catalog
                 for name, adapter in self.engine_adapters.items()
                 if adapter.default_catalog
             }
-        return self._catalogs
+        return self._default_catalog_per_gateway
 
     def _get_engine_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
         if gateway:
@@ -2317,10 +2317,10 @@ def _cleanup_environments(self) -> None:
         expired_environments = self.state_sync.delete_expired_environments()
 
         cleanup_expired_views(
-            adapter=self.engine_adapter,
+            default_adapter=self.engine_adapter,
+            engine_adapters=self.engine_adapters,
             environments=expired_environments,
             console=self.console,
-            engine_adapters=self.engine_adapters,
         )
 
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:

sqlmesh/core/loader.py

@@ -468,7 +468,7 @@ def _load() -> t.List[Model]:
                         default_catalog=self.context.default_catalog,
                         infer_names=self.config.model_naming.infer_names,
                         signal_definitions=signals,
-                        catalogs=self.context.catalogs,
+                        default_catalog_per_gateway=self.context.default_catalog_per_gateway,
                     )
                 except Exception as ex:
                     raise ConfigError(f"Failed to load model definition at '{path}'.\n{ex}")
@@ -526,7 +526,7 @@ def _load_python_models(
                             default_catalog=self.context.default_catalog,
                             infer_names=self.config.model_naming.infer_names,
                             audit_definitions=audits,
-                            catalogs=self.context.catalogs,
+                            default_catalog_per_gateway=self.context.default_catalog_per_gateway,
                         ):
                             if model.enabled:
                                 models[model.fqn] = model

sqlmesh/core/model/decorator.py

@@ -93,6 +93,7 @@ def models(
         path: Path,
         module_path: Path,
         dialect: t.Optional[str] = None,
+        default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
         **loader_kwargs: t.Any,
     ) -> t.List[Model]:
         return create_models_from_blueprints(
@@ -103,6 +104,7 @@ def models(
             path=path,
             module_path=module_path,
             dialect=dialect,
+            default_catalog_per_gateway=default_catalog_per_gateway,
             **loader_kwargs,
         )
 

sqlmesh/core/model/definition.py

@@ -1886,6 +1886,7 @@ def create_models_from_blueprints(
     path: Path = Path(),
     module_path: Path = Path(),
     dialect: DialectType = None,
+    default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
     **loader_kwargs: t.Any,
 ) -> t.List[Model]:
     model_blueprints: t.List[Model] = []
@@ -1907,11 +1908,10 @@ def create_models_from_blueprints(
         else:
             gateway_name = None
 
-        # We pop to avoid pydantic validation issues since catalogs is not a model property
         if (
-            (catalogs := loader_kwargs.pop("catalogs", None))
+            default_catalog_per_gateway
             and gateway_name
-            and (catalog := catalogs.get(gateway_name))
+            and (catalog := default_catalog_per_gateway.get(gateway_name))
         ):
             loader_kwargs["default_catalog"] = catalog
 
@@ -1935,6 +1935,7 @@ def load_sql_based_models(
     path: Path = Path(),
     module_path: Path = Path(),
     dialect: DialectType = None,
+    default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
     **loader_kwargs: t.Any,
 ) -> t.List[Model]:
     gateway: t.Optional[exp.Expression] = None
@@ -1972,6 +1973,7 @@ def load_sql_based_models(
         path=path,
         module_path=module_path,
         dialect=dialect,
+        default_catalog_per_gateway=default_catalog_per_gateway,
         **loader_kwargs,
     )
 

sqlmesh/core/plan/evaluator.py

@@ -424,9 +424,7 @@ def _demote_snapshots(
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
         self.snapshot_evaluator.demote(
-            target_snapshots,
-            environment_naming_info,
-            on_complete=on_complete,
+            target_snapshots, environment_naming_info, on_complete=on_complete
         )
 
     def _restate(self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapshot]) -> None:

sqlmesh/core/snapshot/evaluator.py

@@ -227,21 +227,20 @@ def promote(
             on_complete: A callback to call on each successfully promoted snapshot.
         """
 
-        gateway_by_schema: t.Dict[t.Any, str] = {}
-        tables: t.List[t.Any] = []
+        tables_by_gateway: t.Dict[t.Union[str, None], t.List[exp.Table]] = defaultdict(list)
         for snapshot in target_snapshots:
             if snapshot.is_model and not snapshot.is_symbolic:
+                gateway = (
+                    snapshot.model_gateway if environment_naming_info.gateway_managed else None
+                )
+                adapter = self._get_adapter(gateway)
                 table = snapshot.qualified_view_name.table_for_environment(
-                    environment_naming_info,
-                    dialect=self._get_adapter(snapshot.model_gateway).dialect
-                    if environment_naming_info.gateway_managed
-                    else self.adapter.dialect,
+                    environment_naming_info, dialect=adapter.dialect
                 )
-                tables.append(table)
-                if environment_naming_info.gateway_managed:
-                    table_schema = d.schema_(table.db, catalog=table.catalog)
-                    gateway_by_schema[table_schema] = snapshot.model_gateway or ""
-        self._create_schemas(tables=tables, gateways=gateway_by_schema)
+                tables_by_gateway[gateway].append(table)
+
+        for gateway, tables in tables_by_gateway.items():
+            self._create_schemas(tables=tables, gateway=gateway)
 
         deployability_index = deployability_index or DeployabilityIndex.all_deployable()
         with self.concurrent_context():
@@ -301,8 +300,9 @@ def create(
             allow_destructive_snapshots: Set of snapshots that are allowed to have destructive schema changes.
         """
         snapshots_with_table_names = defaultdict(set)
-        tables_by_schema = defaultdict(set)
-        gateway_by_schema: t.Dict[exp.Table, str] = {}
+        tables_by_gateway_and_schema: t.Dict[t.Union[str, None], t.Dict[exp.Table, set[str]]] = (
+            defaultdict(lambda: defaultdict(set))
+        )
         table_deployability: t.Dict[str, bool] = {}
         allow_destructive_snapshots = allow_destructive_snapshots or set()
 
@@ -324,24 +324,32 @@ def create(
                 snapshots_with_table_names[snapshot].add(table.name)
                 table_deployability[table.name] = is_deployable
                 table_schema = d.schema_(table.db, catalog=table.catalog)
-                tables_by_schema[table_schema].add(table.name)
-                gateway_by_schema[table_schema] = snapshot.model.gateway or ""
+                tables_by_gateway_and_schema[snapshot.model_gateway][table_schema].add(table.name)
 
-        def _get_data_objects(schema: exp.Table, gateway: t.Optional[str] = None) -> t.Set[str]:
+        def _get_data_objects(
+            schema: exp.Table,
+            object_names: t.Optional[t.Set[str]] = None,
+            gateway: t.Optional[str] = None,
+        ) -> t.Set[str]:
             logger.info("Listing data objects in schema %s", schema.sql())
-            objs = self.get_adapter(gateway).get_data_objects(schema, tables_by_schema[schema])
+            objs = self._get_adapter(gateway).get_data_objects(schema, object_names)
             return {obj.name for obj in objs}
 
         with self.concurrent_context():
-            existing_objects = {
-                obj
-                for objs in concurrent_apply_to_values(
-                    list(tables_by_schema),
-                    lambda s: _get_data_objects(s, gateway_by_schema[s]),
-                    self.ddl_concurrent_tasks,
-                )
-                for obj in objs
-            }
+            existing_objects: t.Set[str] = set()
+            for gateway, tables_by_schema in tables_by_gateway_and_schema.items():
+                objs_for_gateway = {
+                    obj
+                    for objs in concurrent_apply_to_values(
+                        list(tables_by_schema),
+                        lambda s: _get_data_objects(
+                            schema=s, object_names=tables_by_schema.get(s), gateway=gateway
+                        ),
+                        self.ddl_concurrent_tasks,
+                    )
+                    for obj in objs
+                }
+                existing_objects.update(objs_for_gateway)
 
         snapshots_to_create = []
         target_deployability_flags: t.Dict[str, t.List[bool]] = defaultdict(list)
@@ -359,7 +367,10 @@ def _get_data_objects(schema: exp.Table, gateway: t.Optional[str] = None) -> t.S
             return
         if on_start:
             on_start(len(snapshots_to_create))
-        self._create_schemas(tables_by_schema, gateway_by_schema)
+
+        for gateway, tables_by_schema in tables_by_gateway_and_schema.items():
+            self._create_schemas(tables=tables_by_schema, gateway=gateway)
+
         self._create_snapshots(
             snapshots_to_create=snapshots_to_create,
             snapshots=snapshots,
@@ -1072,7 +1083,7 @@ def _audit(
     def _create_schemas(
         self,
         tables: t.Iterable[t.Union[exp.Table, str]],
-        gateways: t.Optional[t.Dict[exp.Table, str]] = None,
+        gateway: t.Optional[str] = None,
     ) -> None:
         table_exprs = [exp.to_table(t) for t in tables]
         unique_schemas = {(t.args["db"], t.args.get("catalog")) for t in table_exprs if t and t.db}
@@ -1081,7 +1092,7 @@ def _create_schemas(
         for schema_name, catalog in unique_schemas:
             schema = schema_(schema_name, catalog)
             logger.info("Creating schema '%s'", schema)
-            adapter = self.get_adapter(gateways.get(schema)) if gateways else self.adapter
+            adapter = self._get_adapter(gateway)
             adapter.create_schema(schema)
 
     def get_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:

sqlmesh/core/state_sync/common.py

@@ -21,10 +21,10 @@
 
 
 def cleanup_expired_views(
-    adapter: EngineAdapter,
+    default_adapter: EngineAdapter,
+    engine_adapters: t.Dict[str, EngineAdapter],
     environments: t.List[Environment],
     console: t.Optional[Console] = None,
-    engine_adapters: t.Optional[t.Dict[str, EngineAdapter]] = None,
 ) -> None:
     expired_schema_environments = [
         environment for environment in environments if environment.suffix_target.is_schema
@@ -36,8 +36,8 @@ def cleanup_expired_views(
     # We have to use the corresponding adapter if the virtual layer is gateway managed
     def get_adapter(gateway_managed: bool, gateway: t.Optional[str] = None) -> EngineAdapter:
         if gateway_managed and gateway:
-            return (engine_adapters or {}).get(gateway, adapter)
-        return adapter
+            return engine_adapters.get(gateway, default_adapter)
+        return default_adapter
 
     # Drop the schemas for the expired environments
     for engine_adapter, expired_catalog, expired_schema in {

sqlmesh/engines/commands.py

@@ -140,7 +140,7 @@ def cleanup(
     if isinstance(command_payload, str):
         command_payload = CleanupCommandPayload.parse_raw(command_payload)
 
-    cleanup_expired_views(evaluator.adapter, command_payload.environments)
+    cleanup_expired_views(evaluator.adapter, evaluator.adapters, command_payload.environments)
     evaluator.cleanup(command_payload.tasks)
 
 

tests/core/state_sync/test_state_sync.py

@@ -2635,7 +2635,7 @@ def test_cleanup_expired_views(
         previous_plan_id="test_plan_id",
         catalog_name_override="catalog_override",
     )
-    cleanup_expired_views(adapter, [schema_environment, table_environment])
+    cleanup_expired_views(adapter, {}, [schema_environment, table_environment])
     assert adapter.drop_schema.called
     assert adapter.drop_view.called
     assert adapter.drop_schema.call_args_list == [