Fix: Include unexpired downstream views when cleaning up expired tables

Author: erindru

sqlmesh/core/state_sync/db/snapshot.py

@@ -30,6 +30,7 @@
     SnapshotId,
     SnapshotFingerprint,
     SnapshotChangeCategory,
+    snapshots_to_dag,
 )
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 from sqlmesh.utils.date import now_timestamp, TimeLike, now, to_timestamp
@@ -224,18 +225,51 @@ def _get_expired_snapshots(
 
         if not ignore_ttl:
             expired_query = expired_query.where(
-                (exp.column("updated_ts") + exp.column("ttl_ms")) <= current_ts
+                ((exp.column("updated_ts") + exp.column("ttl_ms")) <= current_ts)
+                # we need to include views even if they havent expired in case one depends on a table that /has/ expired
+                .or_(exp.column("kind_name").eq(ModelKindName.VIEW))
             )
 
-        expired_candidates = {
+        candidates = {
             SnapshotId(name=name, identifier=identifier): SnapshotNameVersion(
                 name=name, version=version
             )
             for name, identifier, version in fetchall(self.engine_adapter, expired_query)
         }
-        if not expired_candidates:
+
+        if not candidates:
             return set(), []
 
+        expired_candidates: t.Dict[SnapshotId, SnapshotNameVersion] = {}
+
+        if ignore_ttl:
+            expired_candidates = candidates
+        else:
+            # Fetch full snapshots because we need to build a dependency tree
+            full_candidates = self.get_snapshots(candidates.keys())
+
+            # Build DAG so we can check if any views that are not expired depend on a table snapshot that is expired
+            dag = snapshots_to_dag(full_candidates.values())
+
+            # remove any non-expired views that dont depend on expired tables
+            for snapshot_id in dag.reversed:
+                snapshot = full_candidates.get(snapshot_id, None)
+                if not snapshot:
+                    continue
+
+                if snapshot.expiration_ts <= current_ts:
+                    # All expired snapshots should be included
+                    expired_candidates[snapshot.snapshot_id] = snapshot.name_version
+                elif snapshot.model_kind_name == ModelKindName.VIEW:
+                    # With non-expired views, check if they have any expired upstream tables
+                    if any(
+                        full_candidates[parent_id].expiration_ts <= current_ts
+                        for parent_id in dag.upstream(snapshot_id)
+                        if full_candidates[parent_id].model_kind_name != ModelKindName.VIEW
+                    ):
+                        # an upstream table has expired, therefore this view is no longer valid and needs to be cleaned up as well
+                        expired_candidates[snapshot.snapshot_id] = snapshot.name_version
+
         promoted_snapshot_ids = {
             snapshot.snapshot_id
             for environment in environments
@@ -253,37 +287,39 @@ def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
             unique_expired_versions, batch_size=self.SNAPSHOT_BATCH_SIZE
         )
         cleanup_targets = []
-        expired_snapshot_ids = set()
+        expired_sv_snapshot_ids = set()
         for versions_batch in version_batches:
-            snapshots = self._get_snapshots_with_same_version(versions_batch)
+            sv_snapshots = self._get_snapshots_with_same_version(versions_batch)
 
             snapshots_by_version = defaultdict(set)
             snapshots_by_dev_version = defaultdict(set)
-            for s in snapshots:
+            for s in sv_snapshots:
                 snapshots_by_version[(s.name, s.version)].add(s.snapshot_id)
                 snapshots_by_dev_version[(s.name, s.dev_version)].add(s.snapshot_id)
 
-            expired_snapshots = [s for s in snapshots if not _is_snapshot_used(s)]
-            expired_snapshot_ids.update([s.snapshot_id for s in expired_snapshots])
+            expired_sv_snapshots = [s for s in sv_snapshots if not _is_snapshot_used(s)]
+            expired_sv_snapshot_ids.update([s.snapshot_id for s in expired_sv_snapshots])
 
-            for snapshot in expired_snapshots:
-                shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
-                shared_version_snapshots.discard(snapshot.snapshot_id)
+            for sv_snapshot in expired_sv_snapshots:
+                shared_version_snapshots = snapshots_by_version[
+                    (sv_snapshot.name, sv_snapshot.version)
+                ]
+                shared_version_snapshots.discard(sv_snapshot.snapshot_id)
 
                 shared_dev_version_snapshots = snapshots_by_dev_version[
-                    (snapshot.name, snapshot.dev_version)
+                    (sv_snapshot.name, sv_snapshot.dev_version)
                 ]
-                shared_dev_version_snapshots.discard(snapshot.snapshot_id)
+                shared_dev_version_snapshots.discard(sv_snapshot.snapshot_id)
 
                 if not shared_dev_version_snapshots:
                     cleanup_targets.append(
                         SnapshotTableCleanupTask(
-                            snapshot=snapshot.full_snapshot.table_info,
+                            snapshot=sv_snapshot.full_snapshot.table_info,
                             dev_table_only=bool(shared_version_snapshots),
                         )
                     )
 
-        return expired_snapshot_ids, cleanup_targets
+        return expired_sv_snapshot_ids, cleanup_targets
 
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         """Deletes snapshots.

tests/core/engine_adapter/integration/test_integration.py

@@ -13,6 +13,7 @@
 import pandas as pd  # noqa: TID253
 import pytest
 import pytz
+import time_machine
 from sqlglot import exp, parse_one
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
@@ -28,8 +29,8 @@
 from sqlmesh.core.model.definition import create_sql_model
 from sqlmesh.core.plan import Plan
 from sqlmesh.core.state_sync.db import EngineAdapterStateSync
-from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory
-from sqlmesh.utils.date import now, to_date, to_time_column
+from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory, SnapshotId
+from sqlmesh.utils.date import now, to_date, to_time_column, to_ds
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
@@ -2672,3 +2673,166 @@ def test_identifier_length_limit(ctx: TestContext):
         match=re.escape(match),
     ):
         adapter.create_table(long_table_name, {"col": exp.DataType.build("int")})
+
+
+def test_janitor_out_of_order_drop(ctx: TestContext, tmp_path: pathlib.Path):
+    """
+    Scenario:
+
+    Ensure that cleaning up expired table snapshots also cleans up any unexpired view snapshots that depend on them
+
+    - We create a A (table) <- B (view)
+    - In dev, we modify A - triggers new version of A and a dev preview of B that both expire in 7 days
+    - We advance time by 3 days
+    - In dev, we modify B - triggers a new version of B that depends on A but expires 3 days after A
+    - We advance time by 5 days so that A has reached its expiry but B has not
+    - We expire dev so that none of these snapshots are promoted and are thus targets for cleanup
+    - We run the janitor
+
+    Expected outcome:
+        - All the dev versions of A and B should be dropped
+        - We should not get a 'ERROR: cannot drop table x because other objects depend on it' on engines that do schema binding
+    """
+
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+    schema = exp.to_table(ctx.schema(TEST_SCHEMA)).this
+
+    (models_dir / "model_a.sql").write_text(f"""
+    MODEL (
+        name {schema}.model_a,
+        kind FULL
+    );
+
+    SELECT 1 as a, 2 as b;
+    """)
+
+    (models_dir / "model_b.sql").write_text(f"""
+    MODEL (
+        name {schema}.model_b,
+        kind VIEW
+    );
+
+    SELECT a from {schema}.model_a;
+    """)
+
+    def _mutate_config(gateway: str, config: Config):
+        # helps with debugging when you can see the SQLMesh state right next to the model tables
+        config.gateways[gateway].state_schema = f"{schema}"
+
+    with time_machine.travel("2020-01-01 00:00:00"):
+        sqlmesh = ctx.create_context(
+            path=tmp_path, config_mutator=_mutate_config, ephemeral_state_connection=False
+        )
+        sqlmesh.plan(auto_apply=True)
+
+    model_a_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_a" in n)
+    # expiry is last updated + ttl
+    assert timedelta(milliseconds=model_a_snapshot.ttl_ms) == timedelta(weeks=1)
+    assert to_ds(model_a_snapshot.updated_ts) == "2020-01-01"
+    assert to_ds(model_a_snapshot.expiration_ts) == "2020-01-08"
+
+    model_b_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_b" in n)
+    assert timedelta(milliseconds=model_b_snapshot.ttl_ms) == timedelta(weeks=1)
+    assert to_ds(model_b_snapshot.updated_ts) == "2020-01-01"
+    assert to_ds(model_b_snapshot.expiration_ts) == "2020-01-08"
+
+    model_a_prod_snapshot = model_a_snapshot
+    model_b_prod_snapshot = model_b_snapshot
+
+    # move forward 1 days
+    # new dev environment - touch models to create new snapshots
+    # model a / b expiry in prod should remain unmodified
+    # model a / b expiry in dev should be as at today
+    with time_machine.travel("2020-01-02 00:00:00"):
+        (models_dir / "model_a.sql").write_text(f"""
+        MODEL (
+            name {schema}.model_a,
+            kind FULL
+        );
+
+        SELECT 1 as a, 2 as b, 3 as c;
+        """)
+
+        sqlmesh = ctx.create_context(
+            path=tmp_path, config_mutator=_mutate_config, ephemeral_state_connection=False
+        )
+        sqlmesh.plan(environment="dev", auto_apply=True)
+
+        # should now have 4 snapshots in state - 2x model a and 2x model b
+        # the new model b is a dev preview because its upstream model changed
+        assert len(sqlmesh.engine_adapter.fetchall(f"select * from {schema}._snapshots")) == 4
+
+        # context just has the two latest
+        assert len(sqlmesh.snapshots) == 2
+
+        # these expire 1 day later than what's in prod
+        model_a_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_a" in n)
+        assert timedelta(milliseconds=model_a_snapshot.ttl_ms) == timedelta(weeks=1)
+        assert to_ds(model_a_snapshot.updated_ts) == "2020-01-02"
+        assert to_ds(model_a_snapshot.expiration_ts) == "2020-01-09"
+
+        model_b_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_b" in n)
+        assert timedelta(milliseconds=model_b_snapshot.ttl_ms) == timedelta(weeks=1)
+        assert to_ds(model_b_snapshot.updated_ts) == "2020-01-02"
+        assert to_ds(model_b_snapshot.expiration_ts) == "2020-01-09"
+
+    # move forward 3 days
+    # touch model b in dev but leave model a
+    # this bumps the model b expiry but model a remains unchanged, so will expire before model b even though model b depends on it
+    with time_machine.travel("2020-01-05 00:00:00"):
+        (models_dir / "model_b.sql").write_text(f"""
+        MODEL (
+            name {schema}.model_b,
+            kind VIEW
+        );
+
+        SELECT a, 'b' as b from {schema}.model_a;
+        """)
+
+        sqlmesh = ctx.create_context(
+            path=tmp_path, config_mutator=_mutate_config, ephemeral_state_connection=False
+        )
+        sqlmesh.plan(environment="dev", auto_apply=True)
+
+        # should now have 5 snapshots in state - 2x model a and 3x model b
+        assert len(sqlmesh.engine_adapter.fetchall(f"select * from {schema}._snapshots")) == 5
+
+        # context just has the two latest
+        assert len(sqlmesh.snapshots) == 2
+
+        # model a expiry should not have changed
+        model_a_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_a" in n)
+        assert timedelta(milliseconds=model_a_snapshot.ttl_ms) == timedelta(weeks=1)
+        assert to_ds(model_a_snapshot.updated_ts) == "2020-01-02"
+        assert to_ds(model_a_snapshot.expiration_ts) == "2020-01-09"
+
+        # model b should now expire well after model a
+        model_b_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_b" in n)
+        assert timedelta(milliseconds=model_b_snapshot.ttl_ms) == timedelta(weeks=1)
+        assert to_ds(model_b_snapshot.updated_ts) == "2020-01-05"
+        assert to_ds(model_b_snapshot.expiration_ts) == "2020-01-12"
+
+    # move forward to date where after model a has expired but before model b has expired
+    # invalidate dev to trigger cleanups
+    # run janitor. model a is expired so will be cleaned up and this will cascade to model b.
+    with time_machine.travel("2020-01-10 00:00:00"):
+        sqlmesh = ctx.create_context(
+            path=tmp_path, config_mutator=_mutate_config, ephemeral_state_connection=False
+        )
+
+        before_snapshots = sqlmesh.engine_adapter.fetchall(
+            f"select name, identifier from {schema}._snapshots"
+        )
+        sqlmesh.invalidate_environment("dev")
+        sqlmesh.run_janitor(ignore_ttl=False)
+        after_snapshots = sqlmesh.engine_adapter.fetchall(
+            f"select name, identifier from {schema}._snapshots"
+        )
+
+        assert len(before_snapshots) != len(after_snapshots)
+
+        # all that's left should be the two snapshots that were in prod
+        assert set(
+            [SnapshotId(name=name, identifier=identifier) for name, identifier in after_snapshots]
+        ) == set([model_a_prod_snapshot.snapshot_id, model_b_prod_snapshot.snapshot_id])

tests/core/state_sync/test_state_sync.py

@@ -1124,6 +1124,86 @@ def test_delete_expired_environments(state_sync: EngineAdapterStateSync, make_sn
     assert state_sync.get_environment_statements(env_a.name) == []
 
 
+def test_get_expired_snapshots_includes_downstream_view_snapshots(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable[..., Snapshot]
+):
+    now_ts = now_timestamp()
+
+    # model_a: table snapshot
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            kind="FULL",
+            query=parse_one("select a, ds"),
+        ),
+    )
+    snapshot_a.ttl = "in 10 seconds"
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_a.updated_ts = now_ts - 15000  # now - 15 seconds = expired
+
+    # model_b: view snapshot that depends on model_a table snapshot
+    # in an actual scenario, this could have been created a few days later so expires a few days after the snapshot it depends on
+    # unlike a table, a view ceases to be valid if the upstream table it points to is dropped
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="b",
+            kind="VIEW",
+            query=parse_one("select *, 'foo' as model_b from a"),
+        ),
+        nodes={"a": snapshot_a.model},
+    )
+    snapshot_b.ttl = "in 10 seconds"
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.updated_ts = now_ts  # now = not expired
+
+    # model_c: table snapshot that depends on model_a table snapshot but since its a table it will still work if model_a is dropped
+    # so should not be considered when cleaning up expired snapshots (as it has not expired)
+    snapshot_c = make_snapshot(
+        SqlModel(
+            name="c",
+            kind="FULL",
+            query=parse_one("select *, 'foo' as model_c from a"),
+        ),
+        nodes={"a": snapshot_a.model},
+    )
+    snapshot_c.ttl = "in 10 seconds"
+    snapshot_c.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_c.updated_ts = now_ts  # now = not expired
+
+    # model_d: view snapshot with no dependency on model a, so should not be dropped if model a is dropped
+    snapshot_d = make_snapshot(
+        SqlModel(
+            name="d",
+            kind="VIEW",
+            query=parse_one("select 'model_d' as d"),
+        ),
+    )
+    snapshot_d.ttl = "in 10 seconds"
+    snapshot_d.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_d.updated_ts = now_ts  # now = not expired
+
+    state_sync.push_snapshots([snapshot_a, snapshot_b, snapshot_c, snapshot_d])
+    cleanup_targets = state_sync.get_expired_snapshots(now_ts)
+
+    # snapshot_a should be cleaned up because it's expired
+    # snapshot_b is unexpired but should be cleaned up because snapshot_a is being cleaned up
+    #   - since it's a view, it will stop working if the snapshot_a table it's pointing to is dropped, so thats why it needs to be dropped too
+    #   - additionally, if it isnt dropped, some engines will throw "cannot drop because other objects depend on it" when trying to drop the snapshot_a table
+    # snapshot_c is unexpired and should not be cleaned up even though its upstream snapshot_a has expired, since it's a table it will still continue to work
+    # snapshot_d is unexpired and should also not be cleaned up. While it's a view, it doesnt depend on anything that *has* expired, so is not a target for cleanup
+    assert len(cleanup_targets) == 2
+
+    snapshot_a_cleanup = next(
+        (t for t in cleanup_targets if t.snapshot.snapshot_id == snapshot_a.snapshot_id), None
+    )
+    assert snapshot_a_cleanup
+
+    snapshot_b_cleanup = next(
+        (t for t in cleanup_targets if t.snapshot.snapshot_id == snapshot_b.snapshot_id), None
+    )
+    assert snapshot_b_cleanup
+
+
 def test_delete_expired_snapshots(state_sync: EngineAdapterStateSync, make_snapshot: t.Callable):
     now_ts = now_timestamp()