diff --git a/sqlmesh/core/state_sync/db/migrator.py b/sqlmesh/core/state_sync/db/migrator.py index aaabeac5d9..dc4022fb1a 100644 --- a/sqlmesh/core/state_sync/db/migrator.py +++ b/sqlmesh/core/state_sync/db/migrator.py @@ -120,19 +120,25 @@ def migrate( if skip_backup: logger.error("Backup was skipped so no rollback was attempted.") else: - self.rollback() + try: + self.rollback() + except Exception: + logger.exception("Migration rollback failed.") - analytics.collector.on_migration_end( - from_sqlmesh_version=versions.sqlmesh_version, - state_sync_type=self.engine_adapter.dialect, - migration_time_sec=time.perf_counter() - migration_start_ts, - error=e, - ) + try: + analytics.collector.on_migration_end( + from_sqlmesh_version=versions.sqlmesh_version, + state_sync_type=self.engine_adapter.dialect, + migration_time_sec=time.perf_counter() - migration_start_ts, + error=e, + ) + except Exception: + logger.exception("Failed to record the migration end analytics event.") self.console.log_migration_status(success=False) if isinstance(e, StateMigrationError): raise - raise SQLMeshError("SQLMesh migration failed.") from e + raise SQLMeshError(f"SQLMesh migration failed: {e}") from e self.console.log_migration_status() diff --git a/tests/core/state_sync/test_state_sync.py b/tests/core/state_sync/test_state_sync.py index 32a359f667..f76211663f 100644 --- a/tests/core/state_sync/test_state_sync.py +++ b/tests/core/state_sync/test_state_sync.py @@ -11,6 +11,7 @@ from pytest_mock.plugin import MockerFixture from sqlglot import exp +from sqlmesh.core import analytics from sqlmesh.core import constants as c from sqlmesh.core.config import EnvironmentSuffixTarget from sqlmesh.core.dialect import parse_one @@ -2854,6 +2855,28 @@ def test_first_migration_failure(duck_conn, mocker: MockerFixture, tmp_path) -> assert not state_sync.engine_adapter.table_exists(state_sync.interval_state.intervals_table) +def test_migration_failure_after_update_versions( + duck_conn, mocker: MockerFixture, tmp_path +) -> None: + state_sync = EngineAdapterStateSync( + create_engine_adapter(lambda: duck_conn, "duckdb"), + schema=c.SQLMESH, + cache_dir=tmp_path / c.CACHE, + ) + # The first call happens inside the try block right after update_versions() and + # raises; the second call happens in the exception handler and succeeds. + mocker.patch.object( + analytics.collector, + "on_migration_end", + side_effect=[Exception("mocked error"), None], + ) + with pytest.raises(SQLMeshError, match="SQLMesh migration failed.") as exc_info: + state_sync.migrate() + assert "mocked error" in str(exc_info.value) + assert isinstance(exc_info.value.__cause__, Exception) + assert "mocked error" in str(exc_info.value.__cause__) + + def test_migrate_rows(state_sync: EngineAdapterStateSync, mocker: MockerFixture) -> None: state_sync.engine_adapter.replace_query( "sqlmesh._versions",