Repository: TobikoData/sqlmesh
Pull request #3059 (2024-08-27 17:44:57)

Fix: Ignore expired environments in the selector
Otherwise the plan produces unexpected diffs since expired environments are ignored in `context_diff`.


Patch:
diff --git a/sqlmesh/core/selector.py b/sqlmesh/core/selector.py
index 46335be8..7176c7cf 100644
--- a/sqlmesh/core/selector.py
+++ b/sqlmesh/core/selector.py
@@ -68,6 +68,9 @@ class Selector:
             A dictionary of models.
         """
         target_env = self._state_reader.get_environment(Environment.sanitize_name(target_env_name))
+        if target_env and target_env.expired:
+            target_env = None
+
         if not target_env and fallback_env_name:
             target_env = self._state_reader.get_environment(
                 Environment.sanitize_name(fallback_env_name)
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 90514943..31ab318d 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -769,25 +769,14 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         Returns:
             A [start, end) pair.
         """
-        interval_unit = self.node.interval_unit
-        start_ts = to_timestamp(interval_unit.cron_floor(start))
-        if start_ts < to_timestamp(start) and not self.model.allow_partials:
-            start_ts = to_timestamp(interval_unit.cron_next(start_ts))
-
-        if is_date(end):
-            end = to_datetime(end) + timedelta(days=1)
-        end_ts = to_timestamp(interval_unit.cron_floor(end) if not allow_partial else end)
-        if end_ts < start_ts and to_timestamp(end) > to_timestamp(start) and not strict:
-            # This can happen when the interval unit is coarser than the size of the input interval.
-            # For example, if the interval unit is monthly, but the input interval is only 1 hour long.
-            return (start_ts, end_ts)
-
-        if (strict and start_ts >= end_ts) or (start_ts > end_ts):
-            raise ValueError(
-                f"`end` ({to_datetime(end_ts)}) must be greater than `start` ({to_datetime(start_ts)})"
-            )
-
-        return (start_ts, end_ts)
+        return inclusive_exclusive(
+            start,
+            end,
+            self.node.interval_unit,
+            model_allow_partials=self.is_model and self.model.allow_partials,
+            strict=strict,
+            allow_partial=allow_partial,
+        )
 
     def merge_intervals(self, other: t.Union[Snapshot, SnapshotIntervals]) -> None:
         """Inherits intervals from the target snapshot.
@@ -1740,6 +1729,48 @@ def compute_missing_intervals(
     return sorted(missing)
 
 
+@lru_cache(maxsize=None)
+def inclusive_exclusive(
+    start: TimeLike,
+    end: TimeLike,
+    interval_unit: IntervalUnit,
+    model_allow_partials: bool,
+    strict: bool = True,
+    allow_partial: bool = False,
+) -> Interval:
+    """Transform the inclusive start and end into a [start, end) pair.
+
+    Args:
+        start: The start date/time of the interval (inclusive)
+        end: The end date/time of the interval (inclusive)
+        interval_unit: The interval unit.
+        model_allow_partials: Whether or not the model allows partials.
+        strict: Whether to fail when the inclusive start is the same as the exclusive end.
+        allow_partial: Whether the interval can be partial or not.
+
+    Returns:
+        A [start, end) pair.
+    """
+    start_ts = to_timestamp(interval_unit.cron_floor(start))
+    if start_ts < to_timestamp(start) and not model_allow_partials:
+        start_ts = to_timestamp(interval_unit.cron_next(start_ts))
+
+    if is_date(end):
+        end = to_datetime(end) + timedelta(days=1)
+    end_ts = to_timestamp(interval_unit.cron_floor(end) if not allow_partial else end)
+    if end_ts < start_ts and to_timestamp(end) > to_timestamp(start) and not strict:
+        # This can happen when the interval unit is coarser than the size of the input interval.
+        # For example, if the interval unit is monthly, but the input interval is only 1 hour long.
+        return (start_ts, end_ts)
+
+    if (strict and start_ts >= end_ts) or (start_ts > end_ts):
+        raise ValueError(
+            f"`end` ({to_datetime(end_ts)}) must be greater than `start` ({to_datetime(start_ts)})"
+        )
+
+    return (start_ts, end_ts)
+
+
 def earliest_start_date(
     snapshots: t.Collection[Snapshot],
     cache: t.Optional[t.Dict[str, datetime]] = None,
diff --git a/sqlmesh/utils/metaprogramming.py b/sqlmesh/utils/metaprogramming.py
index 31743c7a..2da34688 100644
--- a/sqlmesh/utils/metaprogramming.py
+++ b/sqlmesh/utils/metaprogramming.py
@@ -359,11 +359,12 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
             name = v.__name__
             name = k if name == "<lambda>" else name
 
-            # We can't call getfile on built-in callables
+            # getfile raises a `TypeError` for built-in modules, classes, or functions
             # https://docs.python.org/3/library/inspect.html#inspect.getfile
-            file_path = (
-                Path(inspect.getfile(inspect.unwrap(v))) if not inspect.isbuiltin(v) else None
-            )
+            try:
+                file_path = Path(inspect.getfile(v))
+            except TypeError:
+                file_path = None
 
             if _is_relative_to(file_path, path):
                 serialized[k] = Executable(


Test patch:
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index 59c35ced..3ee4b3d6 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -14,6 +14,7 @@ from sqlmesh.core.model import Model, SqlModel
 from sqlmesh.core.selector import Selector
 from sqlmesh.core.snapshot import SnapshotChangeCategory
 from sqlmesh.utils import UniqueKeyDict
+from sqlmesh.utils.date import now_timestamp
 
 
 @pytest.mark.parametrize(
@@ -185,6 +186,81 @@ def test_select_models(mocker: MockerFixture, make_snapshot, default_catalog: t.
     )
 
 
+def test_select_models_expired_environment(mocker: MockerFixture, make_snapshot):
+    modified_model_v1 = SqlModel(
+        name="db.modified_model",
+        query=d.parse_one("SELECT a + 1 FROM db.added_model"),
+    )
+    modified_model_v2 = SqlModel(
+        name="db.modified_model",
+        query=d.parse_one("SELECT a + 2 FROM db.added_model"),
+    )
+    removed_model = SqlModel(
+        name="db.removed_model",
+        query=d.parse_one("SELECT a FROM db.added_model"),
+    )
+    standalone_audit = StandaloneAudit(
+        name="test_audit", query=d.parse_one("SELECT * FROM added_model WHERE a IS NULL")
+    )
+
+    modified_model_v1_snapshot = make_snapshot(modified_model_v1)
+    modified_model_v1_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    removed_model_snapshot = make_snapshot(removed_model)
+    removed_model_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    standalone_audit_snapshot = make_snapshot(standalone_audit)
+    standalone_audit_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    prod_env = Environment(
+        name="prod",
+        snapshots=[modified_model_v1_snapshot.table_info],
+        start_at="2023-01-01",
+        end_at="2023-02-01",
+        plan_id="test_plan_id",
+    )
+
+    env_name = "test_env"
+    dev_env = Environment(
+        name=env_name,
+        snapshots=[modified_model_v1_snapshot.table_info, removed_model_snapshot.table_info],
+        start_at="2023-01-01",
+        end_at="2023-02-01",
+        plan_id="test_plan_id",
+    )
+
+    state_reader_mock = mocker.Mock()
+    state_reader_mock.get_environment.side_effect = (
+        lambda name: prod_env if name == "prod" else dev_env
+    )
+
+    all_snapshots = {
+        modified_model_v1_snapshot.snapshot_id: modified_model_v1_snapshot,
+        removed_model_snapshot.snapshot_id: removed_model_snapshot,
+    }
+    state_reader_mock.get_snapshots.side_effect = lambda snapshots: {
+        s.snapshot_id: all_snapshots[s.snapshot_id] for s in snapshots
+    }
+
+    local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+    local_models[modified_model_v2.fqn] = modified_model_v2
+    selector = Selector(state_reader_mock, local_models)
+
+    _assert_models_equal(
+        selector.select_models(["*.modified_model"], env_name, fallback_env_name="prod"),
+        {
+            removed_model.fqn: removed_model,
+            modified_model_v2.fqn: modified_model_v2,
+        },
+    )
+
+    dev_env.expiration_ts = now_timestamp() - 1
+    _assert_models_equal(
+        selector.select_models(["*.modified_model"], env_name, fallback_env_name="prod"),
+        {
+            modified_model_v2.fqn: modified_model_v2,
+        },
+    )
+
+
 def test_select_change_schema(mocker: MockerFixture, make_snapshot):
     parent = SqlModel(
         name="db.parent",
diff --git a/tests/utils/test_metaprogramming.py b/tests/utils/test_metaprogramming.py
index 8537851a..be360f9e 100644
--- a/tests/utils/test_metaprogramming.py
+++ b/tests/utils/test_metaprogramming.py
@@ -1,4 +1,5 @@
 import typing as t
+from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -45,7 +46,7 @@ def test_print_exception(mocker: MockerFixture):
 
     expected_message = f"""Traceback (most recent call last):
 
