Ele 1696 upgrade pydantic

elementary/clients/dbt/slim_dbt_runner.py

@@ -17,7 +17,7 @@
 from dbt.parser.manifest import ManifestLoader
 from dbt.tracking import disable_tracking
 from dbt.version import __version__ as dbt_version_string
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from elementary.clients.dbt.base_dbt_runner import BaseDbtRunner
 from elementary.utils.log import get_logger
@@ -65,7 +65,7 @@ class ConfigArgs(BaseModel):
     threads: Optional[int] = 1
     vars: Optional[Union[str, Dict[str, Any]]] = DEFAULT_VARS
 
-    @validator("vars", pre=True)
+    @field_validator("vars", mode="before")
     def validate_vars(cls, vars):
         if not vars:
             return DEFAULT_VARS

elementary/clients/slack/schema.py

@@ -1,6 +1,6 @@
 from typing import List, Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from elementary.utils.log import get_logger
 
@@ -17,7 +17,7 @@ class SlackMessageSchema(BaseModel):
     attachments: Optional[list] = None
     blocks: Optional[list] = None
 
-    @validator("attachments", pre=True)
+    @field_validator("attachments", mode="before")
     def validate_attachments(cls, attachments):
         if (
             isinstance(attachments, list)

elementary/monitor/api/filters/schema.py

@@ -1,9 +1,11 @@
 from typing import List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class FilterSchema(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+
     name: str
     display_name: str
     model_unique_ids: List[Optional[str]] = []
@@ -14,6 +16,8 @@ def add_model_unique_id(self, model_unique_id: Optional[str]):
 
 
 class FiltersSchema(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+
     test_results: List[FilterSchema] = list()
     test_runs: List[FilterSchema] = list()
     model_runs: List[FilterSchema] = list()

elementary/monitor/api/groups/schema.py

@@ -4,8 +4,8 @@
 
 
 class GroupItemSchema(BaseModel):
-    node_id: Optional[str]
-    resource_type: Optional[str]
+    node_id: Optional[str] = None
+    resource_type: Optional[str] = None
 
 
 DbtGroupSchema = Dict[str, dict]

elementary/monitor/api/lineage/schema.py

@@ -1,8 +1,7 @@
-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 pydantic import BaseModel, field_validator
 
 NodeUniqueIdType = str
 NodeType = Literal["model", "source", "exposure"]
@@ -17,11 +16,11 @@ class LineageSchema(BaseModel):
     nodes: Optional[List[LineageNodeSchema]] = None
     edges: Optional[List[Tuple[NodeUniqueIdType, NodeUniqueIdType]]] = None
 
-    @validator("nodes", pre=True, always=True)
+    @field_validator("nodes", mode="before")
     def set_nodes(cls, nodes):
         return nodes or []
 
-    @validator("edges", pre=True, always=True)
+    @field_validator("edges", mode="before")
     def set_edges(cls, edges):
         return edges or []
 

elementary/monitor/api/models/models.py

@@ -1,4 +1,3 @@
-import json
 import os
 import statistics
 from collections import defaultdict
@@ -203,7 +202,7 @@ def _normalize_dbt_artifact_dict(
             SourceSchema: NormalizedSourceSchema,
         }
         artifact_name = artifact.name
-        normalized_artifact = json.loads(artifact.json())
+        normalized_artifact = artifact.model_dump()
         normalized_artifact["model_name"] = artifact_name
         normalized_artifact["normalized_full_path"] = self._normalize_artifact_path(
             artifact

elementary/monitor/api/models/schema.py

@@ -1,8 +1,8 @@
 import os
 import posixpath
-from typing import Dict, List, Optional
+from typing import Dict, List, Literal, Optional
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from elementary.monitor.api.totals_schema import TotalsSchema
 from elementary.monitor.fetchers.models.schema import (
@@ -15,6 +15,8 @@
 
 
 class NormalizedArtifactSchema(ExtendedBaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+
     owners: Optional[List[str]] = []
     tags: Optional[List[str]] = []
     # Should be changed to artifact_name.
@@ -23,32 +25,32 @@ class NormalizedArtifactSchema(ExtendedBaseModel):
     normalized_full_path: str
     fqn: str
 
-    @validator("tags", pre=True)
+    @field_validator("tags", mode="before")
     def load_tags(cls, tags):
         return cls._load_var_to_list(tags)
 
-    @validator("owners", pre=True)
+    @field_validator("owners", mode="before")
     def load_owners(cls, owners):
         return cls._load_var_to_list(owners)
 
-    @validator("normalized_full_path", pre=True)
+    @field_validator("normalized_full_path", mode="before")
     def format_normalized_full_path_sep(cls, normalized_full_path: str) -> str:
         return posixpath.sep.join(normalized_full_path.split(os.path.sep))
 
 
 # NormalizedArtifactSchema must be first in the inheritance order
 class NormalizedModelSchema(NormalizedArtifactSchema, ModelSchema):
-    artifact_type: str = Field("model", const=True)
+    artifact_type: Literal["model"] = "model"
 
 
 # NormalizedArtifactSchema must be first in the inheritance order
 class NormalizedSourceSchema(NormalizedArtifactSchema, SourceSchema):
-    artifact_type: str = Field("source", const=True)
+    artifact_type: Literal["source"] = "source"
 
 
 # NormalizedArtifactSchema must be first in the inheritance order
 class NormalizedExposureSchema(NormalizedArtifactSchema, ExposureSchema):
-    artifact_type: str = Field("exposure", const=True)
+    artifact_type: Literal["exposure"] = "exposure"
 
 
 class ModelCoverageSchema(BaseModel):
@@ -60,11 +62,11 @@ class ModelRunSchema(BaseModel):
     id: str
     time_utc: str
     status: str
-    full_refresh: Optional[bool]
-    materialization: Optional[str]
+    full_refresh: Optional[bool] = None
+    materialization: Optional[str] = None
     execution_time: float
 
-    @validator("time_utc", pre=True)
+    @field_validator("time_utc", mode="before")
     def format_time_utc(cls, time_utc):
         return convert_partial_iso_format_to_full_iso_format(time_utc)
 

elementary/monitor/api/report/report.py

@@ -71,10 +71,10 @@ def get_report_data(
                 test_results.totals, test_runs.totals, models, sources, models_runs.runs
             )
 
-            serializable_groups = groups.dict()
+            serializable_groups = groups.model_dump()
             serializable_models = self._serialize_models(models, sources, exposures)
             serializable_model_runs = self._serialize_models_runs(models_runs.runs)
-            serializable_model_runs_totals = models_runs.dict(include={"totals"})[
+            serializable_model_runs_totals = models_runs.model_dump(include={"totals"})[
                 "totals"
             ]
             serializable_models_coverages = self._serialize_coverages(coverages)
@@ -86,9 +86,9 @@ def get_report_data(
             )
             serializable_test_runs = self._serialize_test_runs(test_runs.runs)
             serializable_test_runs_totals = self._serialize_totals(test_runs.totals)
-            serializable_invocation = test_results.invocation.dict()
-            serializable_filters = filters.dict()
-            serializable_lineage = lineage.dict()
+            serializable_invocation = test_results.invocation.model_dump()
+            serializable_filters = filters.model_dump()
+            serializable_lineage = lineage.model_dump()
 
             models_latest_invocation = invocations_api.get_models_latest_invocation()
             invocations = invocations_api.get_models_latest_invocations_data()
@@ -143,15 +143,15 @@ def _serialize_coverages(
         return {model_id: dict(coverage) for model_id, coverage in coverages.items()}
 
     def _serialize_models_runs(self, models_runs: List[ModelRunsSchema]) -> List[dict]:
-        return [model_runs.dict(by_alias=True) for model_runs in models_runs]
+        return [model_runs.model_dump(by_alias=True) for model_runs in models_runs]
 
     def _serialize_test_results(
         self, test_results: Dict[Optional[str], List[TestResultSchema]]
     ) -> Dict[Optional[str], List[dict]]:
         serializable_test_results = defaultdict(list)
         for model_unique_id, test_result in test_results.items():
             serializable_test_results[model_unique_id].extend(
-                [result.dict() for result in test_result]
+                [result.model_dump() for result in test_result]
             )
         return serializable_test_results
 
@@ -161,7 +161,7 @@ def _serialize_test_runs(
         serializable_test_runs = defaultdict(list)
         for model_unique_id, test_run in test_runs.items():
             serializable_test_runs[model_unique_id].extend(
-                [run.dict() for run in test_run]
+                [run.model_dump() for run in test_run]
             )
         return serializable_test_runs
 
@@ -170,5 +170,5 @@ def _serialize_totals(
     ) -> Dict[Optional[str], dict]:
         serialized_totals = dict()
         for model_unique_id, total in totals.items():
-            serialized_totals[model_unique_id] = total.dict()
+            serialized_totals[model_unique_id] = total.model_dump()
         return serialized_totals

elementary/monitor/api/report/schema.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ReportDataEnvSchema(BaseModel):
@@ -10,6 +10,8 @@ class ReportDataEnvSchema(BaseModel):
 
 
 class ReportDataSchema(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+
     creation_time: Optional[str] = None
     days_back: Optional[int] = None
     models: dict = dict()

elementary/monitor/api/test_management/test_management.py

@@ -16,6 +16,8 @@
 
 
 class TestManagementAPI(APIClient):
+    __test__ = False
+
     def __init__(
         self,
         dbt_runner: BaseDbtRunner,

elementary/monitor/api/tests/schema.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from elementary.monitor.api.totals_schema import TotalsSchema
 from elementary.monitor.fetchers.invocations.schema import DbtInvocationSchema
@@ -12,9 +12,6 @@ class ElementaryTestResultSchema(BaseModel):
     metrics: Optional[Union[list, dict]] = None
     result_description: Optional[str] = None
 
-    class Config:
-        smart_union = True
-
 
 class DbtTestResultSchema(BaseModel):
     display_name: Optional[str] = None
@@ -24,12 +21,12 @@ class DbtTestResultSchema(BaseModel):
 
 
 class InvocationSchema(BaseModel):
-    affected_rows: Optional[int]
+    affected_rows: Optional[int] = None
     time_utc: str
     id: str
     status: str
 
-    @validator("time_utc", pre=True)
+    @field_validator("time_utc", mode="before")
     def format_time_utc(cls, time_utc):
         return convert_partial_iso_format_to_full_iso_format(time_utc)
 
@@ -42,6 +39,8 @@ class InvocationsSchema(BaseModel):
 
 
 class TestMetadataSchema(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+
     test_unique_id: str
     elementary_unique_id: str
     database_name: Optional[str] = None
@@ -69,9 +68,6 @@ class TestResultSchema(BaseModel):
     metadata: TestMetadataSchema
     test_results: Union[DbtTestResultSchema, ElementaryTestResultSchema]
 
-    class Config:
-        smart_union = True
-
 
 class TestResultsWithTotalsSchema(BaseModel):
     results: Dict[Optional[str], List[TestResultSchema]] = dict()
@@ -81,7 +77,7 @@ class TestResultsWithTotalsSchema(BaseModel):
 
 class TestRunSchema(BaseModel):
     metadata: TestMetadataSchema
-    test_runs: Optional[InvocationsSchema]
+    test_runs: Optional[InvocationsSchema] = None
 
 
 class TestRunsWithTotalsSchema(BaseModel):

elementary/monitor/data_monitoring/data_monitoring.py

@@ -47,7 +47,7 @@ def __init__(
                 tracking.register_group(
                     "warehouse",
                     self.warehouse_info.id,
-                    self.warehouse_info.dict(),
+                    self.warehouse_info.model_dump(),
                 )
             tracking.set_env("target_name", latest_invocation.get("target_name"))
             tracking.set_env("dbt_orchestrator", latest_invocation.get("orchestrator"))

elementary/monitor/data_monitoring/report/data_monitoring_report.py

@@ -125,7 +125,7 @@ def get_report_data(
             )
             self.success = False
 
-        report_data_dict = report_data.dict()
+        report_data_dict = report_data.model_dump()
         return report_data_dict
 
     def _add_report_tracking(

elementary/monitor/data_monitoring/schema.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from elementary.monitor.alerts.model import ModelAlert
 from elementary.monitor.alerts.source_freshness import SourceFreshnessAlert
@@ -62,7 +62,7 @@ class SelectorFilterSchema(BaseModel):
     resource_types: Optional[List[ResourceType]] = None
     node_names: Optional[List[str]] = None
 
-    @validator("invocation_time", pre=True)
+    @field_validator("invocation_time", mode="before")
     def format_invocation_time(cls, invocation_time):
         if invocation_time:
             try:

elementary/monitor/fetchers/invocations/schema.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from elementary.utils.json_utils import try_load_json
 from elementary.utils.time import convert_partial_iso_format_to_full_iso_format
@@ -17,11 +17,11 @@ class DbtInvocationSchema(BaseModel):
     job_id: Optional[str] = None
     orchestrator: Optional[str] = None
 
-    @validator("detected_at", pre=True)
+    @field_validator("detected_at", mode="before")
     def format_detected_at(cls, detected_at):
         return convert_partial_iso_format_to_full_iso_format(detected_at)
 
-    @validator("selected", pre=True)
+    @field_validator("selected", mode="before")
     def format_selected(cls, selected):
         selected_list = try_load_json(selected) or []
         return " ".join(selected_list)