Chore!: Use version instead of identifier in the seeds table (#2481)

Author: izeigerman

sqlmesh/core/plan/builder.py

@@ -472,9 +472,6 @@ def _categorize_snapshots(
                                     indirectly_modified,
                                 )
 
-                # set to breaking if an indirect child has no directly modified parents
-                # that need a decision. this can happen when a revert to a parent causes
-                # an indirectly modified snapshot to be created because of a new parent
                 if (
                     not is_directly_modified
                     and not snapshot.version
@@ -484,11 +481,18 @@ def _categorize_snapshots(
                         for upstream in dag.upstream(s_id)
                     )
                 ):
-                    snapshot.categorize_as(
-                        SnapshotChangeCategory.FORWARD_ONLY
-                        if self._is_forward_only_model(s_id)
-                        else SnapshotChangeCategory.INDIRECT_BREAKING
-                    )
+                    if self._context_diff.indirectly_modified(snapshot.name):
+                        # Set to breaking if an indirect child has no directly modified parents
+                        # that need a decision. this can happen when a revert to a parent causes
+                        # an indirectly modified snapshot to be created because of a new parent
+                        snapshot.categorize_as(
+                            SnapshotChangeCategory.FORWARD_ONLY
+                            if self._is_forward_only_model(s_id)
+                            else SnapshotChangeCategory.INDIRECT_BREAKING
+                        )
+                    else:
+                        # Metadata updated.
+                        snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
 
             elif s_id in self._context_diff.added and self._is_new_snapshot(snapshot):
                 snapshot.categorize_as(

sqlmesh/core/state_sync/engine_adapter.py

@@ -144,7 +144,7 @@ def __init__(
 
         self._seed_columns_to_types = {
             "name": exp.DataType.build("text"),
-            "identifier": exp.DataType.build("text"),
+            "version": exp.DataType.build("text"),
             "content": exp.DataType.build("text"),
         }
 
@@ -220,7 +220,7 @@ def _push_snapshots(self, snapshots: t.Iterable[Snapshot], overwrite: bool = Fal
                 seed_contents.append(
                     {
                         "name": snapshot.name,
-                        "identifier": snapshot.identifier,
+                        "version": snapshot.version,
                         "content": seed_model.seed.content,
                     }
                 )
@@ -341,6 +341,10 @@ def _is_snapshot_used(snapshot: Snapshot) -> bool:
                         )
                     )
 
+        seed_deletion_candidates = [t.snapshot for t in cleanup_targets if not t.dev_table_only]
+        if seed_deletion_candidates:
+            self._delete_seeds(seed_deletion_candidates)
+
         return cleanup_targets
 
     def delete_expired_environments(self) -> t.List[Environment]:
@@ -368,7 +372,6 @@ def delete_expired_environments(self) -> t.List[Environment]:
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         for where in self._snapshot_id_filter(snapshot_ids):
             self.engine_adapter.delete_from(self.snapshots_table, where=where)
-            self.engine_adapter.delete_from(self.seeds_table, where=where)
 
     def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[SnapshotId]:
         return self._snapshot_ids_exist(snapshot_ids, self.snapshots_table)
@@ -486,8 +489,8 @@ def _get_snapshots(
                     exp.to_table(self.seeds_table).as_("seeds"),
                     on=exp.and_(
                         exp.column("name", table="snapshots").eq(exp.column("name", table="seeds")),
-                        exp.column("identifier", table="snapshots").eq(
-                            exp.column("identifier", table="seeds")
+                        exp.column("version", table="snapshots").eq(
+                            exp.column("version", table="seeds")
                         ),
                     ),
                     join_type="left",
@@ -998,7 +1001,6 @@ def _migrate_rows(self, promoted_snapshots_only: bool) -> None:
         if not snapshot_mapping:
             logger.info("No changes to snapshots detected")
             return
-        self._migrate_seed_rows(snapshot_mapping)
         self._migrate_environment_rows(environments, snapshot_mapping)
 
     def _migrate_snapshot_rows(
@@ -1147,51 +1149,6 @@ def _visit(
 
         return all_snapshot_mapping
 
-    def _migrate_seed_rows(self, snapshot_mapping: t.Dict[SnapshotId, SnapshotTableInfo]) -> None:
-        # FIXME: This migration won't be necessary if the primary key of the seeds table is changed to
-        # (name, version) instead of (name, identifier).
-        seed_snapshot_ids = [
-            s_id for s_id, table_info in snapshot_mapping.items() if table_info.is_seed
-        ]
-        if not seed_snapshot_ids:
-            logger.info("No seed rows to migrate")
-            return
-
-        logger.info("Migrating seed rows...")
-
-        for where in self._snapshot_id_filter(
-            seed_snapshot_ids, batch_size=self.SNAPSHOT_SEED_MIGRATION_BATCH_SIZE
-        ):
-            seeds = {
-                SnapshotId(name=name, identifier=identifier): content
-                for name, identifier, content in self._fetchall(
-                    exp.select("name", "identifier", "content").from_(self.seeds_table).where(where)
-                )
-            }
-            if not seeds:
-                continue
-
-            new_seeds = {}
-            for snapshot_id, content in seeds.items():
-                new_snapshot_id = snapshot_mapping[snapshot_id].snapshot_id
-                new_seeds[new_snapshot_id] = {
-                    "name": new_snapshot_id.name,
-                    "identifier": new_snapshot_id.identifier,
-                    "content": content,
-                }
-
-            existing_snapshot_ids = self._snapshot_ids_exist(new_seeds, self.seeds_table)
-            seeds_to_push = [
-                s for s_id, s in new_seeds.items() if s_id not in existing_snapshot_ids
-            ]
-
-            if seeds_to_push:
-                self.engine_adapter.insert_append(
-                    self.seeds_table,
-                    pd.DataFrame(seeds_to_push),
-                    columns_to_types=self._seed_columns_to_types,
-                )
-
     def _migrate_environment_rows(
         self,
         environments: t.List[Environment],
@@ -1226,6 +1183,10 @@ def _migrate_environment_rows(
             except Exception:
                 logger.warning("Failed to unpause migrated snapshots", exc_info=True)
 
+    def _delete_seeds(self, snapshots: t.Iterable[SnapshotNameVersionLike]) -> None:
+        for where in self._snapshot_name_version_filter(snapshots, alias=None):
+            self.engine_adapter.delete_from(self.seeds_table, where=where)
+
     def _snapshot_ids_exist(
         self, snapshot_ids: t.Iterable[SnapshotIdLike], table_name: exp.Table
     ) -> t.Set[SnapshotId]:

sqlmesh/migrations/v0049_replace_identifier_with_version_in_seeds_table.py

@@ -0,0 +1,57 @@
+"""Use version instead of identifier in the seeds table."""
+
+from sqlglot import exp
+
+from sqlmesh.utils.migration import index_text_type
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+
+    snapshots_table = "_snapshots"
+    seeds_table = "_seeds"
+    new_seeds_table = f"{seeds_table}_v49"
+
+    if state_sync.schema:
+        snapshots_table = f"{state_sync.schema}.{snapshots_table}"
+        seeds_table = f"{state_sync.schema}.{seeds_table}"
+        new_seeds_table = f"{state_sync.schema}.{new_seeds_table}"
+
+    index_type = index_text_type(engine_adapter.dialect)
+
+    engine_adapter.drop_table(new_seeds_table)
+    engine_adapter.create_state_table(
+        new_seeds_table,
+        {
+            "name": exp.DataType.build(index_type),
+            "version": exp.DataType.build(index_type),
+            "content": exp.DataType.build("text"),
+        },
+        primary_key=("name", "version"),
+    )
+
+    name_col = exp.column("name", table="seeds")
+    version_col = exp.column("version", table="snapshots")
+    query = (
+        exp.select(
+            name_col,
+            version_col,
+            exp.func("MAX", exp.column("content", table="seeds")).as_("content"),
+        )
+        .from_(exp.to_table(seeds_table).as_("seeds"))
+        .join(
+            exp.to_table(snapshots_table).as_("snapshots"),
+            on=exp.and_(
+                exp.column("name", table="seeds").eq(exp.column("name", table="snapshots")),
+                exp.column("identifier", table="seeds").eq(
+                    exp.column("identifier", table="snapshots")
+                ),
+            ),
+        )
+        .where(exp.column("version", table="snapshots").is_(exp.null()).not_())
+        .group_by(name_col, version_col)
+    )
+
+    engine_adapter.insert_append(new_seeds_table, query)
+    engine_adapter.drop_table(seeds_table)
+    engine_adapter.rename_table(new_seeds_table, "_seeds")

tests/core/test_state_sync.py

@@ -1044,6 +1044,40 @@ def test_delete_expired_snapshots(state_sync: EngineAdapterStateSync, make_snaps
     assert not state_sync.get_snapshots(None)
 
 
+def test_delete_expired_snapshots_seed(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
+):
+    now_ts = now_timestamp()
+
+    snapshot = make_snapshot(
+        SeedModel(
+            name="a",
+            kind=SeedKind(path="./path/to/seed"),
+            seed=Seed(content="header\n1\n2"),
+            column_hashes={"header": "hash"},
+            depends_on=set(),
+        ),
+    )
+    snapshot.ttl = "in 10 seconds"
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot.updated_ts = now_ts - 15000
+
+    state_sync.push_snapshots([snapshot])
+    assert set(state_sync.get_snapshots(None)) == {snapshot.snapshot_id}
+    assert state_sync.engine_adapter.fetchall(
+        "SELECT name, version, content FROM sqlmesh._seeds"
+    ) == [
+        (snapshot.name, snapshot.version, snapshot.model.seed.content),
+    ]
+
+    assert state_sync.delete_expired_snapshots() == [
+        SnapshotTableCleanupTask(snapshot=snapshot.table_info, dev_table_only=False),
+    ]
+
+    assert not state_sync.get_snapshots(None)
+    assert not state_sync.engine_adapter.fetchall("SELECT * FROM sqlmesh._seeds")
+
+
 def test_delete_expired_snapshots_batching(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
 ):
@@ -1581,6 +1615,8 @@ def test_migrate_rows(state_sync: EngineAdapterStateSync, mocker: MockerFixture)
         },
     )
 
+    state_sync.engine_adapter.drop_table("sqlmesh._seeds")
+
     old_snapshots = state_sync.engine_adapter.fetchdf("select * from sqlmesh._snapshots")
     old_environments = state_sync.engine_adapter.fetchdf("select * from sqlmesh._environments")
 
@@ -2124,28 +2160,29 @@ def test_snapshot_batching(state_sync, mocker, make_snapshot):
     state_sync.SNAPSHOT_BATCH_SIZE = 2
     state_sync.engine_adapter = mock
 
+    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1")), "1")
+    snapshot_b = make_snapshot(SqlModel(name="a", query=parse_one("select 2")), "2")
+    snapshot_c = make_snapshot(SqlModel(name="a", query=parse_one("select 3")), "3")
+
     state_sync.delete_snapshots(
         (
-            SnapshotId(name="a", identifier="1"),
-            SnapshotId(name="a", identifier="2"),
-            SnapshotId(name="a", identifier="3"),
+            snapshot_a,
+            snapshot_b,
+            snapshot_c,
         )
     )
     calls = mock.delete_from.call_args_list
     assert mock.delete_from.call_args_list == [
         call(
             exp.to_table("sqlmesh._snapshots"),
-            where=parse_one("(name, identifier) in (('a', '1'), ('a', '2'))"),
-        ),
-        call(
-            exp.to_table("sqlmesh._seeds"),
-            where=parse_one("(name, identifier) in (('a', '1'), ('a', '2'))"),
+            where=parse_one(
+                f"(name, identifier) in (('\"a\"', '{snapshot_b.identifier}'), ('\"a\"', '{snapshot_a.identifier}'))"
+            ),
         ),
         call(
             exp.to_table("sqlmesh._snapshots"),
-            where=parse_one("(name, identifier) in (('a', '3'))"),
+            where=parse_one(f"(name, identifier) in (('\"a\"', '{snapshot_c.identifier}'))"),
         ),
-        call(exp.to_table("sqlmesh._seeds"), where=parse_one("(name, identifier) in (('a', '3'))")),
     ]
 
     mock.fetchall.side_effect = [