Repository: TobikoData/sqlmesh
Pull request #2048 (2024-01-30 00:29:26)

Feat: more configuration for format command
Not enough configurability. Expose very nice implementation in Generator to the benefit of all users while preserving existing behavior by default.

eakmanrq: @z3z1ma Can you also add these to the magic command?

https://github.com/TobikoData/sqlmesh/blob/main/sqlmesh/magics.py#L627 



Patch:
diff --git a/sqlmesh/cli/__init__.py b/sqlmesh/cli/__init__.py
index 0167b218..fd9c4e9e 100644
--- a/sqlmesh/cli/__init__.py
+++ b/sqlmesh/cli/__init__.py
@@ -5,6 +5,7 @@ from functools import wraps
 import click
 from sqlglot.errors import SqlglotError
 
+from sqlmesh.core.context import Context
 from sqlmesh.utils import debug_mode_enabled
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.errors import SQLMeshError
@@ -19,15 +20,21 @@ def error_handler(
     func: t.Callable[..., DECORATOR_RETURN_TYPE]
 ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
     @wraps(func)
-    def wrapper(*args: t.Any, **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
+    def wrapper(*args: t.List[t.Any], **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
+        context_or_obj = args[0]
+        sqlmesh_context = (
+            context_or_obj.obj if isinstance(context_or_obj, click.Context) else context_or_obj
+        )
+        if not isinstance(sqlmesh_context, Context):
+            sqlmesh_context = None
         handler = _debug_exception_handler if debug_mode_enabled() else _default_exception_handler
-        return handler(lambda: func(*args, **kwargs))
+        return handler(sqlmesh_context, lambda: func(*args, **kwargs))
 
     return wrapper
 
 
 def _default_exception_handler(
-    func: t.Callable[[], DECORATOR_RETURN_TYPE]
+    context: t.Optional[Context], func: t.Callable[[], DECORATOR_RETURN_TYPE]
 ) -> DECORATOR_RETURN_TYPE:
     try:
         return func()
@@ -36,11 +43,19 @@ def _default_exception_handler(
         raise click.ClickException(f"Failed processing {ex.node}. {cause}")
     except (SQLMeshError, SqlglotError, ValueError) as ex:
         raise click.ClickException(str(ex))
+    finally:
+        if context:
+            context.close()
 
 
-def _debug_exception_handler(func: t.Callable[[], DECORATOR_RETURN_TYPE]) -> DECORATOR_RETURN_TYPE:
+def _debug_exception_handler(
+    context: t.Optional[Context], func: t.Callable[[], DECORATOR_RETURN_TYPE]
+) -> DECORATOR_RETURN_TYPE:
     try:
         return func()
     except Exception:
         logger.exception("Unhandled exception")
         raise
+    finally:
+        if context:
+            context.close()
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index e4716e3b..006134bc 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -201,15 +201,52 @@ def evaluate(
     help="Transpile project models to the specified dialect.",
 )
 @click.option(
-    "--new-line",
+    "--newline",
     is_flag=True,
     help="Include a new line at the end of each file.",
 )
+@click.option(
+    "--normalize",
+    is_flag=True,
+    help="Whether or not to normalize identifiers to lowercase.",
+    default=False,
+)
+@click.option(
+    "--pad",
+    type=int,
+    help="Determines the pad size in a formatted string.",
+    default=2,
+)
+@click.option(
+    "--indent",
+    type=int,
+    help="Determines the indentation size in a formatted string.",
+    default=2,
+)
+@click.option(
+    "--normalize-functions",
+    type=str,
+    help="Whether or not to normalize all function names. Possible values are: 'upper', 'lower'",
+)
+@click.option(
+    "--leading-comma",
+    is_flag=True,
+    help="Determines whether or not the comma is leading or trailing in select expressions. Default is trailing.",
+    default=False,
+)
+@click.option(
+    "--max-text-width",
+    type=int,
+    help="The max number of characters in a segment before creating new lines in pretty mode.",
+    default=80,
+)
 @click.pass_context
 @error_handler
-def format(ctx: click.Context, transpile: t.Optional[str] = None, new_line: bool = False) -> None:
+def format(
+    ctx: click.Context, transpile: t.Optional[str] = None, newline: bool = False, **kwargs: t.Any
+) -> None:
     """Format all SQL models."""
-    ctx.obj.format(transpile, new_line)
+    ctx.obj.format(transpile, newline, **{k: v for k, v in kwargs.items() if v is not None})
 
 
 @cli.command("diff")
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 466fc49e..10952e4a 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -740,7 +740,9 @@ class Context(BaseContext):
 
         return df
 
-    def format(self, transpile: t.Optional[str] = None, newline: bool = False) -> None:
+    def format(
+        self, transpile: t.Optional[str] = None, newline: bool = False, **kwargs: t.Any
+    ) -> None:
         """Format all SQL models."""
         for model in self._models.values():
             if not model._path.suffix == ".sql":
@@ -759,7 +761,9 @@ class Context(BaseContext):
                                 )
                             )
                 file.seek(0)
-                file.write(format_model_expressions(expressions, transpile or model.dialect))
+                file.write(
+                    format_model_expressions(expressions, transpile or model.dialect, **kwargs)
+                )
                 if newline:
                     file.write("\n")
                 file.truncate()
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index b08a4c0e..8d15a409 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -132,7 +132,6 @@ def _parse_id_var(
     any_token: bool = True,
     tokens: t.Optional[t.Collection[TokenType]] = None,
 ) -> t.Optional[exp.Expression]:
-
     if self._prev and self._prev.text == SQLMESH_MACRO_PREFIX and self._match(TokenType.L_BRACE):
         identifier = self.__parse_id_var(any_token=any_token, tokens=tokens)  # type: ignore
         if not self._match(TokenType.R_BRACE):
@@ -505,13 +504,14 @@ def _override(klass: t.Type[Tokenizer | Parser], func: t.Callable) -> None:
 
 
 def format_model_expressions(
-    expressions: t.List[exp.Expression], dialect: t.Optional[str] = None
+    expressions: t.List[exp.Expression], dialect: t.Optional[str] = None, **kwargs: t.Any
 ) -> str:
     """Format a model's expressions into a standardized format.
 
     Args:
         expressions: The model's expressions, must be at least model def + query.
         dialect: The dialect to render the expressions as.
+        **kwargs: Additional keyword arguments to pass to the sql generator.
 
     Returns:
         A string representing the formatted model.
@@ -538,8 +538,8 @@ def format_model_expressions(
 
     return ";\n\n".join(
         [
-            *(statement.sql(pretty=True, dialect=dialect) for statement in statements),
-            query.sql(pretty=True, dialect=dialect),
+            *(statement.sql(pretty=True, dialect=dialect, **kwargs) for statement in statements),
+            query.sql(pretty=True, dialect=dialect, **kwargs),
         ]
     ).strip()
 
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 4d1be08a..83c00254 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -5,7 +5,7 @@ import typing as t
 from collections import defaultdict
 from datetime import datetime, timedelta
 from enum import IntEnum
-from functools import lru_cache
+from functools import cached_property, lru_cache
 
 from pydantic import Field
 from sqlglot import exp
@@ -258,7 +258,7 @@ class SnapshotInfoMixin(ModelKindMixin):
     def is_new_version(self) -> bool:
         raise NotImplementedError
 
-    @property
+    @cached_property
     def fully_qualified_table(self) -> t.Optional[exp.Table]:
         raise NotImplementedError
 
@@ -351,8 +351,6 @@ class SnapshotTableInfo(PydanticModel, SnapshotInfoMixin, frozen=True):
     # This can be removed from this model once Pydantic 1 support is dropped (must remain in `Snapshot` though)
     base_table_name_override: t.Optional[str] = None
 
-    _fully_qualified_table: t.Optional[exp.Table] = None
-
     def __lt__(self, other: SnapshotTableInfo) -> bool:
         return self.name < other.name
 
@@ -368,11 +366,9 @@ class SnapshotTableInfo(PydanticModel, SnapshotInfoMixin, frozen=True):
     def physical_schema(self) -> str:
         return self.physical_schema_
 
-    @property
+    @cached_property
     def fully_qualified_table(self) -> exp.Table:
-        if not self._fully_qualified_table:
-            self._fully_qualified_table = exp.to_table(self.name)
-        return self._fully_qualified_table
+        return exp.to_table(self.name)
 
     @property
     def table_info(self) -> SnapshotTableInfo:
@@ -1030,7 +1026,7 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         """Is restatement disabled for the node"""
         return self.is_model and self.model.disable_restatement
 
-    @property
+    @cached_property
     def fully_qualified_table(self) -> t.Optional[exp.Table]:
         if not self.is_model:
             return None
diff --git a/sqlmesh/core/state_sync/common.py b/sqlmesh/core/state_sync/common.py
index 99e423d9..d2419e03 100644
--- a/sqlmesh/core/state_sync/common.py
+++ b/sqlmesh/core/state_sync/common.py
@@ -15,6 +15,7 @@ from sqlmesh.core.snapshot import (
     SnapshotIdLike,
     SnapshotInfoLike,
     SnapshotNameVersionLike,
+    SnapshotTableInfo,
     start_date,
 )
 from sqlmesh.core.state_sync.base import PromotionResult, StateSync
@@ -128,31 +129,29 @@ class CommonStateSyncMixin(StateSync):
 
         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(
@@ -160,45 +159,32 @@ class CommonStateSyncMixin(StateSync):
                     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()
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index a0a100ec..253511e3 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -618,16 +618,49 @@ class SQLMeshMagics(Magics):
         help="Transpile project models to the specified dialect.",
     )
     @argument(
-        "--new-line",
+        "--newline",
         action="store_true",
         help="The output dialect of the sql string.",
     )
+    @argument(
+        "--normalize",
+        action="store_true",
+        help="Whether or not to normalize identifiers to lowercase.",
+    )
+    @argument(
+        "--pad",
+        type=int,
+        help="Determines the pad size in a formatted string.",
+        default=2,
+    )
+    @argument(
+        "--indent",
+        type=int,
+        help="Determines the indentation size in a formatted string.",
+        default=2,
+    )
+    @argument(
+        "--normalize-functions",
+        type=str,
+        help="Whether or not to normalize all function names. Possible values are: 'upper', 'lower'",
+    )
+    @argument(
+        "--leading-comma",
+        action="store_true",
+        help="Determines whether or not the comma is leading or trailing in select expressions. Default is trailing.",
+    )
+    @argument(
+        "--max-text-width",
+        type=int,
+        help="The max number of characters in a segment before creating new lines in pretty mode.",
+        default=80,
+    )
     @line_magic
     @pass_sqlmesh_context
     def format(self, context: Context, line: str) -> None:
         """Format all SQL models."""
         args = parse_argstring(self.format, line)
-        context.format(args.transpile, args.new_line)
+        context.format(**vars(args))
 
     @magic_arguments()
     @argument("environment", type=str, help="The environment to diff local state against.")
diff --git a/sqlmesh/utils/pydantic.py b/sqlmesh/utils/pydantic.py
index 32681f47..ca7cbc33 100644
--- a/sqlmesh/utils/pydantic.py
+++ b/sqlmesh/utils/pydantic.py
@@ -23,7 +23,9 @@ if t.TYPE_CHECKING:
 
 T = t.TypeVar("T")
 DEFAULT_ARGS = {"exclude_none": True, "by_alias": True}
-PYDANTIC_MAJOR_VERSION = int(pydantic.__version__.split(".")[0])
+PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION = [int(p) for p in pydantic.__version__.split(".")][
+    :2
+]
 
 
 if PYDANTIC_MAJOR_VERSION >= 2:
@@ -105,6 +107,8 @@ class PydanticModel(pydantic.BaseModel):
             smart_union = True
             keep_untouched = (cached_property,)
 
+    _hash_func_mapping: t.ClassVar[t.Dict[t.Type[t.Any], t.Callable[[t.Any], int]]] = {}
+
     def dict(
         self,
         **kwargs: t.Any,
@@ -191,6 +195,27 @@ class PydanticModel(pydantic.BaseModel):
             if predicate(field_info)
         }
 
+    def __eq__(self, other: t.Any) -> bool:
+        if (PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION) < (2, 6):
+            if isinstance(other, pydantic.BaseModel):
+                return self.dict() == other.dict()
+            else:
+                return self.dict() == other
+        return super().__eq__(other)
+
+    def __hash__(self) -> int:
+        if (PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION) < (2, 6):
+            obj = {k: v for k, v in self.__dict__.items() if k in self.all_field_infos()}
+            return hash(self.__class__) + hash(tuple(obj.values()))
+
+        from pydantic._internal._model_construction import (  # type: ignore
+            make_hash_func,
+        )
+
+        if self.__class__ not in PydanticModel._hash_func_mapping:
+            PydanticModel._hash_func_mapping[self.__class__] = make_hash_func(self.__class__)
+        return PydanticModel._hash_func_mapping[self.__class__](self)
+
     def __str__(self) -> str:
         args = []
 


Test patch:
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index 85a77b48..a450a990 100644
--- a/tests/core/test_state_sync.py
+++ b/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,20 +570,52 @@ 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"),
         ),
     )
@@ -593,7 +623,7 @@ def test_promote_snapshots_catalog_name_override_change(
 
     snapshot_b = make_snapshot(
         SqlModel(
-            name="b",
+            name="catalog1.schema.b",
             kind=FullKind(),
             query=parse_one("select * from a"),
         ),
@@ -601,47 +631,95 @@ def test_promote_snapshots_catalog_name_override_change(
     )
     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
diff --git a/tests/utils/test_pydantic.py b/tests/utils/test_pydantic.py
index e0ddf36f..9a7278c3 100644
--- a/tests/utils/test_pydantic.py
+++ b/tests/utils/test_pydantic.py
@@ -1,3 +1,5 @@
+from functools import cached_property
+
 from sqlmesh.utils.date import TimeLike, to_date, to_datetime
 from sqlmesh.utils.pydantic import PYDANTIC_MAJOR_VERSION, PydanticModel
 
@@ -16,3 +18,44 @@ def test_datetime_date_serialization() -> None:
     else:
         assert deserialized_date.ds == target_ds
         assert deserialized_datetime.ds == "2022-01-01T00:00:00+00:00"
+
+
+def test_pydantic_2_equality() -> None:
+    class TestModel(PydanticModel):
+        name: str
+
+        @cached_property
+        def private(self) -> str:
+            return "should be ignored"
+
+    model_a = TestModel(name="a")
+    model_a_duplicate = TestModel(name="a")
+    assert model_a == model_a_duplicate
+    model_b = TestModel(name="b")
+    assert model_a != model_b
+
+
+def test_pydantic_2_hash() -> None:
+    class TestModel(PydanticModel):
+        name: str
+
+        @cached_property
+        def private(self) -> str:
+            return "should be ignored"
+
+    class TestModel2(PydanticModel):
+        name: str
+        field2: str = "test"
+
+        @cached_property
+        def private(self) -> str:
+            return "should be ignored"
+
+    model_a = TestModel(name="a")
+    model_a_duplicate = TestModel(name="a")
+    assert hash(model_a) == hash(model_a_duplicate)
+
+    model_2_a = TestModel2(name="a")
+    model_2_b = TestModel2(name="a")
+    assert hash(model_2_a) == hash(model_2_b)
+    assert hash(model_a) != hash(model_2_a)