fix: changing catalog env override existing catalog matches (#2025)

* fix: changing catalog env override existing catalog matches

* rewrite environment view change detection

Author: eakmanrq

sqlmesh/core/state_sync/common.py

@@ -15,6 +15,7 @@
     SnapshotIdLike,
     SnapshotInfoLike,
     SnapshotNameVersionLike,
+    SnapshotTableInfo,
     start_date,
 )
 from sqlmesh.core.state_sync.base import PromotionResult, StateSync
@@ -128,77 +129,62 @@ def promote(
 
         existing_environment = self._get_environment(environment.name, lock_for_update=True)
 
-        environment_suffix_target_changed = False
-        environment_catalog_changed = False
-        removed_environment_naming_info = environment.naming_info
+        existing_table_infos = (
+            {table_info.name: table_info for table_info in existing_environment.promoted_snapshots}
+            if existing_environment
+            else {}
+        )
+        table_infos = {table_info.name: table_info for table_info in environment.promoted_snapshots}
+        views_that_changed_location: t.Set[SnapshotTableInfo] = set()
         if existing_environment:
+            views_that_changed_location = {
+                existing_table_info
+                for name, existing_table_info in existing_table_infos.items()
+                if name in table_infos
+                and existing_table_info.qualified_view_name.for_environment(
+                    existing_environment.naming_info
+                )
+                != table_infos[name].qualified_view_name.for_environment(environment.naming_info)
+            }
             if environment.previous_plan_id != existing_environment.plan_id:
                 raise SQLMeshError(
                     f"Plan '{environment.plan_id}' is no longer valid for the target environment '{environment.name}'. "
                     f"Expected previous plan ID: '{environment.previous_plan_id}', actual previous plan ID: '{existing_environment.plan_id}'. "
                     "Please recreate the plan and try again"
                 )
-
-            environment_suffix_target_changed = (
-                environment.suffix_target != existing_environment.suffix_target
-            )
-            if environment_suffix_target_changed:
-                removed_environment_naming_info.suffix_target = existing_environment.suffix_target
-
-            environment_catalog_changed = (
-                environment.catalog_name_override != existing_environment.catalog_name_override
-            )
-            if environment_catalog_changed:
-                removed_environment_naming_info.catalog_name_override = (
-                    existing_environment.catalog_name_override
-                )
-
             if no_gaps_snapshot_names != set():
                 snapshots = self._get_snapshots(environment.snapshots).values()
                 self._ensure_no_gaps(
                     snapshots,
                     existing_environment,
                     no_gaps_snapshot_names,
                 )
-
-            existing_table_infos = {
-                table_info.name: table_info
-                for table_info in existing_environment.promoted_snapshots
-            }
-
             demoted_snapshots = set(existing_environment.snapshots) - set(environment.snapshots)
             for demoted_snapshot in self._get_snapshots(demoted_snapshots).values():
                 # Update the updated_at attribute.
                 self._update_snapshot(demoted_snapshot)
-        else:
-            existing_table_infos = {}
 
         missing_models = set(existing_table_infos) - {
             snapshot.name for snapshot in environment.promoted_snapshots
         }
 
-        table_infos = set(environment.promoted_snapshots)
-        if (
-            existing_environment
-            and existing_environment.finalized_ts
-            and not environment_suffix_target_changed
-            and not environment_catalog_changed
-        ):
+        added_table_infos = set(table_infos.values())
+        if existing_environment and existing_environment.finalized_ts:
             # Only promote new snapshots.
-            table_infos -= set(existing_environment.promoted_snapshots)
+            added_table_infos -= set(existing_environment.promoted_snapshots)
 
         self._update_environment(environment)
 
-        removed = (
-            list(existing_table_infos.values())
-            if environment_suffix_target_changed or environment_catalog_changed
-            else [existing_table_infos[name] for name in missing_models]
+        removed = {existing_table_infos[name] for name in missing_models}.union(
+            views_that_changed_location
         )
 
         return PromotionResult(
-            added=sorted(table_infos),
-            removed=removed,
-            removed_environment_naming_info=removed_environment_naming_info if removed else None,
+            added=sorted(added_table_infos),
+            removed=list(removed),
+            removed_environment_naming_info=existing_environment.naming_info
+            if removed and existing_environment
+            else None,
         )
 
     @transactional()

tests/core/test_state_sync.py

@@ -473,13 +473,11 @@ def test_promote_snapshots(state_sync: EngineAdapterStateSync, make_snapshot: t.
         [snapshot_a, snapshot_b_old, snapshot_c],
         "prod",
     )
-    assert set(promotion_result.added) == set(
-        [
-            snapshot_a.table_info,
-            snapshot_b_old.table_info,
-            snapshot_c.table_info,
-        ]
-    )
+    assert set(promotion_result.added) == {
+        snapshot_a.table_info,
+        snapshot_b_old.table_info,
+        snapshot_c.table_info,
+    }
     assert not promotion_result.removed
     assert not promotion_result.removed_environment_naming_info
 
@@ -550,7 +548,7 @@ def test_promote_snapshots_suffix_change(
         environment_suffix_target=EnvironmentSuffixTarget.TABLE,
     )
 
-    assert set(promotion_result.added) == set([snapshot_a.table_info, snapshot_b.table_info])
+    assert set(promotion_result.added) == {snapshot_a.table_info, snapshot_b.table_info}
     assert not promotion_result.removed
     assert not promotion_result.removed_environment_naming_info
 
@@ -572,76 +570,156 @@ def test_promote_snapshots_suffix_change(
     )
 
     # We still only add the snapshots that are included in the promotion
-    assert set(promotion_result.added) == set([snapshot_b.table_info, snapshot_c.table_info])
-    # We also remove b because of the suffix target change. The new one will be created in the new suffix target
-    assert set(promotion_result.removed) == set([snapshot_a.table_info, snapshot_b.table_info])
+    assert set(promotion_result.added) == {snapshot_b.table_info, snapshot_c.table_info}
+    # B does not get removed because the suffix target change doesn't affect it due to running in prod.
+    assert set(promotion_result.removed) == {snapshot_a.table_info}
     # Make sure the removed suffix target is correctly seen as table
-    assert promotion_result.removed_environment_naming_info
+    assert promotion_result.removed_environment_naming_info is not None
     assert promotion_result.removed_environment_naming_info.suffix_target.is_table
 
+    promotion_result = promote_snapshots(
+        state_sync,
+        [snapshot_b, snapshot_c],
+        "dev",
+        environment_suffix_target=EnvironmentSuffixTarget.SCHEMA,
+    )
+
+    # We still only add the snapshots that are included in the promotion
+    assert set(promotion_result.added) == {snapshot_b.table_info, snapshot_c.table_info}
+    assert len(promotion_result.removed) == 0
+    assert promotion_result.removed_environment_naming_info is None
+
+    promotion_result = promote_snapshots(
+        state_sync,
+        [snapshot_b, snapshot_c],
+        "dev",
+        environment_suffix_target=EnvironmentSuffixTarget.TABLE,
+    )
+
+    # All snapshots are promoted due to suffix target change
+    assert set(promotion_result.added) == {
+        snapshot_b.table_info,
+        snapshot_c.table_info,
+    }
+    # All snapshots are removed due to suffix target change
+    assert set(promotion_result.removed) == {
+        snapshot_b.table_info,
+        snapshot_c.table_info,
+    }
+    assert promotion_result.removed_environment_naming_info is not None
+    assert promotion_result.removed_environment_naming_info.suffix_target.is_schema
+
 
 def test_promote_snapshots_catalog_name_override_change(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
 ):
     snapshot_a = make_snapshot(
         SqlModel(
-            name="a",
+            name="catalog1.schema.a",
             query=parse_one("select 1, ds"),
         ),
     )
     snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
 
     snapshot_b = make_snapshot(
         SqlModel(
-            name="b",
+            name="catalog1.schema.b",
             kind=FullKind(),
             query=parse_one("select * from a"),
         ),
         nodes={"a": snapshot_a.model},
     )
     snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
 
-    state_sync.push_snapshots([snapshot_a, snapshot_b])
+    snapshot_c = make_snapshot(
+        SqlModel(
+            name="catalog2.schema.c",
+            kind=FullKind(),
+            query=parse_one("select * from a"),
+        ),
+        nodes={"a": snapshot_a.model},
+    )
+    snapshot_c.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    state_sync.push_snapshots([snapshot_a, snapshot_b, snapshot_c])
 
     promotion_result = promote_snapshots(
         state_sync,
-        [snapshot_a, snapshot_b],
+        [snapshot_a, snapshot_b, snapshot_c],
         "prod",
-        environment_suffix_target=EnvironmentSuffixTarget.TABLE,
         environment_catalog_mapping={},
     )
 
-    assert set(promotion_result.added) == set([snapshot_a.table_info, snapshot_b.table_info])
+    assert set(promotion_result.added) == {
+        snapshot_a.table_info,
+        snapshot_b.table_info,
+        snapshot_c.table_info,
+    }
     assert not promotion_result.removed
     assert not promotion_result.removed_environment_naming_info
 
-    snapshot_c = make_snapshot(
+    snapshot_d = make_snapshot(
         SqlModel(
-            name="c",
+            name="catalog1.schema.d",
             query=parse_one("select 3, ds"),
         ),
     )
-    snapshot_c.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_d.categorize_as(SnapshotChangeCategory.BREAKING)
 
-    state_sync.push_snapshots([snapshot_c])
+    state_sync.push_snapshots([snapshot_d])
 
     promotion_result = promote_snapshots(
         state_sync,
-        [snapshot_b, snapshot_c],
+        [snapshot_b, snapshot_c, snapshot_d],
         "prod",
         environment_catalog_mapping={
-            re.compile("^prod$"): "prod_catalog",
+            re.compile("^prod$"): "catalog1",
         },
     )
 
-    # We still only add the snapshots that are included in the promotion
-    assert set(promotion_result.added) == set([snapshot_b.table_info, snapshot_c.table_info])
-    # We also remove b because of the catalog change. The new one will be created in the new catalog
-    assert set(promotion_result.removed) == set([snapshot_a.table_info, snapshot_b.table_info])
+    # We still only add the snapshots that are included in the promotion which means removing A
+    assert set(promotion_result.added) == {
+        snapshot_b.table_info,
+        snapshot_c.table_info,
+        snapshot_d.table_info,
+    }
+    # C is removed because of the catalog change. The new one will be created in the new catalog.
+    # B is not removed because it's catalog did not change and therefore removing would actually result
+    # in dropping what we just added.
+    # A is removed because it was explicitly removed from the promotion.
+    assert set(promotion_result.removed) == {snapshot_a.table_info, snapshot_c.table_info}
     # Make sure the removed suffix target correctly has the old catalog name set
     assert promotion_result.removed_environment_naming_info
     assert promotion_result.removed_environment_naming_info.catalog_name_override is None
 
+    promotion_result = promote_snapshots(
+        state_sync,
+        [snapshot_b, snapshot_c, snapshot_d],
+        "prod",
+        environment_catalog_mapping={
+            re.compile("^prod$"): "catalog2",
+        },
+    )
+
+    # All are added since their catalog was changed
+    assert set(promotion_result.added) == {
+        snapshot_b.table_info,
+        snapshot_c.table_info,
+        snapshot_d.table_info,
+    }
+    # All are removed since there were moved from their old catalog location
+    # Note that C has a catalog set in the model definition of `catalog2` which is what we moved to so you might think
+    # it shouldn't be removed, but its actual catalog was `catalog1` because of the previous override so therefore
+    # it should be removed from `catalog1`.
+    assert set(promotion_result.removed) == {
+        snapshot_b.table_info,
+        snapshot_c.table_info,
+        snapshot_d.table_info,
+    }
+    # Make sure the removed suffix target correctly has the old catalog name set
+    assert promotion_result.removed_environment_naming_info
+    assert promotion_result.removed_environment_naming_info.catalog_name_override == "catalog1"
+
 
 def test_promote_snapshots_parent_plan_id_mismatch(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable