feat(cli): report state backend versions on rollback and info (#6088)

Signed-off-by: Adegbite Ayoade <tripleaceme@gmail.com>

Author: tripleaceme

sqlmesh/core/context.py

@@ -108,6 +108,7 @@
     CachingStateSync,
     StateReader,
     StateSync,
+    Versions,
 )
 from sqlmesh.core.janitor import cleanup_expired_views, delete_expired_snapshots
 from sqlmesh.core.table_diff import TableDiff
@@ -2610,15 +2611,18 @@ def migrate(self) -> None:
         """
         self.notification_target_manager.notify(NotificationEvent.MIGRATION_START)
         self._load_materializations()
+        state_sync = self._new_state_sync()
+        previous_versions = self._state_versions(state_sync)
         try:
-            self._new_state_sync().migrate(
+            state_sync.migrate(
                 promoted_snapshots_only=self.config.migration.promoted_snapshots_only,
             )
         except Exception as e:
             self.notification_target_manager.notify(
                 NotificationEvent.MIGRATION_FAILURE, traceback.format_exc()
             )
             raise e
+        self._print_state_versions(self._state_versions(state_sync), previous_versions)
         self.notification_target_manager.notify(NotificationEvent.MIGRATION_END)
 
     @python_api_analytics
@@ -2627,7 +2631,10 @@ def rollback(self) -> None:
 
         Please contact your SQLMesh administrator before doing this. This action cannot be undone.
         """
-        self._new_state_sync().rollback()
+        state_sync = self._new_state_sync()
+        previous_versions = self._state_versions(state_sync)
+        state_sync.rollback()
+        self._print_state_versions(self._state_versions(state_sync), previous_versions)
 
     @python_api_analytics
     def create_external_models(self, strict: bool = False) -> None:
@@ -2701,6 +2708,12 @@ def print_info(
         if state_connection:
             self._try_connection("state backend", state_connection.connection_validator())
 
+        if verbosity >= Verbosity.VERBOSE:
+            try:
+                self._print_state_versions(self._state_versions())
+            except Exception as ex:
+                self.console.log_error(f"Failed to fetch the state backend versions. {ex}")
+
     @python_api_analytics
     def print_environment_names(self) -> None:
         """Prints all environment names along with expiry datetime."""
@@ -3290,6 +3303,25 @@ def _try_connection(self, connection_name: str, validator: t.Callable[[], None])
         except Exception as ex:
             self.console.log_error(f"{connection_name} connection failed. {ex}")
 
+    def _state_versions(self, state_sync: t.Optional[StateSync] = None) -> Versions:
+        """Returns the versions recorded in the state backend without validating them."""
+        return (state_sync or self._new_state_sync()).get_versions(validate=False)
+
+    def _print_state_versions(
+        self, versions: Versions, previous_versions: t.Optional[Versions] = None
+    ) -> None:
+        """Prints the state backend versions, optionally alongside the ones they replaced."""
+        self.console.log_status_update("\nState backend versions:")
+        for label, attribute in (
+            ("Schema version", "schema_version"),
+            ("SQLGlot version", "sqlglot_version"),
+            ("SQLMesh version", "sqlmesh_version"),
+        ):
+            version = getattr(versions, attribute)
+            if previous_versions is not None:
+                version = f"{getattr(previous_versions, attribute)} -> {version}"
+            self.console.log_status_update(f"{label}: {version}")
+
     def _new_state_sync(self) -> StateSync:
         return self._provided_state_sync or self._scheduler.create_state_sync(self)
 

sqlmesh/core/state_sync/db/migrator.py

@@ -96,6 +96,15 @@ def migrate(
             migrate_rows = self._apply_migrations(schema, skip_backup)
 
             if not migrate_rows and major_minor(SQLMESH_VERSION) == versions.minor_sqlmesh_version:
+                # Nothing to migrate, but a patch-level bump still leaves the recorded versions
+                # behind what is actually running, so they are brought up to date here. The
+                # schema version is carried over rather than defaulted, since no migration ran
+                # and moving it could hide one that is genuinely needed later.
+                if (
+                    versions.sqlmesh_version != SQLMESH_VERSION
+                    or versions.sqlglot_version != SQLGLOT_VERSION
+                ):
+                    self.version_state.update_versions(schema_version=versions.schema_version)
                 return
 
             if migrate_rows:

tests/cli/test_cli.py

@@ -10,11 +10,15 @@
 
 from click import ClickException
 from click.testing import CliRunner
+from sqlglot import __version__ as SQLGLOT_VERSION
 from sqlmesh import RuntimeEnv
+from sqlmesh._version import __version__ as SQLMESH_VERSION
 from sqlmesh.cli.project_init import ProjectTemplate, init_example_project
 from sqlmesh.cli.main import cli
 from sqlmesh.core.context import Context
+from sqlmesh.core.state_sync.base import SCHEMA_VERSION
 from sqlmesh.integrations.dlt import generate_dlt_models
+from sqlmesh.utils import major_minor
 from sqlmesh.utils.date import now_ds, time_like_to_str, timedelta, to_datetime, yesterday_ds
 from sqlmesh.core.config.connection import DIALECT_TO_TYPE
 
@@ -1024,6 +1028,107 @@ def test_info_on_new_project_does_not_create_state_sync(runner, tmp_path):
     assert not context.engine_adapter.table_exists("sqlmesh._versions")
 
 
+def test_info_state_versions(runner, tmp_path):
+    create_example_project(tmp_path)
+    init_prod_and_backfill(runner, tmp_path)
+
+    result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "info"])
+    assert result.exit_code == 0
+    assert "State backend versions" not in result.output
+
+    result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "info", "-v"])
+    assert result.exit_code == 0
+    assert "State backend versions" in result.output
+    assert f"Schema version: {SCHEMA_VERSION}" in result.output
+    assert f"SQLGlot version: {SQLGLOT_VERSION}" in result.output
+    assert f"SQLMesh version: {SQLMESH_VERSION}" in result.output
+
+
+def test_rollback_state_versions(runner, tmp_path):
+    create_example_project(tmp_path)
+    init_prod_and_backfill(runner, tmp_path)
+
+    context = Context(paths=tmp_path)
+    state_sync = context._new_state_sync()
+    # Back up the current state, then pretend the state was migrated by a newer SQLMesh.
+    state_sync.migrator._backup_state()
+    state_sync.version_state.update_versions(
+        schema_version=SCHEMA_VERSION + 1,
+        sqlglot_version="9999.0.0",
+        sqlmesh_version="9999.0.0",
+    )
+    context.close()
+
+    result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "rollback"])
+    assert result.exit_code == 0
+    assert "State backend versions" in result.output
+    assert f"Schema version: {SCHEMA_VERSION + 1} -> {SCHEMA_VERSION}" in result.output
+    assert f"SQLGlot version: 9999.0.0 -> {SQLGLOT_VERSION}" in result.output
+    assert f"SQLMesh version: 9999.0.0 -> {SQLMESH_VERSION}" in result.output
+
+
+def test_migrate_state_versions(runner, tmp_path):
+    create_example_project(tmp_path)
+    init_prod_and_backfill(runner, tmp_path)
+
+    context = Context(paths=tmp_path)
+    # Pretend the state was written by an older patch release of the same minor version, which
+    # is the case `migrate` used to leave untouched.
+    major, minor = major_minor(SQLMESH_VERSION)
+    older_sqlmesh = f"{major}.{minor}.dev0"
+    context._new_state_sync().version_state.update_versions(
+        sqlglot_version="0.0.1",
+        sqlmesh_version=older_sqlmesh,
+    )
+    context.close()
+
+    result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "migrate"])
+    assert result.exit_code == 0
+    assert "State backend versions" in result.output
+    assert f"SQLGlot version: 0.0.1 -> {SQLGLOT_VERSION}" in result.output
+    assert f"SQLMesh version: {older_sqlmesh} -> {SQLMESH_VERSION}" in result.output
+
+
+def test_migrate_updates_versions_after_a_patch_bump(runner, tmp_path):
+    """A patch bump leaves the minor version equal, but the recorded versions must still move.
+
+    Both minor versions have to match the installed ones, otherwise `_apply_migrations` reports
+    rows to migrate and the early return this covers is never reached.
+    """
+    create_example_project(tmp_path)
+    init_prod_and_backfill(runner, tmp_path)
+
+    sqlmesh_major, sqlmesh_minor = major_minor(SQLMESH_VERSION)
+    sqlglot_major, sqlglot_minor = major_minor(SQLGLOT_VERSION)
+    context = Context(paths=tmp_path)
+    context._new_state_sync().version_state.update_versions(
+        sqlglot_version=f"{sqlglot_major}.{sqlglot_minor}.dev0",
+        sqlmesh_version=f"{sqlmesh_major}.{sqlmesh_minor}.dev0",
+    )
+    context.close()
+
+    assert (
+        runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "migrate"]).exit_code
+        == 0
+    )
+
+    context = Context(paths=tmp_path)
+    versions = context._new_state_sync().get_versions(validate=False)
+    context.close()
+    assert versions.sqlmesh_version == SQLMESH_VERSION
+    assert versions.sqlglot_version == SQLGLOT_VERSION
+
+
+def test_rollback_without_backup_does_not_print_state_versions(runner, tmp_path):
+    create_example_project(tmp_path)
+    init_prod_and_backfill(runner, tmp_path)
+
+    result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "rollback"])
+    assert result.exit_code == 1
+    assert "There are no prior migrations to roll back to." in result.output
+    assert "State backend versions" not in result.output
+
+
 def test_dlt_pipeline_errors(runner, tmp_path):
     # Error if no pipeline is provided
     result = runner.invoke(cli, ["--paths", tmp_path, "init", "-t", "dlt", "duckdb"])

tests/core/state_sync/test_state_sync.py

