Merge pull request #1426 from elementary-data/ELE-1696-support-pydant…

…ic-2

support pydantic 2 in elementary

elementary/clients/dbt/slim_dbt_runner.py

@@ -17,10 +17,10 @@
 from dbt.parser.manifest import Manifest, ManifestLoader
 from dbt.tracking import disable_tracking
 from dbt.version import __version__ as dbt_version_string
-from pydantic import BaseModel, validator
 
 from elementary.clients.dbt.base_dbt_runner import BaseDbtRunner
 from elementary.utils.log import get_logger
+from elementary.utils.pydantic_shim import BaseModel, validator
 
 logger = get_logger(__name__)
 

elementary/clients/slack/schema.py

@@ -1,8 +1,7 @@
 from typing import List, Optional
 
-from pydantic import BaseModel, validator
-
 from elementary.utils.log import get_logger
+from elementary.utils.pydantic_shim import BaseModel, validator
 
 logger = get_logger(__name__)
 

elementary/monitor/api/filters/schema.py

@@ -1,6 +1,6 @@
 from typing import List, Optional
 
-from pydantic import BaseModel
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class FilterSchema(BaseModel):

elementary/monitor/api/groups/schema.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class GroupItemSchema(BaseModel):

elementary/monitor/api/lineage/schema.py

@@ -1,9 +1,9 @@
 import re
-from typing import List, Optional, Tuple
+from typing import List, Literal, Optional, Tuple
 
 import networkx as nx
-from pydantic import BaseModel, validator
-from pydantic.typing import Literal
+
+from elementary.utils.pydantic_shim import BaseModel, validator
 
 NodeUniqueIdType = str
 NodeType = Literal["model", "source", "exposure"]

elementary/monitor/api/models/schema.py

@@ -2,14 +2,13 @@
 import posixpath
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, Field, validator
-
 from elementary.monitor.api.totals_schema import TotalsSchema
 from elementary.monitor.fetchers.models.schema import (
     ExposureSchema,
     ModelSchema,
     SourceSchema,
 )
+from elementary.utils.pydantic_shim import BaseModel, Field, validator
 from elementary.utils.schema import ExtendedBaseModel
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 
@@ -38,17 +37,17 @@ def format_normalized_full_path_sep(cls, normalized_full_path: str) -> str:
 
 # NormalizedArtifactSchema must be first in the inheritance order
 class NormalizedModelSchema(NormalizedArtifactSchema, ModelSchema):
-    artifact_type: str = Field("model", const=True)
+    artifact_type: str = Field("model", const=True)  # type: ignore  # noqa
 
 
 # NormalizedArtifactSchema must be first in the inheritance order
 class NormalizedSourceSchema(NormalizedArtifactSchema, SourceSchema):
-    artifact_type: str = Field("source", const=True)
+    artifact_type: str = Field("source", const=True)  # type: ignore  # noqa
 
 
 # NormalizedArtifactSchema must be first in the inheritance order
 class NormalizedExposureSchema(NormalizedArtifactSchema, ExposureSchema):
-    artifact_type: str = Field("exposure", const=True)
+    artifact_type: str = Field("exposure", const=True)  # type: ignore  # noqa
 
 
 class ModelCoverageSchema(BaseModel):

elementary/monitor/api/report/schema.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import BaseModel
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class ReportDataEnvSchema(BaseModel):

elementary/monitor/api/selector/schema.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from pydantic import BaseModel
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class SelectorSchema(BaseModel):

elementary/monitor/api/source_freshnesses/schema.py

@@ -1,8 +1,7 @@
 from typing import List, Optional
 
-from pydantic import BaseModel, Field, validator
-
 from elementary.monitor.api.totals_schema import TotalsSchema
+from elementary.utils.pydantic_shim import BaseModel, Field, validator
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 
 

elementary/monitor/api/tests/schema.py

@@ -1,8 +1,7 @@
 from typing import List, Optional, Union
 
-from pydantic import BaseModel, Field, validator
-
 from elementary.monitor.api.totals_schema import TotalsSchema
+from elementary.utils.pydantic_shim import BaseModel, Field, validator
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 
 

elementary/monitor/api/totals_schema.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class TotalsSchema(BaseModel):

elementary/monitor/data_monitoring/alerts/integrations/utils/report_link.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Optional
 
-from pydantic import BaseModel
+from elementary.utils.pydantic_shim import BaseModel
 
 TEST_RUNS_LINK_TEXT = "View test runs"
 MODEL_RUNS_LINK_TEXT = "View model runs"

elementary/monitor/data_monitoring/alerts/schema.py

@@ -1,8 +1,7 @@
 from typing import List
 
-from pydantic import BaseModel
-
 from elementary.monitor.fetchers.alerts.schema.pending_alerts import PendingAlertSchema
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class SortedAlertsSchema(BaseModel):

elementary/monitor/data_monitoring/schema.py

@@ -3,9 +3,8 @@
 from enum import Enum
 from typing import Any, List, Optional, Pattern, Tuple
 
-from pydantic import BaseModel, Field, validator
-
 from elementary.utils.log import get_logger
+from elementary.utils.pydantic_shim import BaseModel, Field, validator
 from elementary.utils.time import DATETIME_FORMAT, convert_local_time_to_timezone
 
 logger = get_logger(__name__)

elementary/monitor/fetchers/alerts/schema/alert_data.py

@@ -1,8 +1,6 @@
 from datetime import datetime
 from typing import Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field, validator
-
 from elementary.monitor.alerts.model_alert import ModelAlertModel
 from elementary.monitor.alerts.source_freshness_alert import SourceFreshnessAlertModel
 from elementary.monitor.alerts.test_alert import TestAlertModel
