Fix: Drop the dev table clone if the schema migration fails (#2830)

TobikoData · Jun 26, 2024 · 322d742 · 322d742
1 parent 4bde314
commit 322d742
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
@@ -688,6 +688,9 @@ def _create_snapshot(
                         snapshot, alter_expressions, allow_destructive_snapshots
                     )
                     self.adapter.alter_table(alter_expressions)
+                except Exception:
+                    self.adapter.drop_table(target_table_name)
+                    raise
                 finally:
                     self.adapter.drop_table(tmp_table_name)
             else:

diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
@@ -1134,6 +1134,57 @@ def test_create_clone_in_dev(mocker: MockerFixture, adapter_mock, make_snapshot)
     )
 
 
+def test_drop_clone_in_dev_when_migration_fails(mocker: MockerFixture, adapter_mock, make_snapshot):
+    adapter_mock.SUPPORTS_CLONING = True
+    adapter_mock.get_alter_expressions.return_value = []
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    adapter_mock.alter_table.side_effect = Exception("Migration failed")
+
+    model = load_sql_based_model(
+        parse(  # type: ignore
+            """
+            MODEL (
+                name test_schema.test_model,
+                kind INCREMENTAL_BY_TIME_RANGE (
+                    time_column ds
+                )
+            );
+
+            SELECT 1::INT as a, ds::DATE FROM a;
+            """
+        ),
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+    snapshot.previous_versions = snapshot.all_versions
+
+    evaluator.create([snapshot], {})
+
+    adapter_mock.clone_table.assert_called_once_with(
+        f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp",
+        f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}",
+        replace=True,
+    )
+
+    adapter_mock.get_alter_expressions.assert_called_once_with(
+        f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp",
+        f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp__schema_migration_source",
+    )
+
+    adapter_mock.alter_table.assert_called_once_with([])
+
+    adapter_mock.drop_table.assert_has_calls(
+        [
+            call(f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp"),
+            call(
+                f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp__schema_migration_source"
+            ),
+        ]
+    )
+
+
 def test_create_clone_in_dev_self_referencing(mocker: MockerFixture, adapter_mock, make_snapshot):
     adapter_mock.SUPPORTS_CLONING = True
     adapter_mock.get_alter_expressions.return_value = []