Refactors; revise to keep model gateway in SnapshotTableInfo

Author: themisvaltinos

sqlmesh/core/context.py

@@ -2311,23 +2311,11 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
     def _cleanup_environments(self) -> None:
         expired_environments = self.state_sync.delete_expired_environments()
 
-        environment_snapshot_adapters: t.Dict[str, t.Dict[str, EngineAdapter]] = {}
-        for environment in expired_environments:
-            snapshot_adapters: t.Dict[str, EngineAdapter] = {}
-            if environment.gateway_managed_virtual_layer:
-                snapshots = self.state_sync.get_snapshots(environment.snapshots)
-                for snapshot_id, snapshot in snapshots.items():
-                    if snapshot.is_model and not snapshot.is_symbolic:
-                        snapshot_adapters[snapshot_id.name] = self._get_engine_adapter(
-                            snapshot.model_gateway
-                        )
-                environment_snapshot_adapters[environment.name] = snapshot_adapters
-
         cleanup_expired_views(
             adapter=self.engine_adapter,
             environments=expired_environments,
             console=self.console,
-            environment_snapshot_adapters=environment_snapshot_adapters,
+            engine_adapters=self.engine_adapters,
         )
 
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:

sqlmesh/core/context_diff.py

@@ -121,11 +121,7 @@ def create(
         env = state_reader.get_environment(environment)
 
         create_from_env_exists = False
-        if (
-            env is None
-            or env.expired
-            or env.gateway_managed_virtual_layer != gateway_managed_virtual_layer
-        ):
+        if env is None or env.expired or env.gateway_managed != gateway_managed_virtual_layer:
             env = state_reader.get_environment(create_from.lower())
 
             if not env and create_from != c.PROD:

sqlmesh/core/environment.py

@@ -16,7 +16,7 @@
 from sqlmesh.utils.date import TimeLike, now_timestamp
 from sqlmesh.utils.jinja import JinjaMacroRegistry
 from sqlmesh.utils.metaprogramming import Executable
-from sqlmesh.utils.pydantic import PydanticModel, field_validator
+from sqlmesh.utils.pydantic import PydanticModel, field_validator, ValidationInfo
 
 T = t.TypeVar("T", bound="EnvironmentNamingInfo")
 PydanticType = t.TypeVar("PydanticType", bound="PydanticModel")
@@ -32,30 +32,27 @@ class EnvironmentNamingInfo(PydanticModel):
         catalog_name_override: The name of the catalog to use for this environment if an override was provided
         normalize_name: Indicates whether the environment's name will be normalized. For example, if it's
             `dev`, then it will become `DEV` when targeting Snowflake.
-        gateway_managed_virtual_layer: Determines whether the virtual layer's views are created by the model-specific
+        gateway_managed: Determines whether the virtual layer's views are created by the model-specific
             gateways, otherwise the default gateway is used. Default: False.
     """
 
     name: str = c.PROD
     suffix_target: EnvironmentSuffixTarget = Field(default=EnvironmentSuffixTarget.SCHEMA)
     catalog_name_override: t.Optional[str] = None
     normalize_name: bool = True
-    gateway_managed_virtual_layer: bool = False
+    gateway_managed: bool = False
 
     @field_validator("name", mode="before")
     @classmethod
     def _sanitize_name(cls, v: str) -> str:
         return word_characters_only(v).lower()
 
-    @field_validator("normalize_name", mode="before")
+    @field_validator("normalize_name", "gateway_managed", mode="before")
     @classmethod
-    def _validate_normalize_name(cls, v: t.Any) -> bool:
-        return True if v is None else bool(v)
-
-    @field_validator("gateway_managed_virtual_layer", mode="before")
-    @classmethod
-    def _validate_gateway_managed_virtual_layer(cls, v: t.Any) -> bool:
-        return False if v is None else bool(v)
+    def _validate_boolean_field(cls, v: t.Any, info: ValidationInfo) -> bool:
+        if v is None:
+            return info.field_name == "normalize_name"
+        return bool(v)
 
     @t.overload
     @classmethod
@@ -202,7 +199,7 @@ def naming_info(self) -> EnvironmentNamingInfo:
             suffix_target=self.suffix_target,
             catalog_name_override=self.catalog_name_override,
             normalize_name=self.normalize_name,
-            gateway_managed_virtual_layer=self.gateway_managed_virtual_layer,
+            gateway_managed=self.gateway_managed,
         )
 
     @property

sqlmesh/core/plan/builder.py

@@ -151,7 +151,7 @@ def __init__(
             name=self._context_diff.environment,
             suffix_target=environment_suffix_target,
             normalize_name=self._context_diff.normalize_environment_name,
-            gateway_managed_virtual_layer=self._context_diff.gateway_managed_virtual_layer,
+            gateway_managed=self._context_diff.gateway_managed_virtual_layer,
         )
 
         self._latest_plan: t.Optional[Plan] = None

sqlmesh/core/plan/evaluator.py

@@ -423,14 +423,8 @@ def _demote_snapshots(
         environment_naming_info: EnvironmentNamingInfo,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
-        # In a multi virtual layer setup we need the gateway info from the snapshots for demotion
-        snapshots_to_demote: t.List[Snapshot] = []
-        if environment_naming_info.gateway_managed_virtual_layer:
-            removed_snapshots = self.state_sync.get_snapshots(target_snapshots)
-            snapshots_to_demote = [removed_snapshots[s.snapshot_id] for s in target_snapshots]
-
         self.snapshot_evaluator.demote(
-            snapshots_to_demote or target_snapshots,
+            target_snapshots,
             environment_naming_info,
             on_complete=on_complete,
         )

sqlmesh/core/snapshot/definition.py

@@ -477,6 +477,7 @@ class SnapshotTableInfo(PydanticModel, SnapshotInfoMixin, frozen=True):
     base_table_name_override: t.Optional[str] = None
     custom_materialization: t.Optional[str] = None
     dev_table_suffix: str
+    model_gateway: t.Optional[str] = None
 
     def __lt__(self, other: SnapshotTableInfo) -> bool:
         return self.name < other.name
@@ -1177,6 +1178,7 @@ def table_info(self) -> SnapshotTableInfo:
             node_type=self.node_type,
             custom_materialization=custom_materialization,
             dev_table_suffix=self.dev_table_suffix,
+            model_gateway=self.model_gateway,
         )
 
     @property
@@ -1336,7 +1338,6 @@ def __getstate__(self) -> t.Dict[t.Any, t.Any]:
 class SnapshotTableCleanupTask(PydanticModel):
     snapshot: SnapshotTableInfo
     dev_table_only: bool
-    gateway: t.Optional[str] = None
 
 
 SnapshotIdLike = t.Union[SnapshotId, SnapshotTableInfo, Snapshot]

sqlmesh/core/snapshot/evaluator.py

@@ -234,11 +234,11 @@ def promote(
                 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_virtual_layer
+                    if environment_naming_info.gateway_managed
                     else self.adapter.dialect,
                 )
                 tables.append(table)
-                if environment_naming_info.gateway_managed_virtual_layer:
+                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)
@@ -437,15 +437,14 @@ def cleanup(
         snapshots_to_dev_table_only = {
             t.snapshot.snapshot_id: t.dev_table_only for t in target_snapshots
         }
-        snapshot_gateways = {t.snapshot.snapshot_id: t.gateway for t in target_snapshots}
 
         with self.concurrent_context():
             concurrent_apply_to_snapshots(
                 [t.snapshot for t in target_snapshots],
                 lambda s: self._cleanup_snapshot(
                     s,
                     snapshots_to_dev_table_only[s.snapshot_id],
-                    self._get_adapter(snapshot_gateways[s.snapshot_id]),
+                    self._get_adapter(s.model_gateway),
                     on_complete,
                 ),
                 self.ddl_concurrent_tasks,
@@ -928,7 +927,7 @@ def _promote_snapshot(
         if snapshot.is_model:
             adapter = (
                 self._get_adapter(snapshot.model_gateway)
-                if environment_naming_info.gateway_managed_virtual_layer
+                if environment_naming_info.gateway_managed
                 else self.adapter
             )
             table_name = snapshot.table_name(deployability_index.is_representative(snapshot))
@@ -965,8 +964,7 @@ def _demote_snapshot(
     ) -> None:
         adapter = (
             self._get_adapter(snapshot.model_gateway)
-            if environment_naming_info.gateway_managed_virtual_layer
-            and isinstance(snapshot, Snapshot)
+            if environment_naming_info.gateway_managed
             else self.adapter
         )
         view_name = snapshot.qualified_view_name.for_environment(

sqlmesh/core/state_sync/common.py

@@ -23,8 +23,8 @@
 def cleanup_expired_views(
     adapter: EngineAdapter,
     environments: t.List[Environment],
-    environment_snapshot_adapters: t.Optional[t.Dict[str, t.Dict[str, EngineAdapter]]] = None,
     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
@@ -33,19 +33,16 @@ def cleanup_expired_views(
         environment for environment in environments if environment.suffix_target.is_table
     ]
 
+    # 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
+
     # Drop the schemas for the expired environments
-    # Note: We have to use the corresponding adapter if it is a gateway managed virtual layer
     for engine_adapter, expired_catalog, expired_schema in {
         (
-            (
-                engine_adapter := (
-                    (environment_dict.get(snapshot.name) or adapter)
-                    if environment.gateway_managed_virtual_layer
-                    and environment_snapshot_adapters
-                    and (environment_dict := environment_snapshot_adapters.get(environment.name))
-                    else adapter
-                )
-            ),
+            (engine_adapter := get_adapter(environment.gateway_managed, snapshot.model_gateway)),
             snapshot.qualified_view_name.catalog_for_environment(
                 environment.naming_info, dialect=engine_adapter.dialect
             ),
@@ -74,15 +71,7 @@ def cleanup_expired_views(
     # Drop the views for the expired environments
     for engine_adapter, expired_view in {
         (
-            (
-                engine_adapter := (
-                    (environment_dict.get(snapshot.name) or adapter)
-                    if environment.gateway_managed_virtual_layer
-                    and environment_snapshot_adapters
-                    and (environment_dict := environment_snapshot_adapters.get(environment.name))
-                    else adapter
-                )
-            ),
+            (engine_adapter := get_adapter(environment.gateway_managed, snapshot.model_gateway)),
             snapshot.qualified_view_name.for_environment(
                 environment.naming_info, dialect=engine_adapter.dialect
             ),

sqlmesh/core/state_sync/db/environment.py

@@ -48,7 +48,7 @@ def __init__(
             "catalog_name_override": exp.DataType.build("text"),
             "previous_finalized_snapshots": exp.DataType.build(blob_type),
             "normalize_name": exp.DataType.build("boolean"),
-            "gateway_managed_virtual_layer": exp.DataType.build("boolean"),
+            "gateway_managed": exp.DataType.build("boolean"),
             "requirements": exp.DataType.build(blob_type),
         }
 
@@ -168,7 +168,6 @@ def delete_expired_environments(self) -> t.List[Environment]:
         Returns:
             A list of deleted environments.
         """
-
         now_ts = now_timestamp()
         filter_expr = exp.LTE(
             this=exp.column("expiration_ts"),
@@ -182,10 +181,8 @@ def delete_expired_environments(self) -> t.List[Environment]:
                 lock_for_update=True,
             ),
         )
-
         environments = [self._environment_from_row(r) for r in rows]
 
-        # Delete the expired environments
         self.engine_adapter.delete_from(
             self.environments_table,
             where=filter_expr,
@@ -331,7 +328,7 @@ def _environment_to_df(environment: Environment) -> pd.DataFrame:
                     else None
                 ),
                 "normalize_name": environment.normalize_name,
-                "gateway_managed_virtual_layer": environment.gateway_managed_virtual_layer,
+                "gateway_managed": environment.gateway_managed,
                 "requirements": json.dumps(environment.requirements),
             }
         ]

sqlmesh/core/state_sync/db/facade.py

@@ -201,8 +201,7 @@ def promote(
             }
             if (
                 not existing_environment.expired
-                and existing_environment.gateway_managed_virtual_layer
-                == environment.gateway_managed_virtual_layer
+                and existing_environment.gateway_managed == environment.gateway_managed
             ):
                 if environment.previous_plan_id != existing_environment.plan_id:
                     raise ConflictingPlanError(