refactor(pydantic): remove Pydantic v1 legacy code and migrate fully to v2

metadata-ingestion/setup.cfg

@@ -1,8 +1,7 @@
 [mypy]
 plugins =
     ./tests/test_helpers/sqlalchemy_mypy_plugin.py,
-    pydantic.mypy,
-    pydantic.v1.mypy
+    pydantic.mypy
 exclude = ^(venv/|build/|dist/|examples/transforms/setup.py)
 ignore_missing_imports = yes
 namespace_packages = no
@@ -75,8 +74,6 @@ filterwarnings =
     ignore:Deprecated call to \`pkg_resources.declare_namespace:DeprecationWarning
     ignore:pkg_resources is deprecated as an API:DeprecationWarning
     ignore:Did not recognize type:sqlalchemy.exc.SAWarning
-    # TODO: We should remove this and start fixing the deprecations.
-    ignore::pydantic.warnings.PydanticDeprecatedSince20
     ignore::datahub.configuration.common.ConfigurationWarning
     ignore:The new datahub SDK:datahub.errors.ExperimentalWarning
     # We should not be unexpectedly seeing API tracing warnings.

metadata-ingestion/src/datahub/api/entities/assertion/assertion.py

@@ -1,12 +1,15 @@
 from abc import abstractmethod
 from typing import Optional
 
+from pydantic import BaseModel, Field
+
 from datahub.api.entities.assertion.assertion_trigger import AssertionTrigger
-from datahub.configuration.pydantic_migration_helpers import v1_ConfigModel, v1_Field
 from datahub.metadata.com.linkedin.pegasus2avro.assertion import AssertionInfo
 
 
-class BaseAssertionProtocol(v1_ConfigModel):
+class BaseAssertionProtocol(BaseModel):
+    model_config = {"extra": "forbid"}
+
     @abstractmethod
     def get_id(self) -> str:
         pass
@@ -24,15 +27,17 @@ def get_assertion_trigger(
         pass
 
 
-class BaseAssertion(v1_ConfigModel):
-    id_raw: Optional[str] = v1_Field(
+class BaseAssertion(BaseModel):
+    model_config = {"extra": "forbid"}
+
+    id_raw: Optional[str] = Field(
         default=None,
         description="The raw id of the assertion."
         "If provided, this is used when creating identifier for this assertion"
         "along with assertion type and entity.",
     )
 
-    id: Optional[str] = v1_Field(
+    id: Optional[str] = Field(
         default=None,
         description="The id of the assertion."
         "If provided, this is used as identifier for this assertion."
@@ -41,17 +46,14 @@ class BaseAssertion(v1_ConfigModel):
 
     description: Optional[str] = None
 
-    # Can contain metadata extracted from datahub. e.g.
-    # - entity qualified name
-    # - entity schema
     meta: Optional[dict] = None
 
 
 class BaseEntityAssertion(BaseAssertion):
-    entity: str = v1_Field(
+    entity: str = Field(
         description="The entity urn that the assertion is associated with"
     )
 
-    trigger: Optional[AssertionTrigger] = v1_Field(
+    trigger: Optional[AssertionTrigger] = Field(
         default=None, description="The trigger schedule for assertion", alias="schedule"
     )

metadata-ingestion/src/datahub/api/entities/assertion/assertion_config_spec.py

@@ -1,13 +1,13 @@
 from typing import List, Optional
 
+from pydantic import BaseModel, Field
 from ruamel.yaml import YAML
 from typing_extensions import Literal
 
 from datahub.api.entities.assertion.datahub_assertion import DataHubAssertion
-from datahub.configuration.pydantic_migration_helpers import v1_ConfigModel, v1_Field
 
 
-class AssertionsConfigSpec(v1_ConfigModel):
+class AssertionsConfigSpec(BaseModel):
     """
     Declarative configuration specification for datahub assertions.
 
@@ -18,9 +18,11 @@ class AssertionsConfigSpec(v1_ConfigModel):
     In future, this would invoke datahub GraphQL API to upsert assertions.
     """
 
+    model_config = {"extra": "forbid"}
+
     version: Literal[1]
 
-    id: Optional[str] = v1_Field(
+    id: Optional[str] = Field(
         default=None,
         alias="namespace",
         description="Unique identifier of assertions configuration file",
@@ -34,8 +36,7 @@ def from_yaml(
         file: str,
     ) -> "AssertionsConfigSpec":
         with open(file) as fp:
-            yaml = YAML(typ="rt")  # default, if not specfied, is 'rt' (round-trip)
+            yaml = YAML(typ="rt")
             orig_dictionary = yaml.load(fp)
-            parsed_spec = AssertionsConfigSpec.parse_obj(orig_dictionary)
-            # parsed_spec._original_yaml_dict = orig_dictionary
+            parsed_spec = AssertionsConfigSpec.model_validate(orig_dictionary)
             return parsed_spec

metadata-ingestion/src/datahub/api/entities/assertion/assertion_operator.py

@@ -1,9 +1,9 @@
 import json
 from typing import List, Optional, Union
 
+from pydantic import BaseModel
 from typing_extensions import Literal, Protocol
 
-from datahub.configuration.pydantic_migration_helpers import v1_ConfigModel
 from datahub.metadata.schema_classes import (
     AssertionStdOperatorClass,
     AssertionStdParameterClass,
@@ -61,7 +61,9 @@ def _generate_assertion_std_parameters(
     )
 
 
-class EqualToOperator(v1_ConfigModel):
+class EqualToOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["equal_to"]
     value: Union[str, int, float]
 
@@ -74,7 +76,8 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class NotEqualToOperator(v1_ConfigModel):
+class NotEqualToOperator(BaseModel):
+    model_config = {"extra": "forbid"}
     type: Literal["not_equal_to"]
     value: Union[str, int, float]
 
@@ -87,7 +90,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class BetweenOperator(v1_ConfigModel):
+class BetweenOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["between"]
     min: Union[int, float]
     max: Union[int, float]
@@ -103,7 +108,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         )
 
 
-class LessThanOperator(v1_ConfigModel):
+class LessThanOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["less_than"]
     value: Union[int, float]
 
@@ -116,7 +123,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class GreaterThanOperator(v1_ConfigModel):
+class GreaterThanOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["greater_than"]
     value: Union[int, float]
 
@@ -129,7 +138,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class LessThanOrEqualToOperator(v1_ConfigModel):
+class LessThanOrEqualToOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["less_than_or_equal_to"]
     value: Union[int, float]
 
@@ -142,7 +153,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class GreaterThanOrEqualToOperator(v1_ConfigModel):
+class GreaterThanOrEqualToOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["greater_than_or_equal_to"]
     value: Union[int, float]
 
@@ -155,7 +168,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class InOperator(v1_ConfigModel):
+class InOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["in"]
     value: List[Union[str, float, int]]
 
@@ -168,7 +183,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class NotInOperator(v1_ConfigModel):
+class NotInOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["not_in"]
     value: List[Union[str, float, int]]
 
@@ -181,7 +198,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class IsNullOperator(v1_ConfigModel):
+class IsNullOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["is_null"]
 
     operator: str = AssertionStdOperatorClass.NULL
@@ -193,7 +212,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters()
 
 
-class NotNullOperator(v1_ConfigModel):
+class NotNullOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["is_not_null"]
 
     operator: str = AssertionStdOperatorClass.NOT_NULL
@@ -205,7 +226,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters()
 
 
-class IsTrueOperator(v1_ConfigModel):
+class IsTrueOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["is_true"]
 
     operator: str = AssertionStdOperatorClass.IS_TRUE
@@ -217,7 +240,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters()
 
 
-class IsFalseOperator(v1_ConfigModel):
+class IsFalseOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["is_false"]
 
     operator: str = AssertionStdOperatorClass.IS_FALSE
@@ -229,7 +254,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters()
 
 
-class ContainsOperator(v1_ConfigModel):
+class ContainsOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["contains"]
     value: str
 
@@ -242,7 +269,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class EndsWithOperator(v1_ConfigModel):
+class EndsWithOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["ends_with"]
     value: str
 
@@ -255,7 +284,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class StartsWithOperator(v1_ConfigModel):
+class StartsWithOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["starts_with"]
     value: str
 
@@ -268,7 +299,9 @@ def generate_parameters(self) -> AssertionStdParametersClass:
         return _generate_assertion_std_parameters(value=self.value)
 
 
-class MatchesRegexOperator(v1_ConfigModel):
+class MatchesRegexOperator(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["matches_regex"]
     value: str
 

metadata-ingestion/src/datahub/api/entities/assertion/assertion_trigger.py

@@ -2,51 +2,57 @@
 from typing import Union
 
 import humanfriendly
+from pydantic import BaseModel, Field, RootModel, field_validator
 from typing_extensions import Literal
 
-from datahub.configuration.pydantic_migration_helpers import (
-    v1_ConfigModel,
-    v1_Field,
-    v1_validator,
-)
 
+class CronTrigger(BaseModel):
+    model_config = {"extra": "forbid"}
 
-class CronTrigger(v1_ConfigModel):
     type: Literal["cron"]
-    cron: str = v1_Field(
+    cron: str = Field(
         description="The cron expression to use. See https://crontab.guru/ for help."
     )
-    timezone: str = v1_Field(
+    timezone: str = Field(
         "UTC",
         description="The timezone to use for the cron schedule. Defaults to UTC.",
     )
 
 
-class IntervalTrigger(v1_ConfigModel):
+class IntervalTrigger(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["interval"]
     interval: timedelta
 
-    @v1_validator("interval", pre=True)
+    @field_validator("interval", mode="before")
+    @classmethod
     def lookback_interval_to_timedelta(cls, v):
         if isinstance(v, str):
             seconds = humanfriendly.parse_timespan(v)
             return timedelta(seconds=seconds)
         raise ValueError("Invalid value.")
 
 
-class EntityChangeTrigger(v1_ConfigModel):
+class EntityChangeTrigger(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["on_table_change"]
 
 
-class ManualTrigger(v1_ConfigModel):
+class ManualTrigger(BaseModel):
+    model_config = {"extra": "forbid"}
+
     type: Literal["manual"]
 
 
-class AssertionTrigger(v1_ConfigModel):
-    __root__: Union[
-        CronTrigger, IntervalTrigger, EntityChangeTrigger, ManualTrigger
-    ] = v1_Field(discriminator="type")
+class AssertionTrigger(
+    RootModel[Union[CronTrigger, IntervalTrigger, EntityChangeTrigger, ManualTrigger]]
+):
+    root: Union[CronTrigger, IntervalTrigger, EntityChangeTrigger, ManualTrigger] = (
+        Field(discriminator="type")
+    )
 
     @property
     def trigger(self):
-        return self.__root__
+        return self.root