Repository: TobikoData/sqlmesh
Pull request #2030 (2024-01-26 20:21:35)

Fix: When deleting expired snapshots in janitor only fetch relevant snapshot records
Previously,  the janitor process would fetch all snapshot records to figure out what needs to be cleaned up. This led to a large memory footprint for projects with many models and / or changes.

This update ensures that only records that actually expired are fetched from the state DB.


Patch:
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index c6d9c2c9..4d1be08a 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -1036,6 +1036,10 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
             return None
         return t.cast(Model, self.node).fully_qualified_table
 
+    @property
+    def expiration_ts(self) -> int:
+        return to_timestamp(self.ttl, relative_base=to_datetime(self.updated_ts))
+
     def _ensure_categorized(self) -> None:
         if not self.change_category:
             raise SQLMeshError(f"Snapshot {self.snapshot_id} has not been categorized yet.")
diff --git a/sqlmesh/core/state_sync/common.py b/sqlmesh/core/state_sync/common.py
index 53d98295..99e423d9 100644
--- a/sqlmesh/core/state_sync/common.py
+++ b/sqlmesh/core/state_sync/common.py
@@ -3,7 +3,6 @@ from __future__ import annotations
 import abc
 import logging
 import typing as t
-from collections import defaultdict
 from datetime import datetime
 from functools import wraps
 
@@ -16,15 +15,10 @@ from sqlmesh.core.snapshot import (
     SnapshotIdLike,
     SnapshotInfoLike,
     SnapshotNameVersionLike,
-    SnapshotTableCleanupTask,
     start_date,
 )
-from sqlmesh.core.state_sync.base import (
-    PromotionResult,
-    SnapshotTableCleanupTask,
-    StateSync,
-)
-from sqlmesh.utils.date import TimeLike, now, now_timestamp, to_datetime, to_timestamp
+from sqlmesh.core.state_sync.base import PromotionResult, StateSync
+from sqlmesh.utils.date import TimeLike, now, now_timestamp, to_timestamp
 from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
@@ -227,55 +221,6 @@ class CommonStateSyncMixin(StateSync):
         environment.finalized_ts = now_timestamp()
         self._update_environment(environment)
 