-  File "{__file__}", line 42, in test_print_exception
+  File "{__file__}", line 43, in test_print_exception
     eval("test_fun()", env)
 
   File "<string>", line 1, in <module>
@@ -106,6 +107,11 @@ def noop_metadata() -> None:
 setattr(noop_metadata, c.SQLMESH_METADATA, True)
 
 
+@contextmanager
+def test_context_manager():
+    yield
+
+
 def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2) -> int:
     """DOC STRING"""
     sqlglot.parse_one("1")
@@ -117,6 +123,9 @@ def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2)
     def closure(z: int) -> int:
         return z + Z
 
+    with test_context_manager():
+        pass
+
     return closure(y) + other_func(Y)
 
 
@@ -132,6 +141,7 @@ def test_func_globals() -> None:
         "sqlglot": sqlglot,
         "exp": exp,
         "expressions": exp,
+        "test_context_manager": test_context_manager,
     }
     assert func_globals(other_func) == {
         "X": 1,
@@ -167,6 +177,8 @@ def test_normalize_source() -> None:
 
     def closure(z: int):
         return z + Z
+    with test_context_manager():
+        pass
     return closure(y) + other_func(Y)"""
     )
 
@@ -194,6 +206,7 @@ def test_serialize_env() -> None:
     build_env(main_func, env=env, name="MAIN", path=path)
     env = serialize_env(env, path=path)  # type: ignore
 
+    assert prepare_env(env)
     assert env == {
         "MAIN": Executable(
             name="main_func",
@@ -209,11 +222,19 @@ def test_serialize_env() -> None:
 
     def closure(z: int):
         return z + Z
+    with test_context_manager():
+        pass
     return closure(y) + other_func(Y)""",
         ),
         "X": Executable(payload="1", kind=ExecutableKind.VALUE),
         "Y": Executable(payload="2", kind=ExecutableKind.VALUE),
         "Z": Executable(payload="3", kind=ExecutableKind.VALUE),
