Feat: introduce migration pre-checks

Author: georgesittas

sqlmesh/cli/main.py

@@ -917,12 +917,17 @@ def ui(ctx: click.Context, host: str, port: int, mode: str) -> None:
 
 
 @cli.command("migrate")
+@click.option(
+    "--pre-check",
+    is_flag=True,
+    help="Run pre-checks and display warnings without performing migration",
+)
 @click.pass_context
 @error_handler
 @cli_analytics
-def migrate(ctx: click.Context) -> None:
+def migrate(ctx: click.Context, pre_check: bool) -> None:
     """Migrate SQLMesh to the current running version."""
-    ctx.obj.migrate()
+    ctx.obj.migrate(pre_check_only=pre_check)
 
 
 @cli.command("rollback")

sqlmesh/core/console.py

@@ -497,6 +497,17 @@ def update_env_migration_progress(self, num_tasks: int) -> None:
     def stop_env_migration_progress(self, success: bool = True) -> None:
         """Stop the environment migration progress."""
 
+    @abc.abstractmethod
+    def log_pre_check_warnings(
+        self,
+        pre_check_warnings: t.List[t.Tuple[str, t.List[str]]],
+        pre_check_only: bool,
+    ) -> bool:
+        """
+        Log warnings emitted by pre-check scripts and ask user whether they'd like to
+        proceed with the migration (true) or not (false).
+        """
+
     @abc.abstractmethod
     def plan(
         self,
@@ -662,6 +673,13 @@ def update_env_migration_progress(self, num_tasks: int) -> None:
     def stop_env_migration_progress(self, success: bool = True) -> None:
         pass
 
+    def log_pre_check_warnings(
+        self,
+        pre_check_warnings: t.List[t.Tuple[str, t.List[str]]],
+        pre_check_only: bool,
+    ) -> bool:
+        return True
+
     def start_state_export(
         self,
         output_file: Path,
@@ -1472,6 +1490,33 @@ def stop_env_migration_progress(self, success: bool = True) -> None:
             if success:
                 self.log_success("Environments migrated successfully")
 
+    def log_pre_check_warnings(
+        self,
+        pre_check_warnings: t.List[t.Tuple[str, t.List[str]]],
+        pre_check_only: bool,
+    ) -> bool:
+        if pre_check_warnings:
+            for pre_check, warnings in pre_check_warnings:
+                tree = Tree(f"[bold]Pre-migration warnings for {pre_check}[/bold]")
+                for warning in warnings:
+                    tree.add(f"[yellow]{warning}[/yellow]")
+
+                self._print(tree)
+
+            if pre_check_only:
+                return False
+
+            should_continue = self._confirm("\nDo you want to proceed with the migration?")
+            if not should_continue:
+                self.log_status_update("Migration cancelled.")
+
+            return should_continue
+        if pre_check_only:
+            self.log_status_update("No pre-migration warnings detected.")
+            return False
+
+        return True
+
     def start_state_export(
         self,
         output_file: Path,

sqlmesh/core/context.py

@@ -2319,24 +2319,33 @@ def check_intervals(
         return results
 
     @python_api_analytics
-    def migrate(self) -> None:
+    def migrate(self, pre_check_only: bool = False) -> None:
         """Migrates SQLMesh to the current running version.
 
         Please contact your SQLMesh administrator before doing this.
+
+        Args:
+            pre_check_only: If True, only run pre-checks without performing the migration.
         """
-        self.notification_target_manager.notify(NotificationEvent.MIGRATION_START)
+        if not pre_check_only:
+            self.notification_target_manager.notify(NotificationEvent.MIGRATION_START)
+
         self._load_materializations()
         try:
             self._new_state_sync().migrate(
                 default_catalog=self.default_catalog,
                 promoted_snapshots_only=self.config.migration.promoted_snapshots_only,
+                pre_check_only=pre_check_only,
             )
         except Exception as e:
-            self.notification_target_manager.notify(
-                NotificationEvent.MIGRATION_FAILURE, traceback.format_exc()
-            )
+            if not pre_check_only:
+                self.notification_target_manager.notify(
+                    NotificationEvent.MIGRATION_FAILURE, traceback.format_exc()
+                )
             raise e
-        self.notification_target_manager.notify(NotificationEvent.MIGRATION_END)
+
+        if not pre_check_only:
+            self.notification_target_manager.notify(NotificationEvent.MIGRATION_END)
 
     @python_api_analytics
     def rollback(self) -> None:

sqlmesh/core/state_sync/base.py

@@ -8,7 +8,7 @@
 
 from sqlglot import __version__ as SQLGLOT_VERSION
 
-from sqlmesh import migrations
+from sqlmesh import migrations, pre_checks
 from sqlmesh.core.environment import (
     Environment,
     EnvironmentNamingInfo,
@@ -64,6 +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)
 
 
@@ -456,6 +460,7 @@ def migrate(
         default_catalog: t.Optional[str],
         skip_backup: bool = False,
         promoted_snapshots_only: bool = True,
+        pre_check_only: bool = False,
     ) -> None:
         """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
 

sqlmesh/core/state_sync/db/facade.py

@@ -447,13 +447,15 @@ def migrate(
         default_catalog: t.Optional[str],
         skip_backup: bool = False,
         promoted_snapshots_only: bool = True,
+        pre_check_only: bool = False,
     ) -> None:
         """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
         self.migrator.migrate(
             self,
             default_catalog,
             skip_backup=skip_backup,
             promoted_snapshots_only=promoted_snapshots_only,
+            pre_check_only=pre_check_only,
         )
 
     @transactional()

sqlmesh/core/state_sync/db/migrator.py

@@ -27,6 +27,7 @@
 )
 from sqlmesh.core.state_sync.base import (
     MIGRATIONS,
+    PRE_CHECKS,
 )
 from sqlmesh.core.state_sync.base import StateSync
 from sqlmesh.core.state_sync.db.environment import EnvironmentState
@@ -90,8 +91,22 @@ def migrate(
         default_catalog: t.Optional[str],
         skip_backup: bool = False,
         promoted_snapshots_only: bool = True,
+        pre_check_only: bool = False,
     ) -> None:
-        """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
+        """Migrate the state sync to the latest SQLMesh / SQLGlot version.
+
+        Args:
+            state_sync: The state sync instance.
+            default_catalog: The default catalog.
+            skip_backup: Whether to skip backing up state tables.
+            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)
+        should_migrate = self.console.log_pre_check_warnings(pre_check_warnings, pre_check_only)
+        if not should_migrate:
+            return
+
         versions = self.version_state.get_versions()
         migration_start_ts = time.perf_counter()
 
@@ -153,6 +168,33 @@ 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]]]:
+        """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.
+        """
+        # Get the range of the migrations that would be applied
+        from_version = self.version_state.get_versions().schema_version
+        to_version = len(MIGRATIONS)
+
+        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}")
+                warnings = pre_check(state_sync)
+                if warnings:
+                    pre_check_warnings.append((pre_check_name, warnings))
+
+        return pre_check_warnings
+
     def _apply_migrations(
         self,
         state_sync: StateSync,

sqlmesh/magics.py

@@ -709,11 +709,17 @@ def dag(self, context: Context, line: str) -> None:
         self.display(dag)
 
     @magic_arguments()
+    @argument(
+        "--pre-check",
+        action="store_true",
+        help="Run pre-checks and display warnings without performing migration",
+    )
     @line_magic
     @pass_sqlmesh_context
     def migrate(self, context: Context, line: str) -> None:
         """Migrate SQLMesh to the current running version."""
-        context.migrate()
+        args = parse_argstring(self.migrate, line)
+        context.migrate(pre_check_only=args.pre_check)
         context.console.log_success("Migration complete")
 
     @magic_arguments()