PR feedback

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,11 +225,20 @@ def _get_expired_snapshots(
 
         if not ignore_ttl:
             expired_query = expired_query.where(
-                ((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))
+                (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 or view that /has/ expired
+            # but we only need to do this if there are expired objects to begin with
+
+            expired_record_count = fetchone(
+                self.engine_adapter, exp.select("count(*)").from_(expired_query.subquery())
+            )
+            if expired_record_count and expired_record_count[0]:
+                expired_query = t.cast(
+                    exp.Select, expired_query.or_(exp.column("kind_name").eq(ModelKindName.VIEW))
+                )
+
         candidates = {
             SnapshotId(name=name, identifier=identifier): SnapshotNameVersion(
                 name=name, version=version
@@ -247,29 +257,23 @@ def _get_expired_snapshots(
             # Fetch full snapshots because we need the dependency relationship
             full_candidates = self.get_snapshots(candidates.keys())
 
+            dag = snapshots_to_dag(full_candidates.values())
+
             # Include any non-expired views that depend on expired tables
-            for snapshot_id in full_candidates:
+            for snapshot_id in dag:
                 snapshot = full_candidates.get(snapshot_id, None)
+
                 if not snapshot:
                     continue
 
                 if snapshot.expiration_ts <= current_ts:
-                    # All expired snapshots should be included
+                    # All expired snapshots should be included regardless
                     expired_candidates[snapshot.snapshot_id] = snapshot.name_version
                 elif snapshot.model_kind_name == ModelKindName.VIEW:
-                    # With non-expired views, check if they point to an expired parent table
-                    immediate_parents = [
-                        full_candidates[parent_id]
-                        for parent_id in snapshot.parents
-                        if parent_id in full_candidates
-                    ]
-
-                    if any(
-                        parent.expiration_ts <= current_ts
-                        for parent in immediate_parents
-                        if parent.model_kind_name != ModelKindName.VIEW
-                    ):
-                        # an immediate upstream table has expired, therefore this view is no longer valid and needs to be cleaned up as well
+                    # Check if any of our parents are in the expired list
+                    # This works because we traverse the dag in topological order, so if our parent either directly
+                    # expired or indirectly expired because /its/ parent expired, it will still be in the expired_snapshots list
+                    if any(parent.snapshot_id in expired_candidates for parent in snapshot.parents):
                         expired_candidates[snapshot.snapshot_id] = snapshot.name_version
 
         promoted_snapshot_ids = {

tests/core/engine_adapter/integration/test_integration.py

@@ -24,7 +24,6 @@
 import sqlmesh.core.dialect as d
 from sqlmesh.core.dialect import select_from_values
 from sqlmesh.core.model import Model, load_sql_based_model
-from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.core.engine_adapter.mixins import RowDiffMixin, LogicalMergeMixin
 from sqlmesh.core.model.definition import create_sql_model
@@ -2816,19 +2815,28 @@ def test_janitor_drops_downstream_unexpired_hard_dependencies(
     - 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
+    - In dev, we create B(view) <- C(view) and B(view) <- D(table)
+    - We advance time by 5 days so that A has reached its expiry but B, C and D have 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
+        - C should be dropped as well because it's a view that depends on B which was dropped
+        - D should not be dropped because while it depends on B which was dropped, it's a table so is still valid after B is dropped
         - We should not get a 'ERROR: cannot drop table x because other objects depend on it' on engines that do schema binding
     """
 
-    def _state_sync_engine_adapter(context: Context) -> EngineAdapter:
+    def _all_snapshot_ids(context: Context) -> t.List[SnapshotId]:
         assert isinstance(context.state_sync, CachingStateSync)
         assert isinstance(context.state_sync.state_sync, EngineAdapterStateSync)
-        return context.state_sync.state_sync.engine_adapter
+
+        return [
+            SnapshotId(name=name, identifier=identifier)
+            for name, identifier in context.state_sync.state_sync.engine_adapter.fetchall(
+                "select name, identifier from sqlmesh._snapshots"
+            )
+        ]
 
     models_dir = tmp_path / "models"
     models_dir.mkdir()
@@ -2898,10 +2906,10 @@ def _mutate_config(gateway: str, config: Config):
 
         # 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(_state_sync_engine_adapter(sqlmesh).fetchall(f"select * from sqlmesh._snapshots"))
-            == 4
-        )
+        all_snapshot_ids = _all_snapshot_ids(sqlmesh)
+        assert len(all_snapshot_ids) == 4
+        assert len([s for s in all_snapshot_ids if "model_a" in s.name]) == 2
+        assert len([s for s in all_snapshot_ids if "model_b" in s.name]) == 2
 
         # context just has the two latest
         assert len(sqlmesh.snapshots) == 2
@@ -2930,19 +2938,39 @@ def _mutate_config(gateway: str, config: Config):
         SELECT a, 'b' as b from {schema}.model_a;
         """)
 
+        (models_dir / "model_c.sql").write_text(f"""
+        MODEL (
+            name {schema}.model_c,
+            kind VIEW
+        );
+
+        SELECT a, 'c' as c from {schema}.model_b;
+        """)
+
+        (models_dir / "model_d.sql").write_text(f"""
+        MODEL (
+            name {schema}.model_d,
+            kind FULL
+        );
+
+        SELECT a, 'd' as d from {schema}.model_b;
+        """)
+
         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(_state_sync_engine_adapter(sqlmesh).fetchall(f"select * from sqlmesh._snapshots"))
-            == 5
-        )
+        # should now have 7 snapshots in state - 2x model a, 3x model b, 1x model c and 1x model d
+        all_snapshot_ids = _all_snapshot_ids(sqlmesh)
+        assert len(all_snapshot_ids) == 7
+        assert len([s for s in all_snapshot_ids if "model_a" in s.name]) == 2
+        assert len([s for s in all_snapshot_ids if "model_b" in s.name]) == 3
+        assert len([s for s in all_snapshot_ids if "model_c" in s.name]) == 1
+        assert len([s for s in all_snapshot_ids if "model_d" in s.name]) == 1
 
-        # context just has the two latest
-        assert len(sqlmesh.snapshots) == 2
+        # context just has the 4 latest
+        assert len(sqlmesh.snapshots) == 4
 
         # model a expiry should not have changed
         model_a_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_a" in n)
@@ -2956,26 +2984,41 @@ def _mutate_config(gateway: str, config: Config):
         assert to_ds(model_b_snapshot.updated_ts) == "2020-01-05"
         assert to_ds(model_b_snapshot.expiration_ts) == "2020-01-12"
 
+        # model c should expire at the same time as model b
+        model_c_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_c" in n)
+        assert to_ds(model_c_snapshot.updated_ts) == to_ds(model_b_snapshot.updated_ts)
+        assert to_ds(model_c_snapshot.expiration_ts) == to_ds(model_b_snapshot.expiration_ts)
+
+        # model d should expire at the same time as model b
+        model_d_snapshot = next(s for n, s in sqlmesh.snapshots.items() if "model_d" in n)
+        assert to_ds(model_d_snapshot.updated_ts) == to_ds(model_b_snapshot.updated_ts)
+        assert to_ds(model_d_snapshot.expiration_ts) == to_ds(model_b_snapshot.expiration_ts)
+
     # 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.
+    # run janitor
+    # - table model a is expired so will be cleaned up and this will cascade to view model b
+    # - view model b is not expired, but because it got cascaded to, this will cascade again to view model c
+    # - table model d is a not a view, so even though its parent view model b got dropped, it doesnt need to be dropped
     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 = _state_sync_engine_adapter(sqlmesh).fetchall(
-            f"select name, identifier from sqlmesh._snapshots"
-        )
+        before_snapshot_ids = _all_snapshot_ids(sqlmesh)
+
         sqlmesh.invalidate_environment("dev")
         sqlmesh.run_janitor(ignore_ttl=False)
-        after_snapshots = _state_sync_engine_adapter(sqlmesh).fetchall(
-            f"select name, identifier from sqlmesh._snapshots"
-        )
 
-        assert len(before_snapshots) != len(after_snapshots)
+        after_snapshot_ids = _all_snapshot_ids(sqlmesh)
+
+        assert len(before_snapshot_ids) != len(after_snapshot_ids)
 
-        # 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])
+        # all that's left should be the two original snapshots that were in prod and model d
+        assert set(after_snapshot_ids) == set(
+            [
+                model_a_prod_snapshot.snapshot_id,
+                model_b_prod_snapshot.snapshot_id,
+                model_d_snapshot.snapshot_id,
+            ]
+        )

tests/core/state_sync/test_state_sync.py

@@ -1208,6 +1208,85 @@ def test_get_expired_snapshots_includes_downstream_view_snapshots(
     assert snapshot_b_cleanup
 
 
+def test_get_expired_snapshots_includes_downstream_transitive_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"),
+        ),
+        ttl="in 10 seconds",
+    )
+    snapshot_a.updated_ts = now_ts - 15000  # now - 15 seconds = expired
+    assert not snapshot_a.parents
+
+    # model_b: view snapshot that depends on model_a table snapshot
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="b",
+            kind="VIEW",
+            depends_on=["a"],
+            query=parse_one("select *, 'model_b' as m from a"),
+        ),
+        nodes={'"a"': snapshot_a.model},
+        ttl="in 10 seconds",
+    )
+    assert snapshot_b.parents == (snapshot_a.snapshot_id,)
+
+    # model_c: view snapshot that depends on model_b view snapshot (i.e no direct dependency on the model_a table ansphot).
+    # if the model_b view is dropped, model_c should also be dropped because it's a view that depends on a view being dropped
+    snapshot_c = make_snapshot(
+        SqlModel(
+            name="c",
+            kind="VIEW",
+            depends_on=["b"],
+            query=parse_one("select *, 'model_c' as m from b"),
+        ),
+        nodes={'"a"': snapshot_a.model, '"b"': snapshot_b.model},
+        ttl="in 10 seconds",
+    )
+    assert snapshot_c.parents == (snapshot_b.snapshot_id,)
+
+    # model_d: table snapshot that depends on model_b view snapshot (i.e no direct dependency on the model_a table snapshot).
+    # if the model_b view is dropped, model_d should NOT be dropped because it can still be queried afterwards (and has not expired)
+    snapshot_d = make_snapshot(
+        SqlModel(
+            name="d",
+            kind="FULL",
+            depends_on=["b"],
+            query=parse_one("select *, 'model_d' as m from b"),
+        ),
+        nodes={'"a"': snapshot_a.model, '"b"': snapshot_b.model},
+        ttl="in 10 seconds",
+    )
+    assert snapshot_d.parents == (snapshot_b.snapshot_id,)
+
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_c.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_d.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    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
+    # snapshot_c is unexpired and should be cleaned up because snapshot_b is being cleaned up
+    # snapshot_d is unexpired and should not be cleaned up. Although it depends on snapshot_b which is being cleaned up,
+    #  its a table and not a view so will not be invalid the second snapshot_b is cleaned up
+    assert len(cleanup_targets) == 3
+
+    assert set(t.snapshot.snapshot_id for t in cleanup_targets) == set(
+        [snapshot_a.snapshot_id, snapshot_b.snapshot_id, snapshot_c.snapshot_id]
+    )
+
+
 def test_delete_expired_snapshots(state_sync: EngineAdapterStateSync, make_snapshot: t.Callable):
     now_ts = now_timestamp()