+        "_GeneratorContextManager": Executable(
+            payload="from contextlib import _GeneratorContextManager", kind=ExecutableKind.IMPORT
+        ),
+        "contextmanager": Executable(
+            payload="from contextlib import contextmanager", kind=ExecutableKind.IMPORT
+        ),
         "KLASS_X": Executable(payload="1", kind=ExecutableKind.VALUE),
         "KLASS_Y": Executable(payload="2", kind=ExecutableKind.VALUE),
         "KLASS_Z": Executable(payload="3", kind=ExecutableKind.VALUE),
@@ -255,6 +276,14 @@ class DataClass:
         "expressions": Executable(
             kind=ExecutableKind.IMPORT, payload="import sqlglot.expressions as expressions"
         ),
+        "func": Executable(
+            payload="""@contextmanager
+def test_context_manager():
+    yield""",
+            name="test_context_manager",
+            path="test_metaprogramming.py",
+            alias="func",
+        ),
         "my_lambda": Executable(
             name="my_lambda",
             path="test_metaprogramming.py",
@@ -282,4 +311,9 @@ class DataClass:
     my_lambda()
     return X + a""",
         ),
+        "test_context_manager": Executable(
+            payload="from tests.utils.test_metaprogramming import test_context_manager",
+            kind=ExecutableKind.IMPORT,
+        ),
+        "wraps": Executable(payload="from functools import wraps", kind=ExecutableKind.IMPORT),
     }