From 69db7434b3f66b86f0446cc04a5faa63db091e53 Mon Sep 17 00:00:00 2001 From: George Sittas Date: Wed, 23 Jul 2025 16:56:38 +0300 Subject: [PATCH 1/7] Feat: introduce migration pre-checks --- sqlmesh/cli/main.py | 9 ++++-- sqlmesh/core/console.py | 45 ++++++++++++++++++++++++++ sqlmesh/core/context.py | 21 ++++++++---- sqlmesh/core/state_sync/base.py | 7 +++- sqlmesh/core/state_sync/db/facade.py | 2 ++ sqlmesh/core/state_sync/db/migrator.py | 44 ++++++++++++++++++++++++- sqlmesh/magics.py | 8 ++++- sqlmesh/pre_checks/__init__.py | 0 8 files changed, 125 insertions(+), 11 deletions(-) create mode 100644 sqlmesh/pre_checks/__init__.py diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index 8982efc9f8..83089ed7bc 100644 --- a/sqlmesh/cli/main.py +++ b/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") diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py index cf87fd7443..143d12e746 100644 --- a/sqlmesh/core/console.py +++ b/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, diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index c0d9b21ff8..04520b35c0 100644 --- a/sqlmesh/core/context.py +++ b/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: diff --git a/sqlmesh/core/state_sync/base.py b/sqlmesh/core/state_sync/base.py index 6c2097d760..54cd3bef81 100644 --- a/sqlmesh/core/state_sync/base.py +++ b/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.""" diff --git a/sqlmesh/core/state_sync/db/facade.py b/sqlmesh/core/state_sync/db/facade.py index 779add1cca..765545a334 100644 --- a/sqlmesh/core/state_sync/db/facade.py +++ b/sqlmesh/core/state_sync/db/facade.py @@ -447,6 +447,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.""" self.migrator.migrate( @@ -454,6 +455,7 @@ def migrate( default_catalog, skip_backup=skip_backup, promoted_snapshots_only=promoted_snapshots_only, + pre_check_only=pre_check_only, ) @transactional() diff --git a/sqlmesh/core/state_sync/db/migrator.py b/sqlmesh/core/state_sync/db/migrator.py index 405c0ea667..8ee71364ca 100644 --- a/sqlmesh/core/state_sync/db/migrator.py +++ b/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, diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py index 454b6cd4ce..328ece80bf 100644 --- a/sqlmesh/magics.py +++ b/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() diff --git a/sqlmesh/pre_checks/__init__.py b/sqlmesh/pre_checks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 From ae252fdebeea07f03bc423c27899909e2e25275f Mon Sep 17 00:00:00 2001 From: George Sittas Date: Wed, 30 Jul 2025 18:05:46 +0300 Subject: [PATCH 2/7] Refactor: move pre_check into the migration script --- sqlmesh/core/state_sync/base.py | 9 ++++---- sqlmesh/core/state_sync/db/migrator.py | 30 +++++++++++--------------- sqlmesh/pre_checks/__init__.py | 0 3 files changed, 16 insertions(+), 23 deletions(-) delete mode 100644 sqlmesh/pre_checks/__init__.py diff --git a/sqlmesh/core/state_sync/base.py b/sqlmesh/core/state_sync/base.py index 54cd3bef81..6edfda50ad 100644 --- a/sqlmesh/core/state_sync/base.py +++ b/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): diff --git a/sqlmesh/core/state_sync/db/migrator.py b/sqlmesh/core/state_sync/db/migrator.py index 8ee71364ca..a1484f79fc 100644 --- a/sqlmesh/core/state_sync/db/migrator.py +++ b/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 diff --git a/sqlmesh/pre_checks/__init__.py b/sqlmesh/pre_checks/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 From 0cc02efd40d3068391b8283eb610534059c74acd Mon Sep 17 00:00:00 2001 From: George Sittas Date: Wed, 30 Jul 2025 18:30:29 +0300 Subject: [PATCH 3/7] Alter version table to add pre_check_version --- sqlmesh/core/state_sync/base.py | 11 ++++++--- sqlmesh/core/state_sync/db/version.py | 9 ++++++- .../migrations/v0087_add_pre_check_version.py | 24 +++++++++++++++++++ 3 files changed, 40 insertions(+), 4 deletions(-) create mode 100644 sqlmesh/migrations/v0087_add_pre_check_version.py diff --git a/sqlmesh/core/state_sync/base.py b/sqlmesh/core/state_sync/base.py index 6edfda50ad..fa9841686b 100644 --- a/sqlmesh/core/state_sync/base.py +++ b/sqlmesh/core/state_sync/base.py @@ -40,6 +40,7 @@ class Versions(PydanticModel): schema_version: int = 0 sqlglot_version: str = "0.0.0" sqlmesh_version: str = "0.0.0" + pre_check_version: int = 0 @property def minor_sqlglot_version(self) -> t.Tuple[int, int]: @@ -54,9 +55,9 @@ def minor_sqlmesh_version(self) -> t.Tuple[int, int]: def _package_version_validator(cls, v: t.Any) -> str: return "0.0.0" if v is None else str(v) - @field_validator("schema_version", mode="before") + @field_validator("schema_version", "pre_check_version", mode="before") @classmethod - def _schema_version_validator(cls, v: t.Any) -> int: + def _int_version_validator(cls, v: t.Any) -> int: return 0 if v is None else int(v) @@ -66,7 +67,11 @@ def _schema_version_validator(cls, v: t.Any) -> int: ] SCHEMA_VERSION: int = len(MIGRATIONS) PRE_CHECK_VERSION: int = ( - max(idx for idx, migration in enumerate(MIGRATIONS) if hasattr(migration, "pre_check")) + 1 + max( + [idx for idx, migration in enumerate(MIGRATIONS) if hasattr(migration, "pre_check")], + default=-1, + ) + + 1 ) diff --git a/sqlmesh/core/state_sync/db/version.py b/sqlmesh/core/state_sync/db/version.py index 873e1633df..9374fd90a5 100644 --- a/sqlmesh/core/state_sync/db/version.py +++ b/sqlmesh/core/state_sync/db/version.py @@ -13,6 +13,7 @@ SQLMESH_VERSION, ) from sqlmesh.core.state_sync.base import ( + PRE_CHECK_VERSION, SCHEMA_VERSION, Versions, ) @@ -31,6 +32,7 @@ def __init__(self, engine_adapter: EngineAdapter, schema: t.Optional[str] = None "schema_version": exp.DataType.build("int"), "sqlglot_version": exp.DataType.build(index_type), "sqlmesh_version": exp.DataType.build(index_type), + "pre_check_version": exp.DataType.build("int"), } def update_versions( @@ -38,6 +40,7 @@ def update_versions( schema_version: int = SCHEMA_VERSION, sqlglot_version: str = SQLGLOT_VERSION, sqlmesh_version: str = SQLMESH_VERSION, + pre_check_version: int = PRE_CHECK_VERSION, ) -> None: import pandas as pd @@ -51,6 +54,7 @@ def update_versions( "schema_version": schema_version, "sqlglot_version": sqlglot_version, "sqlmesh_version": sqlmesh_version, + "pre_check_version": pre_check_version, } ] ), @@ -69,5 +73,8 @@ def get_versions(self) -> Versions: return no_version return Versions( - schema_version=row[0], sqlglot_version=row[1], sqlmesh_version=seq_get(row, 2) + schema_version=row[0], + sqlglot_version=row[1], + sqlmesh_version=seq_get(row, 2), + pre_check_version=seq_get(row, 3), ) diff --git a/sqlmesh/migrations/v0087_add_pre_check_version.py b/sqlmesh/migrations/v0087_add_pre_check_version.py new file mode 100644 index 0000000000..d470b66f02 --- /dev/null +++ b/sqlmesh/migrations/v0087_add_pre_check_version.py @@ -0,0 +1,24 @@ +"""Add new 'pre_check_version' column to the version state table.""" + +from sqlglot import exp + + +def migrate(state_sync, **kwargs): # type: ignore + engine_adapter = state_sync.engine_adapter + schema = state_sync.schema + versions_table = "_versions" + if schema: + versions_table = f"{schema}.{versions_table}" + + alter_table_exp = exp.Alter( + this=exp.to_table(versions_table), + kind="TABLE", + actions=[ + exp.ColumnDef( + this=exp.to_column("pre_check_version"), + kind=exp.DataType.build("int"), + ) + ], + ) + + engine_adapter.execute(alter_table_exp) From f9dbc72c9193b541bd8e1be7ee122f67915564c5 Mon Sep 17 00:00:00 2001 From: George Sittas Date: Wed, 30 Jul 2025 21:50:47 +0300 Subject: [PATCH 4/7] Add tests --- tests/core/state_sync/test_state_sync.py | 90 +++++++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/tests/core/state_sync/test_state_sync.py b/tests/core/state_sync/test_state_sync.py index a5a6969e38..a064771961 100644 --- a/tests/core/state_sync/test_state_sync.py +++ b/tests/core/state_sync/test_state_sync.py @@ -2,6 +2,7 @@ import logging import re import typing as t +from types import ModuleType from unittest.mock import call, patch import duckdb # noqa: TID253 @@ -11,8 +12,16 @@ from pytest_mock.plugin import MockerFixture from sqlglot import exp +from sqlmesh.cli.project_init import init_example_project from sqlmesh.core import constants as c -from sqlmesh.core.config import EnvironmentSuffixTarget +from sqlmesh.core.config import ( + Config, + DuckDBConnectionConfig, + EnvironmentSuffixTarget, + GatewayConfig, + ModelDefaultsConfig, +) +from sqlmesh.core.context import Context from sqlmesh.core.dialect import parse_one, schema_ from sqlmesh.core.engine_adapter import create_engine_adapter from sqlmesh.core.environment import Environment, EnvironmentStatements @@ -48,6 +57,7 @@ ) from sqlmesh.utils.date import now_timestamp, to_datetime, to_timestamp from sqlmesh.utils.errors import SQLMeshError +from tests.utils.test_helpers import use_terminal_console pytestmark = pytest.mark.slow @@ -3629,3 +3639,81 @@ def test_update_environment_statements(state_sync: EngineAdapterStateSync): "@grant_schema_usage()", "@grant_select_privileges()", ] + + +@use_terminal_console +def test_pre_checks(tmp_path, mocker): + init_example_project(tmp_path, engine_type="duckdb") + + db_path = str(tmp_path / "db.db") + config = Config( + gateways={"main": GatewayConfig(connection=DuckDBConnectionConfig(database=db_path))}, + model_defaults=ModelDefaultsConfig(dialect="duckdb"), + ) + context = Context(paths=tmp_path, config=config) + context.plan(auto_apply=True, no_prompts=True) + + def mock_migrate(state_sync, **kwargs): + pass + + def mock_pre_check_with_warnings(state_sync): + return [ + "Warning: This migration will break compatibility with older versions", + "Warning: You must update all model configurations before applying this migration", + "Warning: Existing snapshots will need to be rebuilt", + ] + + def mock_pre_check_without_warnings(state_sync): + return [] + + # Create a mock migration module with a pre_check function + mock_migration = ModuleType("v9999_test_pre_check") + + setattr(mock_migration, "migrate", mock_migrate) + setattr(mock_migration, "pre_check", mock_pre_check_with_warnings) + + versions_before_migrate = context.state_sync.get_versions() + + import sqlmesh.core.state_sync as state_sync + + test_migrations = state_sync.db.migrator.MIGRATIONS + [mock_migration] + + # Test 1: Pre-check warnings are properly collected and displayed, user rejects migration + with ( + patch.object(state_sync.db.migrator, "MIGRATIONS", test_migrations), + patch.object(context.console, "_confirm", return_value=False), + ): + console = context.console + log_pre_check_warnings_spy = mocker.spy(console, "log_pre_check_warnings") + + context.migrate(pre_check_only=False) + + calls = log_pre_check_warnings_spy.mock_calls + assert len(calls) == 1 + + pre_check_warnings = calls[0].args[0] + assert len(pre_check_warnings) == 1 + + assert pre_check_warnings[0][0] == "v9999_test_pre_check" + assert len(pre_check_warnings[0][1]) == 3 + assert all(warning.startswith("Warning:") for warning in pre_check_warnings[0][1]) + + assert context.state_sync.get_versions() == versions_before_migrate + + update_versions_spy = mocker.spy(state_sync.db.version.VersionState, "update_versions") + + # Test 2: User accepts migration after being notified about pre-check warnings + with ( + patch.object(state_sync.db.migrator, "MIGRATIONS", test_migrations), + patch.object(context.console, "_confirm", return_value=True), + ): + context.migrate(pre_check_only=False) + assert len(update_versions_spy.mock_calls) == 1 + + # Test 3: Pre-check without warning should automatically reuslt in a migration + setattr(mock_migration, "pre_check", mock_pre_check_without_warnings) + with patch.object(state_sync.db.migrator, "MIGRATIONS", test_migrations): + # Since the version module's SCHEMA_VERSION, etc, weren't patched, the old versions + # are still used, so the following should result in hitting the update_versions path + context.migrate(pre_check_only=False) + assert len(update_versions_spy.mock_calls) == 2 From c669436a5b91ffe6899a343542d4ed83efa64668 Mon Sep 17 00:00:00 2001 From: George Sittas Date: Wed, 30 Jul 2025 22:14:12 +0300 Subject: [PATCH 5/7] PR feedback --- sqlmesh/core/console.py | 27 ++++++------------------ sqlmesh/core/state_sync/db/migrator.py | 4 ++-- tests/core/state_sync/test_state_sync.py | 7 ++---- 3 files changed, 11 insertions(+), 27 deletions(-) diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py index 143d12e746..6ad18334e6 100644 --- a/sqlmesh/core/console.py +++ b/sqlmesh/core/console.py @@ -498,13 +498,9 @@ 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: + def log_pre_check_warnings(self, pre_check_warnings: t.List[str], pre_check_only: bool) -> bool: """ - Log warnings emitted by pre-check scripts and ask user whether they'd like to + Log warnings emitted by pre-checks and ask user whether they'd like to proceed with the migration (true) or not (false). """ @@ -673,11 +669,7 @@ 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: + def log_pre_check_warnings(self, pre_check_warnings: t.List[str], pre_check_only: bool) -> bool: return True def start_state_export( @@ -1490,16 +1482,11 @@ 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: + def log_pre_check_warnings(self, pre_check_warnings: 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]") + tree = Tree(f"[bold]Pre-migration warnings[/bold]") + for warning in pre_check_warnings: + tree.add(f"[yellow]{warning}[/yellow]") self._print(tree) diff --git a/sqlmesh/core/state_sync/db/migrator.py b/sqlmesh/core/state_sync/db/migrator.py index a1484f79fc..770ce57f66 100644 --- a/sqlmesh/core/state_sync/db/migrator.py +++ b/sqlmesh/core/state_sync/db/migrator.py @@ -165,7 +165,7 @@ 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[str]: """Run pre-checks for migrations between specified versions. Args: @@ -185,7 +185,7 @@ def _run_pre_checks(self, state_sync: StateSync) -> t.List[t.Tuple[str, t.List[s logger.info(f"Running pre-check for {migration_name}") warnings = pre_check(state_sync) if warnings: - pre_check_warnings.append((migration_name, warnings)) + pre_check_warnings.extend(warnings) return pre_check_warnings diff --git a/tests/core/state_sync/test_state_sync.py b/tests/core/state_sync/test_state_sync.py index a064771961..7e27634c2b 100644 --- a/tests/core/state_sync/test_state_sync.py +++ b/tests/core/state_sync/test_state_sync.py @@ -3692,11 +3692,8 @@ def mock_pre_check_without_warnings(state_sync): assert len(calls) == 1 pre_check_warnings = calls[0].args[0] - assert len(pre_check_warnings) == 1 - - assert pre_check_warnings[0][0] == "v9999_test_pre_check" - assert len(pre_check_warnings[0][1]) == 3 - assert all(warning.startswith("Warning:") for warning in pre_check_warnings[0][1]) + assert len(pre_check_warnings) == 3 + assert all(warning.startswith("Warning:") for warning in pre_check_warnings) assert context.state_sync.get_versions() == versions_before_migrate From d4003183df67476b2609d4b06777da3a86fd4d1e Mon Sep 17 00:00:00 2001 From: George Sittas Date: Thu, 31 Jul 2025 19:37:08 +0300 Subject: [PATCH 6/7] Rename migration script to v0088 --- ...87_add_pre_check_version.py => v0088_add_pre_check_version.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename sqlmesh/migrations/{v0087_add_pre_check_version.py => v0088_add_pre_check_version.py} (100%) diff --git a/sqlmesh/migrations/v0087_add_pre_check_version.py b/sqlmesh/migrations/v0088_add_pre_check_version.py similarity index 100% rename from sqlmesh/migrations/v0087_add_pre_check_version.py rename to sqlmesh/migrations/v0088_add_pre_check_version.py From fd045e04bf3c5eb03a6ce190b6aa496f8ca124f1 Mon Sep 17 00:00:00 2001 From: George Sittas Date: Tue, 12 Aug 2025 16:08:01 +0300 Subject: [PATCH 7/7] Rename pre-check migration script --- ...88_add_pre_check_version.py => v0089_add_pre_check_version.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename sqlmesh/migrations/{v0088_add_pre_check_version.py => v0089_add_pre_check_version.py} (100%) diff --git a/sqlmesh/migrations/v0088_add_pre_check_version.py b/sqlmesh/migrations/v0089_add_pre_check_version.py similarity index 100% rename from sqlmesh/migrations/v0088_add_pre_check_version.py rename to sqlmesh/migrations/v0089_add_pre_check_version.py