Chore: Split migration script implementations into DDL and DML (#5307)

Author: izeigerman

sqlmesh/core/context.py

@@ -587,6 +587,7 @@ def state_sync(self) -> StateSync:
             self._state_sync = self._new_state_sync()
 
             if self._state_sync.get_versions(validate=False).schema_version == 0:
+                self.console.log_status_update("Initializing new project state...")
                 self._state_sync.migrate(default_catalog=self.default_catalog)
             self._state_sync.get_versions()
             self._state_sync = CachingStateSync(self._state_sync)  # type: ignore

sqlmesh/core/state_sync/db/environment.py

@@ -285,11 +285,13 @@ def get_environment_statements(self, environment: str) -> t.List[EnvironmentStat
         return []
 
     def _environment_from_row(self, row: t.Tuple[str, ...]) -> Environment:
-        return Environment(**{field: row[i] for i, field in enumerate(Environment.all_fields())})
+        return Environment(
+            **{field: row[i] for i, field in enumerate(sorted(Environment.all_fields()))}
+        )
 
     def _environment_summmary_from_row(self, row: t.Tuple[str, ...]) -> EnvironmentSummary:
         return EnvironmentSummary(
-            **{field: row[i] for i, field in enumerate(EnvironmentSummary.all_fields())}
+            **{field: row[i] for i, field in enumerate(sorted(EnvironmentSummary.all_fields()))}
         )
 
     def _environments_query(
@@ -298,7 +300,7 @@ def _environments_query(
         lock_for_update: bool = False,
         required_fields: t.Optional[t.List[str]] = None,
     ) -> exp.Select:
-        query_fields = required_fields if required_fields else Environment.all_fields()
+        query_fields = required_fields if required_fields else sorted(Environment.all_fields())
         query = (
             exp.select(*(exp.to_identifier(field) for field in query_fields))
             .from_(self.environments_table)
@@ -328,7 +330,7 @@ def _fetch_environment_summaries(
                 self.engine_adapter,
                 self._environments_query(
                     where=where,
-                    required_fields=list(EnvironmentSummary.all_fields()),
+                    required_fields=sorted(EnvironmentSummary.all_fields()),
                 ),
             )
         ]

sqlmesh/core/state_sync/db/migrator.py

@@ -173,9 +173,14 @@ def _apply_migrations(
 
         snapshot_count_before = self.snapshot_state.count() if versions.schema_version else None
 
+        state_table_exist = any(self.engine_adapter.table_exists(t) for t in self._state_tables)
+
         for migration in migrations:
             logger.info(f"Applying migration {migration}")
-            migration.migrate(state_sync, default_catalog=default_catalog)
+            migration.migrate_schemas(state_sync, default_catalog=default_catalog)
+            if state_table_exist:
+                # No need to run DML for the initial migration since all tables are empty
+                migration.migrate_rows(state_sync, default_catalog=default_catalog)
 
         snapshot_count_after = self.snapshot_state.count()
 

sqlmesh/migrations/v0001_init.py

@@ -9,7 +9,7 @@
 from sqlmesh.utils.migration import index_text_type
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
     engine_adapter = state_sync.engine_adapter
     schema = state_sync.schema
     snapshots_table = "_snapshots"
@@ -58,3 +58,7 @@ def migrate(state_sync, **kwargs):  # type: ignore
             "sqlglot_version": exp.DataType.build("text"),
         },
     )
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
+    pass

sqlmesh/migrations/v0002_remove_identify.py

@@ -1,5 +1,9 @@
 """Remove identify=True kwarg for rendering sql"""
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
+    pass
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
     pass

sqlmesh/migrations/v0003_move_batch_size.py

@@ -5,7 +5,11 @@
 from sqlglot import exp
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
+    pass
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
     snapshots_table = "_snapshots"
     if state_sync.schema:
         snapshots_table = f"{state_sync.schema}.{snapshots_table}"

sqlmesh/migrations/v0004_environmnent_add_finalized_at.py

@@ -3,7 +3,7 @@
 from sqlglot import exp
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
     engine_adapter = state_sync.engine_adapter
     environments_table = "_environments"
     if state_sync.schema:
@@ -21,3 +21,7 @@ def migrate(state_sync, **kwargs):  # type: ignore
     )
 
     engine_adapter.execute(alter_table_exp)
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
+    pass

sqlmesh/migrations/v0005_create_seed_table.py

@@ -5,7 +5,7 @@
 from sqlmesh.utils.migration import index_text_type
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
     engine_adapter = state_sync.engine_adapter
     seeds_table = "_seeds"
     if state_sync.schema:
@@ -22,3 +22,7 @@ def migrate(state_sync, **kwargs):  # type: ignore
         },
         primary_key=("name", "identifier"),
     )
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
+    pass

sqlmesh/migrations/v0006_change_seed_hash.py

@@ -1,5 +1,9 @@
 """Seed hashes moved from to_string to to_json for performance."""
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
+    pass
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
     pass

sqlmesh/migrations/v0007_env_table_info_to_kind.py

@@ -12,7 +12,11 @@ def _hash(data):  # type: ignore
     return str(zlib.crc32(";".join("" if d is None else d for d in data).encode("utf-8")))
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
+    pass
+
+
+def migrate_rows(state_sync, **kwargs):  # type: ignore
     import pandas as pd
 
     engine_adapter = state_sync.engine_adapter