Refactor: move pre_check into the migration script

Author: georgesittas

sqlmesh/core/state_sync/base.py

@@ -8,7 +8,7 @@
 
 from sqlglot import __version__ as SQLGLOT_VERSION
 
-from sqlmesh import migrations, pre_checks
+from sqlmesh import migrations
 from sqlmesh.core.environment import (
     Environment,
     EnvironmentNamingInfo,
@@ -64,11 +64,10 @@ def _schema_version_validator(cls, v: t.Any) -> int:
     importlib.import_module(f"sqlmesh.migrations.{migration}")
     for migration in sorted(info.name for info in pkgutil.iter_modules(migrations.__path__))
 ]
-PRE_CHECKS = {
-    pre_check: importlib.import_module(f"sqlmesh.pre_checks.{pre_check}")
-    for pre_check in sorted(info.name for info in pkgutil.iter_modules(pre_checks.__path__))
-}
 SCHEMA_VERSION: int = len(MIGRATIONS)
+PRE_CHECK_VERSION: int = (
+    max(idx for idx, migration in enumerate(MIGRATIONS) if hasattr(migration, "pre_check")) + 1
+)
 
 
 class PromotionResult(PydanticModel):

sqlmesh/core/state_sync/db/migrator.py

@@ -25,10 +25,7 @@
 from sqlmesh.core.snapshot.definition import (
     _parents_from_node,
 )
-from sqlmesh.core.state_sync.base import (
-    MIGRATIONS,
-    PRE_CHECKS,
-)
+from sqlmesh.core.state_sync.base import MIGRATIONS
 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
@@ -102,7 +99,7 @@ def migrate(
             promoted_snapshots_only: Whether to migrate only promoted snapshots.
             pre_check_only: If True, only run pre-checks without performing migration.
         """
-        pre_check_warnings = self.run_pre_checks(state_sync)
+        pre_check_warnings = self._run_pre_checks(state_sync)
         should_migrate = self.console.log_pre_check_warnings(pre_check_warnings, pre_check_only)
         if not should_migrate:
             return
@@ -168,30 +165,27 @@ def rollback(self) -> None:
 
         logger.info("Migration rollback successful.")
 
-    def run_pre_checks(self, state_sync: StateSync) -> t.List[t.Tuple[str, t.List[str]]]:
+    def _run_pre_checks(self, state_sync: StateSync) -> t.List[t.Tuple[str, t.List[str]]]:
         """Run pre-checks for migrations between specified versions.
 
         Args:
             state_sync: The state sync instance.
 
         Returns:
-            A list of pairs comprising the executed pre-checks and the corresponding warnings.
+            A list of pairs comprising the migration name containing the executed pre-checks
+            and the corresponding warnings.
         """
-        # Get the range of the migrations that would be applied
-        from_version = self.version_state.get_versions().schema_version
-        to_version = len(MIGRATIONS)
+        versions = self.version_state.get_versions()
+        migrations = MIGRATIONS[versions.schema_version :]
 
         pre_check_warnings = []
-        for i in range(from_version, to_version):
-            # Assumption: pre-check and migration names match
-            pre_check_name = MIGRATIONS[i].__name__.split(".")[-1]
-            pre_check_module = PRE_CHECKS.get(pre_check_name)
-
-            if callable(pre_check := getattr(pre_check_module, "pre_check", None)):
-                logger.info(f"Running pre-check for {pre_check_name}")
+        for migration in migrations:
+            if callable(pre_check := getattr(migration, "pre_check", None)):
+                migration_name = migration.__name__.split(".")[-1]
+                logger.info(f"Running pre-check for {migration_name}")
                 warnings = pre_check(state_sync)
                 if warnings:
-                    pre_check_warnings.append((pre_check_name, warnings))
+                    pre_check_warnings.append((migration_name, warnings))
 
         return pre_check_warnings