Repository: TobikoData/sqlmesh
Pull request #2494 (2024-04-23 17:00:47)

Chore!: Use version instead of identifier in the seeds table
This is a re-introduction of #2481 which I had to revert because of issues with MSSQL and MySQL surfaced in the engine adapter integration tests. I've added fixes for table renaming API to support these engines.

Previous description: 
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/docs/index.md b/docs/index.md
index cbece990..e5ecc7f8 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -9,7 +9,7 @@ The experience of developing and deploying data pipelines is more uncertain and
 Here are some challenges that data teams run into, especially when data sizes increase or the number of data users expands:
 
 1. Data pipelines are fragmented and fragile
-    * Data pipelines generally consist of Python or SQL scripts that implicitly depend upon each other through tables. Changes to upstream scripts that break downstream dependencies are usually only detected at run time.
+    * Data pipelines generally consist of Python or SQL scripts that implicitly depend upon each other through tables. Changes to upstream scripts that break downstream consumers are usually only detected at run time.
 
 1. Data quality checks are not sufficient
     * The data community has settled on data quality checks as the "solution" for testing data pipelines. Although data quality checks are great for detecting large unexpected data changes, they are expensive to run, and they have trouble validating exact logic.
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 0eb6c633..b85f193a 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -1611,7 +1611,7 @@ class EngineAdapter:
                 raise UnsupportedCatalogOperationError(
                     "Tried to rename table across catalogs which is not supported"
                 )
-        self.execute(exp.rename_table(old_table_name, new_table_name))
+        self._rename_table(old_table_name, new_table_name)
 
     def get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
@@ -2029,6 +2029,13 @@ class EngineAdapter:
                     exc_info=True,
                 )
 
+    def _rename_table(
+        self,
+        old_table_name: TableName,
+        new_table_name: TableName,
+    ) -> None:
+        self.execute(exp.rename_table(old_table_name, new_table_name))
+
     @classmethod
     def _select_columns(cls, columns: t.Iterable[str]) -> exp.Select:
         return exp.select(*(exp.column(c, quoted=True) for c in columns))
diff --git a/sqlmesh/core/engine_adapter/mssql.py b/sqlmesh/core/engine_adapter/mssql.py
index 1988441e..54bafcb5 100644
--- a/sqlmesh/core/engine_adapter/mssql.py
+++ b/sqlmesh/core/engine_adapter/mssql.py
@@ -222,3 +222,12 @@ class MSSQLEngineAdapter(
     def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.Any) -> str:
         sql = super()._to_sql(expression, quote=quote, **kwargs)
         return f"{sql};"
+
+    def _rename_table(
+        self,
+        old_table_name: TableName,
+        new_table_name: TableName,
+    ) -> None:
+        # The function that renames tables in MSSQL takes string literals as arguments instead of identifiers,
+        # so we shouldn't quote the identifiers.
+        self.execute(exp.rename_table(old_table_name, new_table_name), quote_identifiers=False)
diff --git a/sqlmesh/core/engine_adapter/mysql.py b/sqlmesh/core/engine_adapter/mysql.py
index 51963ce0..5bc0d724 100644
--- a/sqlmesh/core/engine_adapter/mysql.py
+++ b/sqlmesh/core/engine_adapter/mysql.py
@@ -142,3 +142,16 @@ class MySQLEngineAdapter(
                 f"Column comments for table '{table.alias_or_name}' not registered - this may be due to limited permissions.",
                 exc_info=True,
             )
+
+    def _rename_table(
+        self,
+        old_table_name: TableName,
+        new_table_name: TableName,
+    ) -> None:
+        old_table = exp.to_table(old_table_name)
+        new_table = exp.to_table(new_table_name)
+        if not new_table.db and old_table.db:
+            # In MySQL you need to provide the full target table name.
+            # Therefore we use the old schema name if one is not provided explicitly in the new name.
+            new_table.set("db", old_table.args["db"])
+        return super()._rename_table(old_table, new_table)
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/engine_adapter/test_mssql.py b/tests/core/engine_adapter/test_mssql.py
index ed15d4f9..5181db32 100644
--- a/tests/core/engine_adapter/test_mssql.py
+++ b/tests/core/engine_adapter/test_mssql.py
@@ -516,3 +516,16 @@ def test_df_dates(make_mocked_engine_adapter: t.Callable):
             }
         )
     )
+
+
+def test_rename_table(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+
+    adapter.rename_table("test_schema.old_name", "new_name")
+    adapter.rename_table("old_name", "new_name")
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        "EXEC sp_rename 'test_schema.old_name', 'new_name';",
+        "EXEC sp_rename 'old_name', 'new_name';",
+    ]
diff --git a/tests/core/engine_adapter/test_mysql.py b/tests/core/engine_adapter/test_mysql.py
index e0f07de1..348e5752 100644
--- a/tests/core/engine_adapter/test_mysql.py
+++ b/tests/core/engine_adapter/test_mysql.py
@@ -62,3 +62,18 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
         f"ALTER TABLE `test_table` COMMENT = '{truncated_table_comment}'",
         f"ALTER TABLE `test_table` MODIFY `a` INT COMMENT '{truncated_column_comment}'",
     ]
+
+
+def test_rename_table(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(MySQLEngineAdapter)
+
+    adapter.rename_table("test_schema.old_name", "new_name")
+    adapter.rename_table("test_schema.old_name", "new_test_schema.new_name")
+    adapter.rename_table("old_name", "new_name")
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        "ALTER TABLE `test_schema`.`old_name` RENAME TO `test_schema`.`new_name`",
+        "ALTER TABLE `test_schema`.`old_name` RENAME TO `new_test_schema`.`new_name`",
+        "ALTER TABLE `old_name` RENAME TO `new_name`",
+    ]
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index b52b8b7a..c0d9c032 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -1045,6 +1045,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
 ):
@@ -1582,6 +1616,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")
 
@@ -2125,28 +2161,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 = [