Repository: TobikoData/sqlmesh
Pull request #2035 (2024-01-28 02:12:26)

fix!: ignore private attributes doing equality pydantic
We have an issue that Pydantic includes private attributes in the equality check:
https://github.com/pydantic/pydantic/blob/b785d5b1db4eccb922dd2bf17faea5fc842e6a57/pydantic/main.py#L880

We use these for caching values. When we do equality checks our intent is to check if the public attributes of the class are the same and not if the cached attributes are equal. 

This PR has a way of solving this problem but involves mutating the actual objects themselves so it would be best if we had a solution that didn't involve that. Are there Pydantic config properties that can help us here? @izeigerman do you have any ideas? I want to find a solution that works for all Pydantic classes instead requiring each one to provide an override if they introduce private attributes. 


Patch:
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/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/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)