-    @transactional()
-    def delete_expired_snapshots(self) -> t.List[SnapshotTableCleanupTask]:
-        current_time = now(minute_floor=False)
-
-        snapshots = self._get_snapshots(hydrate_intervals=False)
-
-        snapshots_by_version = defaultdict(set)
-        snapshots_by_temp_version = defaultdict(set)
-        for s in snapshots.values():
-            snapshots_by_version[(s.name, s.version)].add(s.snapshot_id)
-            snapshots_by_temp_version[(s.name, s.temp_version_get_or_generate())].add(s.snapshot_id)
-
-        promoted_snapshot_ids = {
-            snapshot.snapshot_id
-            for environment in self.get_environments()
-            for snapshot in environment.snapshots
-        }
-
-        def _is_snapshot_used(snapshot: Snapshot) -> bool:
-            return (
-                snapshot.snapshot_id in promoted_snapshot_ids
-                or to_datetime(snapshot.ttl, relative_base=to_datetime(snapshot.updated_ts))
-                > current_time
-            )
-
-        expired_snapshots = [s for s in snapshots.values() if not _is_snapshot_used(s)]
-
-        if expired_snapshots:
-            self.delete_snapshots(expired_snapshots)
-
-        cleanup_targets = []
-        for snapshot in expired_snapshots:
-            shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
-            shared_version_snapshots.discard(snapshot.snapshot_id)
-
-            shared_temp_version_snapshots = snapshots_by_temp_version[
-                (snapshot.name, snapshot.temp_version_get_or_generate())
-            ]
-            shared_temp_version_snapshots.discard(snapshot.snapshot_id)
-
-            if not shared_temp_version_snapshots:
-                cleanup_targets.append(
-                    SnapshotTableCleanupTask(
-                        snapshot=snapshot.table_info, dev_table_only=bool(shared_version_snapshots)
-                    )
-                )
-
-        return cleanup_targets
-
     @transactional()
     def unpause_snapshots(
         self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index 2e2f6209..0f6a36b0 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -44,7 +44,9 @@ from sqlmesh.core.snapshot import (
     SnapshotIdLike,
     SnapshotInfoLike,
     SnapshotIntervals,
+    SnapshotNameVersion,
     SnapshotNameVersionLike,
+    SnapshotTableCleanupTask,
     fingerprint_from_node,
 )
 from sqlmesh.core.snapshot.definition import (
@@ -55,7 +57,7 @@ from sqlmesh.core.snapshot.definition import (
 )
 from sqlmesh.core.state_sync.base import MIGRATIONS, SCHEMA_VERSION, StateSync, Versions
 from sqlmesh.core.state_sync.common import CommonStateSyncMixin, transactional
-from sqlmesh.utils import major_minor, random_id
+from sqlmesh.utils import major_minor, random_id, unique
 from sqlmesh.utils.date import TimeLike, now_timestamp, time_like_to_str
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import parse_obj_as
@@ -63,6 +65,9 @@ from sqlmesh.utils.pydantic import parse_obj_as
 logger = logging.getLogger(__name__)
 
 
+T = t.TypeVar("T")
+
+
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
 
@@ -115,6 +120,7 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
             "version": exp.DataType.build("text"),
             "snapshot": exp.DataType.build("text"),
             "kind_name": exp.DataType.build("text"),
+            "expiration_ts": exp.DataType.build("bigint"),
         }
 
         self._environment_columns_to_types = {
@@ -256,6 +262,74 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
             where=filter_expr,
         )
 
+    @transactional()
+    def delete_expired_snapshots(self) -> t.List[SnapshotTableCleanupTask]:
+        current_ts = now_timestamp(minute_floor=False)
+
+        expired_query = (
+            exp.select("name", "identifier", "version")
+            .from_(self.snapshots_table)
+            .where(exp.column("expiration_ts") <= current_ts)
+        )
+        expired_candidates = {
+            SnapshotId(name=name, identifier=identifier): SnapshotNameVersion(
+                name=name, version=version
+            )
+            for name, identifier, version in self.engine_adapter.fetchall(expired_query)
+        }
+        if not expired_candidates:
+            return []
+
+        promoted_snapshot_ids = {
+            snapshot.snapshot_id
+            for environment in self.get_environments()
+            for snapshot in environment.snapshots
+        }
+
+        def _is_snapshot_used(snapshot: Snapshot) -> bool:
+            return (
+                snapshot.snapshot_id in promoted_snapshot_ids
+                or snapshot.snapshot_id not in expired_candidates
+            )
+
+        unique_expired_versions = unique(expired_candidates.values())
+        version_batches = self._snapshot_batches(unique_expired_versions)
+        cleanup_targets = []
+        for versions_batch in version_batches:
+            snapshots = self._get_snapshots_with_same_version(versions_batch)
+
+            snapshots_by_version = defaultdict(set)
+            snapshots_by_temp_version = defaultdict(set)
+            for s in snapshots:
+                snapshots_by_version[(s.name, s.version)].add(s.snapshot_id)
+                snapshots_by_temp_version[(s.name, s.temp_version_get_or_generate())].add(
+                    s.snapshot_id
+                )
+
+            expired_snapshots = [s for s in snapshots if not _is_snapshot_used(s)]
+
+            if expired_snapshots:
+                self.delete_snapshots(expired_snapshots)
+
+            for snapshot in expired_snapshots:
+                shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
+                shared_version_snapshots.discard(snapshot.snapshot_id)
+
+                shared_temp_version_snapshots = snapshots_by_temp_version[
+                    (snapshot.name, snapshot.temp_version_get_or_generate())
+                ]
+                shared_temp_version_snapshots.discard(snapshot.snapshot_id)
+
+                if not shared_temp_version_snapshots:
+                    cleanup_targets.append(
+                        SnapshotTableCleanupTask(
+                            snapshot=snapshot.table_info,
+                            dev_table_only=bool(shared_version_snapshots),
+                        )
+                    )
+
+        return cleanup_targets
+
     def delete_expired_environments(self) -> t.List[Environment]:
         now_ts = now_timestamp()
         filter_expr = exp.LTE(
@@ -339,7 +413,7 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
         for where in self._snapshot_id_filter([snapshot.snapshot_id]):
             self.engine_adapter.update_table(
                 self.snapshots_table,
-                {"snapshot": _snapshot_to_json(snapshot)},
+                {"snapshot": _snapshot_to_json(snapshot), "expiration_ts": snapshot.expiration_ts},
                 where=where,
                 contains_json=True,
             )
@@ -1033,11 +1107,7 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
         name_identifiers = sorted(
             {(snapshot_id.name, snapshot_id.identifier) for snapshot_id in snapshot_ids}
         )
-
-        batches = [
-            name_identifiers[i : i + self.SNAPSHOT_BATCH_SIZE]
-            for i in range(0, len(name_identifiers), self.SNAPSHOT_BATCH_SIZE)
-        ]
+        batches = self._snapshot_batches(name_identifiers)
 
         if not name_identifiers:
             yield exp.false()
@@ -1068,11 +1138,7 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
         self, snapshot_name_versions: t.Iterable[SnapshotNameVersionLike], alias: str = "snapshots"
     ) -> t.Iterator[exp.Condition]:
         name_versions = sorted({(s.name, s.version) for s in snapshot_name_versions})
-
-        batches = [
-            name_versions[i : i + self.SNAPSHOT_BATCH_SIZE]
-            for i in range(0, len(name_versions), self.SNAPSHOT_BATCH_SIZE)
-        ]
+        batches = self._snapshot_batches(name_versions)
 
         if not name_versions:
             return exp.false()
@@ -1099,6 +1165,11 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                     ]
                 )
 
+    def _snapshot_batches(self, l: t.List[T]) -> t.List[t.List[T]]:
+        return [
+            l[i : i + self.SNAPSHOT_BATCH_SIZE] for i in range(0, len(l), self.SNAPSHOT_BATCH_SIZE)
+        ]
+
     @contextlib.contextmanager
     def _transaction(self) -> t.Iterator[None]:
         with self.engine_adapter.transaction():
