Repository: TobikoData/sqlmesh (python)

Issue:
Allow rollbacks for SQLMesh updates
As part of our go-live readiness checks we simulated a migration and rollback from SQLMesh 0.125.6 to 0.126.2 and then was unable to rollback.

The migration process seems to only create backup tables if there's a schema version change or if there's a sqlglot major/minor change.
If SQLMesh has been updated without one of these 2 things also being satisfied then the backup tables are not created.
Because the backup tables don't exist then SQLMesh throws an exception when trying to rollback.
We then cannot proceed without manual intervention.

I believe I've found a simple one line fix, PR to follow.

Pull request description:
fix: consider sqlmesh major/minor for migrate and rollback
Fixes #3445

I was considering on top of this if it would be sensible to extend/copy the circle ci migration test.

Patch:
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index e97c5c95..1639fd88 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -674,8 +674,6 @@ def format_model_expressions(
     if len(expressions) == 1 and is_meta_expression(expressions[0]):
         return expressions[0].sql(pretty=True, dialect=dialect)
 
-    *statements, query = expressions
-
     if rewrite_casts:
 
         def cast_to_colon(node: exp.Expression) -> exp.Expression:
@@ -696,14 +694,16 @@ def format_model_expressions(
             exp.replace_children(node, cast_to_colon)
             return node
 
-        query = query.copy()
-        exp.replace_children(query, cast_to_colon)
+        new_expressions = []
+        for expression in expressions:
+            expression = expression.copy()
+            exp.replace_children(expression, cast_to_colon)
+            new_expressions.append(expression)
+
+        expressions = new_expressions
 
     return ";\n\n".join(
-        [
-            *(statement.sql(pretty=True, dialect=dialect, **kwargs) for statement in statements),
-            query.sql(pretty=True, dialect=dialect, **kwargs),
-        ]
+        expression.sql(pretty=True, dialect=dialect, **kwargs) for expression in expressions
     ).strip()
 
 
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index 9afb0aab..3c477d7a 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -1307,16 +1307,24 @@ class EngineAdapterStateSync(StateSync):
     ) -> bool:
         versions = self.get_versions(validate=False)
         migrations = MIGRATIONS[versions.schema_version :]
-
-        migrate_rows = migrations or major_minor(SQLGLOT_VERSION) != versions.minor_sqlglot_version
-        if not skip_backup and migrate_rows:
+        should_backup = any(
+            [
+                migrations,
+                major_minor(SQLGLOT_VERSION) != versions.minor_sqlglot_version,
+                major_minor(SQLMESH_VERSION) != versions.minor_sqlmesh_version,
+            ]
+        )
+        if not skip_backup and should_backup:
             self._backup_state()
 
         for migration in migrations:
             logger.info(f"Applying migration {migration}")
             migration.migrate(self, default_catalog=default_catalog)
 
-        return bool(migrate_rows)
+        migrate_snapshots_and_environments = (
+            bool(migrations) or major_minor(SQLGLOT_VERSION) != versions.minor_sqlglot_version
+        )
+        return migrate_snapshots_and_environments
 
     def _migrate_rows(self, promoted_snapshots_only: bool) -> None:
         logger.info("Fetching environments")


Test patch:
diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
index 195c5b6a..1d6ca0fe 100644
--- a/tests/core/test_dialect.py
+++ b/tests/core/test_dialect.py
@@ -76,7 +76,12 @@ def test_format_model_expressions():
     sum(n + 1)::int as n, -- n
     o,
     p + 1,
-    CAST(x as int)::int,
+    CAST(x as int)::int,;
+
+@IF(
+    @runtime_stage = 'creating',
+    GRANT SELECT ON foo.bar TO "bla"
+)
     """
         )
     )
@@ -151,7 +156,9 @@ SELECT
   SUM(n + 1)::INT AS n, /* n */
   o,
   p + 1,
-  x::INT::INT"""
+  x::INT::INT;
+
+@IF(@runtime_stage = 'creating', GRANT SELECT ON foo.bar TO "bla")"""
     )
 
     x = format_model_expressions(