fix(ingest): update pydantic helpers to address unique name issue (#8324)

Co-authored-by: Aseem Bansal <asmbansal2@gmail.com>

metadata-ingestion/src/datahub/configuration/pydantic_field_deprecation.py

@@ -19,4 +19,8 @@ def pydantic_field_deprecated(field: str, message: Optional[str] = None) -> clas
             warnings.warn(output, ConfigurationWarning, stacklevel=2)
         return values
 
+    # Hack: Pydantic maintains unique list of validators by referring its __name__.
+    # https://github.com/pydantic/pydantic/blob/v1.10.9/pydantic/main.py#L264
+    # This hack ensures that multiple field deprecated do not overwrite each other.
+    _validate_deprecated.__name__ = f"{_validate_deprecated.__name__}_{field}"
     return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_deprecated)

metadata-ingestion/src/datahub/configuration/validate_field_removal.py

@@ -10,7 +10,7 @@ def pydantic_removed_field(
     field: str,
     print_warning: bool = True,
 ) -> classmethod:
-    def _validate_field_rename(cls: Type, values: dict) -> dict:
+    def _validate_field_removal(cls: Type, values: dict) -> dict:
         if field in values:
             if print_warning:
                 warnings.warn(
@@ -21,4 +21,8 @@ def pydantic_removed_field(
             values.pop(field)
         return values
 
-    return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_field_rename)
+    # Hack: Pydantic maintains unique list of validators by referring its __name__.
+    # https://github.com/pydantic/pydantic/blob/v1.10.9/pydantic/main.py#L264
+    # This hack ensures that multiple field removals do not overwrite each other.
+    _validate_field_removal.__name__ = f"{_validate_field_removal.__name__}_{field}"
+    return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_field_removal)

metadata-ingestion/tests/unit/test_pydantic_validators.py

@@ -1,15 +1,20 @@
+from typing import Optional
+
 import pytest
 from pydantic import ValidationError
 
 from datahub.configuration.common import ConfigModel
+from datahub.configuration.pydantic_field_deprecation import pydantic_field_deprecated
+from datahub.configuration.validate_field_removal import pydantic_removed_field
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
+from datahub.utilities.global_warning_util import get_global_warnings
 
 
 def test_field_rename():
     class TestModel(ConfigModel):
         b: str
 
-        _validate_deprecated = pydantic_renamed_field("a", "b")
+        _validate_rename = pydantic_renamed_field("a", "b")
 
     v = TestModel.parse_obj({"b": "original"})
     assert v.b == "original"
@@ -54,3 +59,37 @@ def test_field_multiple_fields_rename():
 
     with pytest.raises(ValidationError):
         TestModel.parse_obj({})
+
+
+def test_field_remove():
+    class TestModel(ConfigModel):
+        b: str
+
+        _validate_removed_r1 = pydantic_removed_field("r1")
+        _validate_removed_r2 = pydantic_removed_field("r2")
+
+    v = TestModel.parse_obj({"b": "original"})
+    assert v.b == "original"
+
+    v = TestModel.parse_obj({"b": "original", "r1": "removed", "r2": "removed"})
+    assert v.b == "original"
+
+
+def test_field_deprecated():
+    class TestModel(ConfigModel):
+        d1: Optional[str]
+        d2: Optional[str]
+        b: str
+
+        _validate_deprecated_d1 = pydantic_field_deprecated("d1")
+        _validate_deprecated_d2 = pydantic_field_deprecated("d2")
+
+    v = TestModel.parse_obj({"b": "original"})
+    assert v.b == "original"
+
+    v = TestModel.parse_obj({"b": "original", "d1": "deprecated", "d2": "deprecated"})
+    assert v.b == "original"
+    assert v.d1 == "deprecated"
+    assert v.d2 == "deprecated"
+    assert any(["d1 is deprecated" in warning for warning in get_global_warnings()])
+    assert any(["d2 is deprecated" in warning for warning in get_global_warnings()])