move virtual_environment_model attribute into model meta

Author: izeigerman

sqlmesh/core/context.py

@@ -2909,7 +2909,6 @@ def _nodes_to_snapshots(self, nodes: t.Dict[str, Node]) -> t.Dict[str, Snapshot]
                 config = self.config_for_node(node)
                 kwargs["ttl"] = config.snapshot_ttl
                 kwargs["table_naming_convention"] = config.physical_table_naming_convention
-                kwargs["virtual_environment_mode"] = config.virtual_environment_mode
 
             snapshot = Snapshot.from_node(
                 node,

sqlmesh/core/loader.py

@@ -603,6 +603,7 @@ def _load_sql_models(
                 infer_names=self.config.model_naming.infer_names,
                 signal_definitions=signals,
                 default_catalog_per_gateway=self.context.default_catalog_per_gateway,
+                virtual_environment_mode=self.config.virtual_environment_mode,
                 **loading_default_kwargs or {},
             )
 
@@ -683,6 +684,7 @@ def _load_python_models(
                             audit_definitions=audits,
                             signal_definitions=signals,
                             default_catalog_per_gateway=self.context.default_catalog_per_gateway,
+                            virtual_environment_mode=self.config.virtual_environment_mode,
                         ):
                             if model.enabled:
                                 models[model.fqn] = model

sqlmesh/core/model/decorator.py

@@ -8,6 +8,7 @@
 from sqlglot import exp
 from sqlglot.dialects.dialect import DialectType
 
+from sqlmesh.core.config.common import VirtualEnvironmentMode
 from sqlmesh.core.macros import MacroRegistry
 from sqlmesh.core.signal import SignalRegistry
 from sqlmesh.utils.jinja import JinjaMacroRegistry
@@ -154,6 +155,7 @@ def model(
         variables: t.Optional[t.Dict[str, t.Any]] = None,
         infer_names: t.Optional[bool] = False,
         blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
+        virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
     ) -> Model:
         """Get the model registered by this function."""
         env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
@@ -228,6 +230,7 @@ def model(
             "audit_definitions": audit_definitions,
             "signal_definitions": signal_definitions,
             "blueprint_variables": blueprint_variables,
+            "virtual_environment_mode": virtual_environment_mode,
             **rendered_fields,
         }
 

sqlmesh/core/model/definition.py

@@ -1062,6 +1062,7 @@ def _data_hash_values(self) -> t.List[str]:
             self.gateway,
             self.interval_unit.value if self.interval_unit is not None else None,
             str(self.optimize_query) if self.optimize_query is not None else None,
+            self.virtual_environment_mode.value,
         ]
 
         for column_name, column_type in (self.columns_to_types_ or {}).items():
@@ -2954,6 +2955,7 @@ def render_expression(
         )
     ),
     "formatting": str,
+    "virtual_environment_mode": lambda value: exp.Literal.string(value.value),
 }
 
 

sqlmesh/core/model/meta.py

@@ -10,6 +10,7 @@
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core import dialect as d
+from sqlmesh.core.config.common import VirtualEnvironmentMode
 from sqlmesh.core.config.linter import LinterConfig
 from sqlmesh.core.dialect import normalize_model_name, extract_func_call
 from sqlmesh.core.model.common import (
@@ -83,6 +84,7 @@ class ModelMeta(_Node):
         default=None, exclude=True, alias="ignored_rules"
     )
     formatting: t.Optional[bool] = Field(default=None, exclude=True)
+    virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default
 
     _bool_validator = bool_validator
     _model_kind_validator = model_kind_validator

sqlmesh/core/plan/builder.py