@@ -13,6 +11,7 @@
     unpack_and_flatten_and_dedup_list_of_strings,
     unpack_and_flatten_str_to_list,
 )
+from elementary.utils.pydantic_shim import BaseModel, Field, validator
 
 ALERTS_CONFIG_KEY = "alerts_config"
 CHANNEL_KEY = "channel"
@@ -138,7 +137,7 @@ class TestAlertDataSchema(BaseAlertDataSchema):
     severity: str
     test_meta: Optional[Dict] = None
     elementary_unique_id: str
-    resource_type: ResourceType = Field(ResourceType.TEST, const=True)
+    resource_type: ResourceType = Field(ResourceType.TEST, const=True)  # type: ignore  # noqa
 
     @property
     def flatten_model_meta(self) -> Dict:
@@ -230,7 +229,7 @@ class ModelAlertDataSchema(BaseAlertDataSchema):
     materialization: str
     full_refresh: bool
     message: Optional[str] = None
-    resource_type: ResourceType = Field(ResourceType.MODEL, const=True)
+    resource_type: ResourceType = Field(ResourceType.MODEL, const=True)  # type: ignore  # noqa
 
     def format_alert(
         self,
@@ -282,7 +281,7 @@ class SourceFreshnessAlertDataSchema(BaseAlertDataSchema):
     path: str
     error: Optional[str] = None
     freshness_description: Optional[str] = None
-    resource_type: ResourceType = Field(ResourceType.SOURCE_FRESHNESS, const=True)
+    resource_type: ResourceType = Field(ResourceType.SOURCE_FRESHNESS, const=True)  # type: ignore  # noqa
 
     def format_alert(
         self,

elementary/monitor/fetchers/alerts/schema/pending_alerts.py

@@ -2,14 +2,13 @@
 from enum import Enum
 from typing import Optional, Union
 
-from pydantic import BaseModel, root_validator
-
 from elementary.monitor.fetchers.alerts.schema.alert_data import (
     ModelAlertDataSchema,
     SourceFreshnessAlertDataSchema,
     TestAlertDataSchema,
 )
 from elementary.utils.json_utils import try_load_json
+from elementary.utils.pydantic_shim import BaseModel, root_validator
 
 ALERTS_CONFIG_KEY = "alerts_config"
 CHANNEL_KEY = "channel"

elementary/monitor/fetchers/invocations/schema.py

@@ -1,8 +1,7 @@
 from typing import Optional
 
-from pydantic import BaseModel, validator
-
 from elementary.utils.json_utils import try_load_json
+from elementary.utils.pydantic_shim import BaseModel, validator
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 
 

elementary/monitor/fetchers/models/schema.py

@@ -2,8 +2,7 @@
 import posixpath
 from typing import Any, Dict, List, Optional, TypeVar
 
-from pydantic import Field, validator
-
+from elementary.utils.pydantic_shim import Field, validator
 from elementary.utils.schema import ExtendedBaseModel
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 

elementary/monitor/fetchers/source_freshnesses/schema.py

@@ -1,7 +1,6 @@
 from typing import List, Optional
 
-from pydantic import validator
-
+from elementary.utils.pydantic_shim import validator
 from elementary.utils.schema import ExtendedBaseModel
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 

elementary/monitor/fetchers/test_management/schema.py

@@ -1,7 +1,6 @@
 from typing import List, Literal, Optional
 
-from pydantic import BaseModel, Field, validator
-
+from elementary.utils.pydantic_shim import BaseModel, Field, validator
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 
 

elementary/monitor/fetchers/tests/schema.py

@@ -1,7 +1,6 @@
 from typing import List, Optional, Union
 
-from pydantic import Field, validator
-
+from elementary.utils.pydantic_shim import Field, validator
 from elementary.utils.schema import ExtendedBaseModel
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
 

elementary/utils/pydantic_shim.py

@@ -0,0 +1,32 @@
+"""
+copied from: https://github.com/dbt-labs/dbt-semantic-interfaces/blob/main/dsi_pydantic_shim.py
+"""
+
+from importlib.metadata import version
+
+pydantic_version = version("pydantic")
+# Pydantic uses semantic versioning, i.e. <major>.<minor>.<patch>, and we need to know the major
+pydantic_major = pydantic_version.split(".")[0]
+
+if pydantic_major == "1":
+    from pydantic import (  # type: ignore  # noqa
+        BaseModel,
+        Extra,
+        Field,
+        create_model,
+        root_validator,
+        validator,
+    )
+elif pydantic_major == "2":
+    from pydantic.v1 import (  # type: ignore  # noqa
+        BaseModel,
+        Extra,
+        Field,
+        create_model,
+        root_validator,
+        validator,
+    )
+else:
+    raise RuntimeError(
+        f"Currently only pydantic 1 and 2 are supported, found pydantic {pydantic_version}"
+    )

elementary/utils/schema.py

@@ -1,9 +1,8 @@
 import json
 from typing import Union
 
-from pydantic import BaseModel
-
 from elementary.utils.json_utils import try_load_json
+from elementary.utils.pydantic_shim import BaseModel
 
 
 class ExtendedBaseModel(BaseModel):

mypy.ini

@@ -14,3 +14,6 @@ ignore_missing_imports = True
 
 [mypy-posthog]
 ignore_missing_imports = True
+
+[mypy-networkx]
+ignore_missing_imports = True

pyproject.toml

@@ -27,7 +27,8 @@ google-cloud-storage = "<3.0.0"
 alive-progress = "<=2.3.1"
 slack-sdk = ">=3.20.1,<4.0.0"
 pytest-parametrization = ">=2022.2.1"
-pydantic = "<2.0"
+
+pydantic = "<3.0"
 networkx = ">=2.3,<3"
 packaging = ">=20.9,<=23.1"
 azure-storage-blob = ">=12.11.0"