fix: remove state sync from migrate (#5502)

Author: eakmanrq

sqlmesh/core/state_sync/db/facade.py

@@ -469,7 +469,7 @@ def migrate(
     ) -> None:
         """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
         self.migrator.migrate(
-            self,
+            self.schema,
             skip_backup=skip_backup,
             promoted_snapshots_only=promoted_snapshots_only,
         )

sqlmesh/core/state_sync/db/migrator.py

@@ -30,7 +30,6 @@
     MIN_SCHEMA_VERSION,
     MIN_SQLMESH_VERSION,
 )
-from sqlmesh.core.state_sync.base import StateSync
 from sqlmesh.core.state_sync.db.environment import EnvironmentState
 from sqlmesh.core.state_sync.db.interval import IntervalState
 from sqlmesh.core.state_sync.db.snapshot import SnapshotState
@@ -85,7 +84,7 @@ def __init__(
 
     def migrate(
         self,
-        state_sync: StateSync,
+        schema: t.Optional[str],
         skip_backup: bool = False,
         promoted_snapshots_only: bool = True,
     ) -> None:
@@ -94,7 +93,7 @@ def migrate(
         migration_start_ts = time.perf_counter()
 
         try:
-            migrate_rows = self._apply_migrations(state_sync, skip_backup)
+            migrate_rows = self._apply_migrations(schema, skip_backup)
 
             if not migrate_rows and major_minor(SQLMESH_VERSION) == versions.minor_sqlmesh_version:
                 return
@@ -153,7 +152,7 @@ def rollback(self) -> None:
 
     def _apply_migrations(
         self,
-        state_sync: StateSync,
+        schema: t.Optional[str],
         skip_backup: bool,
     ) -> bool:
         versions = self.version_state.get_versions()
@@ -184,10 +183,10 @@ def _apply_migrations(
 
         for migration in migrations:
             logger.info(f"Applying migration {migration}")
-            migration.migrate_schemas(state_sync)
+            migration.migrate_schemas(engine_adapter=self.engine_adapter, schema=schema)
             if state_table_exist:
                 # No need to run DML for the initial migration since all tables are empty
-                migration.migrate_rows(state_sync)
+                migration.migrate_rows(engine_adapter=self.engine_adapter, schema=schema)
 
         snapshot_count_after = self.snapshot_state.count()
 

sqlmesh/migrations/v0000_baseline.py

@@ -4,15 +4,12 @@
 from sqlmesh.utils.migration import blob_text_type, index_text_type
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
-    schema = state_sync.schema
-    engine_adapter = state_sync.engine_adapter
-
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     intervals_table = "_intervals"
     snapshots_table = "_snapshots"
     environments_table = "_environments"
     versions_table = "_versions"
-    if state_sync.schema:
+    if schema:
         engine_adapter.create_schema(schema)
         intervals_table = f"{schema}.{intervals_table}"
         snapshots_table = f"{schema}.{snapshots_table}"
@@ -94,5 +91,5 @@ def migrate_schemas(state_sync, **kwargs):  # type: ignore
     engine_adapter.create_index(intervals_table, "_intervals_name_version_idx", ("name", "version"))
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     pass

sqlmesh/migrations/v0061_mysql_fix_blob_text_type.py

@@ -9,12 +9,9 @@
 from sqlmesh.utils.migration import blob_text_type
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
-    engine_adapter = state_sync.engine_adapter
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     if engine_adapter.dialect != "mysql":
         return
-
-    schema = state_sync.schema
     environments_table = "_environments"
     snapshots_table = "_snapshots"
 
@@ -46,5 +43,5 @@ def migrate_schemas(state_sync, **kwargs):  # type: ignore
         engine_adapter.execute(alter_table_exp)
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     pass

sqlmesh/migrations/v0062_add_model_gateway.py

@@ -1,9 +1,9 @@
 """Add the gateway model attribute."""
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     pass
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     pass

sqlmesh/migrations/v0063_change_signals.py

@@ -7,15 +7,13 @@
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     pass
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     import pandas as pd
 
-    engine_adapter = state_sync.engine_adapter
-    schema = state_sync.schema
     snapshots_table = "_snapshots"
     index_type = index_text_type(engine_adapter.dialect)
     if schema:

sqlmesh/migrations/v0064_join_when_matched_strings.py

@@ -7,15 +7,13 @@
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     pass
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     import pandas as pd
 
-    engine_adapter = state_sync.engine_adapter
-    schema = state_sync.schema
     snapshots_table = "_snapshots"
     index_type = index_text_type(engine_adapter.dialect)
     if schema:

sqlmesh/migrations/v0065_add_model_optimize.py

@@ -1,9 +1,9 @@
 """Add the optimize_query model attribute."""
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     pass
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     pass

sqlmesh/migrations/v0066_add_auto_restatements.py

@@ -5,9 +5,7 @@
 from sqlmesh.utils.migration import index_text_type
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
-    engine_adapter = state_sync.engine_adapter
-    schema = state_sync.schema
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     auto_restatements_table = "_auto_restatements"
     intervals_table = "_intervals"
 
@@ -40,9 +38,7 @@ def migrate_schemas(state_sync, **kwargs):  # type: ignore
     engine_adapter.execute(alter_table_exp)
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
-    engine_adapter = state_sync.engine_adapter
-    schema = state_sync.schema
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     intervals_table = "_intervals"
 
     if schema:

sqlmesh/migrations/v0067_add_tsql_date_full_precision.py

@@ -1,9 +1,9 @@
 """Add full precision for tsql to support nanoseconds."""
 
 
-def migrate_schemas(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(engine_adapter, schema, **kwargs):  # type: ignore
     pass
 
 
-def migrate_rows(state_sync, **kwargs):  # type: ignore
+def migrate_rows(engine_adapter, schema, **kwargs):  # type: ignore
     pass