@@ -601,7 +601,7 @@ def _categorize_snapshots(
                 # If the model kind changes mark as breaking
                 if snapshot.is_model and snapshot.name in self._context_diff.modified_snapshots:
                     _, old = self._context_diff.modified_snapshots[snapshot.name]
-                    if _is_breaking_kind_change(old, snapshot):
+                    if _should_force_breaking_change(old, snapshot):
                         category = SnapshotChangeCategory.BREAKING
 
                 snapshot.categorize_as(category)
@@ -766,7 +766,7 @@ def _is_forward_only_change(self, s_id: SnapshotId) -> bool:
         if snapshot.name in self._context_diff.modified_snapshots:
             _, old = self._context_diff.modified_snapshots[snapshot.name]
             # If the model kind has changed in a breaking way, then we can't consider this to be a forward-only change.
-            if snapshot.is_model and _is_breaking_kind_change(old, snapshot):
+            if snapshot.is_model and _should_force_breaking_change(old, snapshot):
                 return False
         return (
             snapshot.is_model
@@ -884,7 +884,10 @@ def _modified_and_added_snapshots(self) -> t.List[Snapshot]:
         ]
 
 
-def _is_breaking_kind_change(old: Snapshot, new: Snapshot) -> bool:
+def _should_force_breaking_change(old: Snapshot, new: Snapshot) -> bool:
+    if old.virtual_environment_mode != new.virtual_environment_mode:
+        # If the virtual environment mode has changed, then it's a breaking change
+        return True
     if old.model.kind.name == new.model.kind.name:
         # If the kind hasn't changed, then it's not a breaking change
         return False

sqlmesh/core/snapshot/definition.py

@@ -337,7 +337,6 @@ class SnapshotInfoMixin(ModelKindMixin):
     base_table_name_override: t.Optional[str]
     dev_table_suffix: str
     table_naming_convention: TableNamingConvention
-    virtual_environment_mode: VirtualEnvironmentMode
 
     @cached_property
     def identifier(self) -> str:
@@ -383,6 +382,10 @@ def is_new_version(self) -> bool:
     def fully_qualified_table(self) -> t.Optional[exp.Table]:
         raise NotImplementedError
 
+    @property
+    def virtual_environment_mode(self) -> VirtualEnvironmentMode:
+        raise NotImplementedError
+
     @property
     def is_forward_only(self) -> bool:
         return self.change_category == SnapshotChangeCategory.FORWARD_ONLY
@@ -494,8 +497,10 @@ class SnapshotTableInfo(PydanticModel, SnapshotInfoMixin, frozen=True):
     custom_materialization: t.Optional[str] = None
     dev_table_suffix: str
     model_gateway: t.Optional[str] = None
-    table_naming_convention: TableNamingConvention = Field(default=TableNamingConvention.default)
-    virtual_environment_mode: VirtualEnvironmentMode = Field(default=VirtualEnvironmentMode.default)
+    table_naming_convention: TableNamingConvention = TableNamingConvention.default
+    virtual_environment_mode_: VirtualEnvironmentMode = Field(
+        default=VirtualEnvironmentMode.default, alias="virtual_environment_mode"
+    )
 
     def __lt__(self, other: SnapshotTableInfo) -> bool:
         return self.name < other.name
@@ -527,6 +532,10 @@ def table_info(self) -> SnapshotTableInfo:
         """Helper method to return self."""
         return self
 
+    @property
+    def virtual_environment_mode(self) -> VirtualEnvironmentMode:
+        return self.virtual_environment_mode_
+
     @property
     def data_version(self) -> SnapshotDataVersion:
         return SnapshotDataVersion(
@@ -624,7 +633,6 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
     table_naming_convention_: TableNamingConvention = Field(
         default=TableNamingConvention.default, alias="table_naming_convention"
     )
-    virtual_environment_mode: VirtualEnvironmentMode = Field(default=VirtualEnvironmentMode.default)
 
     @field_validator("ttl")
     @classmethod
@@ -677,7 +685,6 @@ def from_node(
         version: t.Optional[str] = None,
         cache: t.Optional[t.Dict[str, SnapshotFingerprint]] = None,
         table_naming_convention: TableNamingConvention = TableNamingConvention.default,
-        virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
     ) -> Snapshot:
         """Creates a new snapshot for a node.
 
@@ -689,7 +696,6 @@ def from_node(
             version: The version that a snapshot is associated with. Usually set during the planning phase.
             cache: Cache of node name to fingerprints.
             table_naming_convention: Convention to follow when generating the physical table name
-            virtual_environment_mode: Mode for handling virtual environments
 
         Returns:
             The newly created snapshot.
@@ -722,7 +728,6 @@ def from_node(
             ttl=ttl,
             version=version,
             table_naming_convention=table_naming_convention,
-            virtual_environment_mode=virtual_environment_mode,
         )
 
     def __eq__(self, other: t.Any) -> bool:
@@ -1384,6 +1389,12 @@ def custom_materialization(self) -> t.Optional[str]:
             return t.cast(CustomKind, self.model.kind).materialization
         return None
 
+    @property
+    def virtual_environment_mode(self) -> VirtualEnvironmentMode:
+        return (
+            self.model.virtual_environment_mode if self.is_model else VirtualEnvironmentMode.default
+        )
+
     def _ensure_categorized(self) -> None:
         if not self.change_category:
             raise SQLMeshError(f"Snapshot {self.snapshot_id} has not been categorized yet.")

tests/core/test_model.py

@@ -1907,7 +1907,8 @@ def test_render_definition_with_defaults():
             dialect spark,
             kind VIEW (
                 materialized FALSE
-            )
+            ),
+            virtual_environment_mode 'full'
         );
 
         {query}
@@ -5730,7 +5731,7 @@ def test_default_catalog_sql(assert_exp_eq):
     The system is not designed to actually support having an engine that doesn't support default catalog
     to start supporting it or the reverse of that. If that did happen then bugs would occur.
     """
-    HASH_WITH_CATALOG = "516937963"
+    HASH_WITH_CATALOG = "1269513823"
 
     # Test setting default catalog doesn't change hash if it matches existing logic
     expressions = d.parse(
@@ -5896,7 +5897,7 @@ def test_default_catalog_sql(assert_exp_eq):
 
 
 def test_default_catalog_python():
-    HASH_WITH_CATALOG = "770057346"
+    HASH_WITH_CATALOG = "2728996410"
 
     @model(name="db.table", kind="full", columns={'"COL"': "int"})
     def my_model(context, **kwargs):
@@ -5988,7 +5989,7 @@ def test_default_catalog_external_model():
     Since external models fqns are the only thing affected by default catalog, and when they change new snapshots
     are made, the hash will be the same across different names.
     """
-    EXPECTED_HASH = "3614876346"
+    EXPECTED_HASH = "763256265"
 
     model = create_external_model("db.table", columns={"a": "int", "limit": "int"})
     assert model.default_catalog is None

tests/core/test_plan_stages.py

@@ -1399,19 +1399,19 @@ def test_build_plan_stages_virtual_environment_mode_filtering(
             name="full_model",
             query=parse_one("select 1, ds"),
             kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+            virtual_environment_mode=VirtualEnvironmentMode.FULL,
         )
     )
-    snapshot_full.virtual_environment_mode = VirtualEnvironmentMode.FULL
     snapshot_full.categorize_as(SnapshotChangeCategory.BREAKING)
 
     snapshot_dev_only = make_snapshot(
         SqlModel(
             name="dev_only_model",
             query=parse_one("select 2, ds"),
             kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+            virtual_environment_mode=VirtualEnvironmentMode.DEV_ONLY,
         )
     )
-    snapshot_dev_only.virtual_environment_mode = VirtualEnvironmentMode.DEV_ONLY
     snapshot_dev_only.categorize_as(SnapshotChangeCategory.BREAKING)
 
     # Mock state reader
@@ -1593,9 +1593,9 @@ def test_build_plan_stages_virtual_environment_mode_no_updates(
             name="dev_only_model",
             query=parse_one("select 1, ds"),
             kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+            virtual_environment_mode=VirtualEnvironmentMode.DEV_ONLY,
         )
     )
-    snapshot_dev_only.virtual_environment_mode = VirtualEnvironmentMode.DEV_ONLY
     snapshot_dev_only.categorize_as(SnapshotChangeCategory.BREAKING)
 
     # Mock state reader