Repository: TobikoData/sqlmesh
Pull request #2481 (2024-04-20 18:51:57)

Chore!: Use version instead of identifier in the seeds table
The contents of the seed are only different if the snapshot's version is different too. Therefore, it doesn't make sense to store seed content by the (name, identifier) tuple since it leads to duplication and more complexity and inefficiencies downstream. For example, we don't need to migrate seed records if we were to store them by (name, version), since only identifiers change during the migration.


Patch:
diff --git a/sqlmesh/core/plan/builder.py b/sqlmesh/core/plan/builder.py
index 0d83fb55..09419c19 100644
--- a/sqlmesh/core/plan/builder.py
+++ b/sqlmesh/core/plan/builder.py
@@ -472,9 +472,6 @@ class PlanBuilder:
                                     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 @@ class PlanBuilder:
                         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(
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index b9163634..b2960aed 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -144,7 +144,7 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
 
         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 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                 seed_contents.append(
                     {
                         "name": snapshot.name,
-                        "identifier": snapshot.identifier,
+                        "version": snapshot.version,
                         "content": seed_model.seed.content,
                     }
                 )
@@ -341,6 +341,10 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                         )
                     )
 
+        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 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
     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 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                     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 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
         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 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
 
         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 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
             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]:
diff --git a/sqlmesh/migrations/v0049_replace_identifier_with_version_in_seeds_table.py b/sqlmesh/migrations/v0049_replace_identifier_with_version_in_seeds_table.py
new file mode 100644
index 00000000..5b0fb958
--- /dev/null
+++ b/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")


Test patch:
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index 8de2b07d..47786317 100644
--- a/tests/core/test_state_sync.py
+++ b/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 = [