@@ -1154,6 +1225,7 @@ def _snapshots_to_df(snapshots: t.Iterable[Snapshot]) -> pd.DataFrame:
                 "version": snapshot.version,
                 "snapshot": _snapshot_to_json(snapshot),
                 "kind_name": snapshot.model_kind_name.value if snapshot.model_kind_name else None,
+                "expiration_ts": snapshot.expiration_ts,
             }
             for snapshot in snapshots
         ]
diff --git a/sqlmesh/migrations/v0038_add_expiration_ts_to_snapshot.py b/sqlmesh/migrations/v0038_add_expiration_ts_to_snapshot.py
new file mode 100644
index 00000000..636eab99
--- /dev/null
+++ b/sqlmesh/migrations/v0038_add_expiration_ts_to_snapshot.py
@@ -0,0 +1,69 @@
+"""Add the expiration_ts column to the snapshots table."""
+import json
+
+import pandas as pd
+from sqlglot import exp
+
+from sqlmesh.utils.date import to_datetime, to_timestamp
+from sqlmesh.utils.migration import index_text_type
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    index_type = index_text_type(engine_adapter.dialect)
+
+    alter_table_exp = exp.AlterTable(
+        this=exp.to_table(snapshots_table),
+        actions=[
+            exp.ColumnDef(
+                this=exp.to_column("expiration_ts"),
+                kind=exp.DataType.build("bigint"),
+            )
+        ],
+    )
+    engine_adapter.execute(alter_table_exp)
+
+    new_snapshots = []
+
+    for name, identifier, version, snapshot, kind_name in engine_adapter.fetchall(
+        exp.select("name", "identifier", "version", "snapshot", "kind_name").from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+
+        updated_ts = parsed_snapshot["updated_ts"]
+        ttl = parsed_snapshot["ttl"]
+        expiration_ts = to_timestamp(ttl, relative_base=to_datetime(updated_ts))
+
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": snapshot,
+                "kind_name": kind_name,
+                "expiration_ts": expiration_ts,
+            }
+        )
+
+    if new_snapshots:
+        engine_adapter.delete_from(snapshots_table, "TRUE")
+
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types={
+                "name": exp.DataType.build(index_type),
+                "identifier": exp.DataType.build(index_type),
+                "version": exp.DataType.build(index_type),
+                "snapshot": exp.DataType.build("text"),
+                "kind_name": exp.DataType.build(index_type),
+                "expiration_ts": exp.DataType.build("bigint"),
+            },
+            contains_json=True,
+        )


Test patch:
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index b7c1ba95..85a77b48 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -876,6 +876,43 @@ def test_delete_expired_snapshots(state_sync: EngineAdapterStateSync, make_snaps
     assert not state_sync.get_snapshots(None)
 
 
+def test_delete_expired_snapshots_batching(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
+):
+    state_sync.SNAPSHOT_BATCH_SIZE = 1
+    now_ts = now_timestamp()
+
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            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
+
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="b",
+            query=parse_one("select a, b, ds"),
+        ),
+    )
+    snapshot_b.ttl = "in 10 seconds"
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.updated_ts = now_ts - 11000
+
+    state_sync.push_snapshots([snapshot_a, snapshot_b])
+    assert set(state_sync.get_snapshots(None)) == {snapshot_a.snapshot_id, snapshot_b.snapshot_id}
+
+    assert state_sync.delete_expired_snapshots() == [
+        SnapshotTableCleanupTask(snapshot=snapshot_a.table_info, dev_table_only=False),
+        SnapshotTableCleanupTask(snapshot=snapshot_b.table_info, dev_table_only=False),
+    ]
+
+    assert not state_sync.get_snapshots(None)
+
+
 def test_delete_expired_snapshots_promoted(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable, mocker: MockerFixture
 ):
@@ -909,8 +946,8 @@ def test_delete_expired_snapshots_promoted(
     env.snapshots = []
     state_sync.promote(env)
 
-    now_mock = mocker.patch("sqlmesh.core.state_sync.common.now")
-    now_mock.return_value = to_datetime(now_timestamp() + 11000)
+    now_timestamp_mock = mocker.patch("sqlmesh.core.state_sync.engine_adapter.now_timestamp")
+    now_timestamp_mock.return_value = now_timestamp() + 11000
 
     assert state_sync.delete_expired_snapshots() == [
         SnapshotTableCleanupTask(snapshot=snapshot.table_info, dev_table_only=False)
@@ -1190,6 +1227,13 @@ def test_migrate(state_sync: EngineAdapterStateSync, mocker: MockerFixture) -> N
         sqlmesh_version=SQLMESH_VERSION,
     )
 
+    assert (
+        state_sync.engine_adapter.fetchone(
+            "SELECT COUNT(*) FROM sqlmesh._snapshots WHERE expiration_ts IS NULL"
+        )[0]
+        == 0
+    )
+
 
 def test_rollback(state_sync: EngineAdapterStateSync, mocker: MockerFixture) -> None:
     with pytest.raises(