@@ -4254,3 +4254,35 @@ def test_state_version_is_too_old(
         match="The current state belongs to an old version of SQLMesh that is no longer supported. Please upgrade to 0.134.0 first before upgrading to.*",
     ):
         state_sync.migrate(skip_backup=True)
+
+
+def test_migrate_patch_bump_preserves_schema_version(
+    state_sync: EngineAdapterStateSync, mocker: MockerFixture
+) -> None:
+    """A run with nothing to migrate must not move the recorded schema version.
+
+    `_apply_migrations` is forced to report no rows so the patch-bump branch is reached with a
+    schema version that differs from the current one. A real state can't get into that shape,
+    which is the point: the schema version is carried over rather than defaulted so that a
+    migration which is genuinely still needed can't be masked.
+    """
+    from sqlmesh import __version__ as SQLMESH_VERSION
+    from sqlmesh.core.state_sync.base import SCHEMA_VERSION
+
+    stale_schema_version = SCHEMA_VERSION - 1
+    state_sync.version_state.update_versions(
+        schema_version=stale_schema_version,
+        sqlglot_version="0.0.1",
+        sqlmesh_version=SQLMESH_VERSION,
+    )
+
+    mocker.patch(
+        "sqlmesh.core.state_sync.db.migrator.StateMigrator._apply_migrations",
+        return_value=False,
+    )
+
+    state_sync.migrate()
+
+    versions = state_sync.get_versions(validate=False)
+    assert versions.schema_version == stale_schema_version
+    assert versions.sqlglot_version == SQLGLOT_VERSION

tests/integrations/jupyter/test_magics.py

@@ -12,8 +12,11 @@
 from IPython.utils.capture import CapturedIO, capture_output
 from pytest_mock.plugin import MockerFixture
 from rich.console import Console as RichConsole
+from sqlglot import __version__ as SQLGLOT_VERSION
 
 from sqlmesh import Context, RuntimeEnv
+from sqlmesh._version import __version__ as SQLMESH_VERSION
+from sqlmesh.core.state_sync.base import SCHEMA_VERSION, Versions
 from sqlmesh.magics import register_magics
 from pathlib import Path
 
@@ -740,14 +743,19 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
 
     assert not output.stdout
     assert not output.stderr
-    assert len(output.outputs) == 6
+    assert len(output.outputs) == 10
+    # No plan has been applied, so the state backend is still empty and reports the defaults.
     assert convert_all_html_output_to_text(output) == [
         "Models: 20",
         "Macros: 8",
         "",
         "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []\n  shared_connection: true",
         "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []\n  shared_connection: true",
         "Data warehouse connection succeeded",
+        "State backend versions:",
+        "Schema version: 0",
+        "SQLGlot version: 0.0.0",
+        "SQLMesh version: 0.0.0",
     ]
     assert get_all_html_output(output) == [
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Models: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">20</span></pre>",
@@ -756,6 +764,10 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
         '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span>  shared_connection: true</pre>',
         '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span>  shared_connection: true</pre>',
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Data warehouse connection <span style=\"color: #008000; text-decoration-color: #008000\">succeeded</span></pre>",
+        "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">State backend versions:</pre>",
+        "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Schema version: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span></pre>",
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">SQLGlot version: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0.0</span>.<span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0</span></pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">SQLMesh version: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0.0</span>.<span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0</span></pre>',
     ]
 
 
@@ -768,25 +780,36 @@ def test_migrate(
 
     assert not output.stdout
     assert not output.stderr
-    assert len(output.outputs) == 1
+    assert len(output.outputs) == 5
+    # The sushi state lives in an in-memory DuckDB database, so the state sync that `migrate`
+    # opens starts empty and the versions move from the defaults to the running ones.
+    empty = Versions()
     assert convert_all_html_output_to_text(output) == [
+        "State backend versions:",
+        f"Schema version: {empty.schema_version} -> {SCHEMA_VERSION}",
+        f"SQLGlot version: {empty.sqlglot_version} -> {SQLGLOT_VERSION}",
+        f"SQLMesh version: {empty.sqlmesh_version} -> {SQLMESH_VERSION}",
         "Migration complete",
     ]
-    assert get_all_html_output(output) == [
-        str(
+    # Rich highlights the numbers inside the version lines, and that markup depends on the
+    # running versions, so only the fixed lines are compared as HTML.
+    html_output = get_all_html_output(output)
+    assert html_output[0] == str(
+        h("pre", {"style": RICH_PRE_STYLE}, "State backend versions:", autoescape=False)
+    )
+    assert html_output[-1] == str(
+        h(
+            "pre",
+            {"style": RICH_PRE_STYLE},
             h(
-                "pre",
-                {"style": RICH_PRE_STYLE},
-                h(
-                    "span",
-                    {"style": SUCCESS_STYLE},
-                    "Migration complete",
-                    autoescape=False,
-                ),
+                "span",
+                {"style": SUCCESS_STYLE},
+                "Migration complete",
                 autoescape=False,
-            )
+            ),
+            autoescape=False,
         )
-    ]
+    )
 
 
 # TODO: Add test for rollback