fix pydantic issues

Author: gaurpulkit

src/datapilot/constants.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class Extra(str, Enum):
+    allow = "allow"
+    forbid = "forbid"
+    ignore = "ignore"

src/datapilot/core/platforms/dbt/insights/checks/check_model_has_labels_keys.py

@@ -76,7 +76,7 @@ def _build_failure_result(self, model_unique_id: str, missing_keys: Sequence[str
     def _check_labels_keys(self, node_id) -> Tuple[int, Set[str]]:
         status_code = 0
         missing_keys = set(self.labels_keys) - set(self.get_node(node_id).label)
-        config = self.get_node(node_id).config.dict() if self.get_node(node_id).config else {}
+        config = self.get_node(node_id).config.model_dump() if self.get_node(node_id).config else {}
         labels = config.get("labels", {})
         label_keys = set(labels.keys())
         extra_keys = set()

src/datapilot/core/platforms/dbt/insights/checks/check_model_has_meta_keys.py

@@ -72,7 +72,7 @@ def _build_failure_result(self, model_unique_id: str, missing_keys: Sequence[str
     def _check_meta_keys(self, node_id) -> Tuple[int, Set[str], Set[str]]:
         status_code = 0
         model = self.get_node(node_id)
-        meta = model.meta.dict() if model.meta else {}
+        meta = model.meta if model.meta else {}
         model_meta_keys = set(meta.keys())
         missing_keys = None
         extra_keys = None

src/datapilot/core/platforms/dbt/insights/checks/check_source_has_freshness.py

@@ -61,7 +61,7 @@ def generate(self, *args, **kwargs) -> List[DBTModelInsightResponse]:
 
     def _check_source_has_freshness(self, source_id: str) -> List[str]:
         source = self.get_node(source_id)
-        freshness = source.freshness.dict() if source.freshness else {}
+        freshness = source.freshness.model_dump() if source.freshness else {}
 
         if not freshness:
             return False

src/datapilot/core/platforms/dbt/insights/checks/check_source_has_labels_keys.py

@@ -74,7 +74,7 @@ def _build_failure_result(self, model_unique_id: str, missing_keys: Sequence[str
     def _check_labels_keys(self, node_id) -> Tuple[int, Set[str]]:
         status_code = 0
         missing_keys = set(self.labels_keys) - set(self.get_node(node_id).label)
-        config = self.get_node(node_id).config.dict() if self.get_node(node_id).config else {}
+        config = self.get_node(node_id).config.model_dump() if self.get_node(node_id).config else {}
         labels = config.get("labels", {})
         label_keys = set(labels.keys())
         extra_keys = set()

src/datapilot/core/platforms/dbt/insights/checks/check_source_has_meta_keys.py

@@ -81,7 +81,7 @@ def generate(self, *args, **kwargs) -> List[DBTModelInsightResponse]:
     def _check_source_has_meta_keys(self, source_unique_id: str):
         status_code = 0
         model = self.get_node(source_unique_id)
-        meta = model.meta.dict() if model.meta else {}
+        meta = model.meta if model.meta else {}
         model_meta_keys = set(meta.keys())
         missing_keys = None
         extra_keys = None

src/datapilot/core/platforms/dbt/schemas/catalog.py

@@ -5,13 +5,16 @@
 from typing import Optional
 from typing import Union
 
-from pydantic.config import Extra
+from pydantic import ConfigDict
 from pydantic.main import BaseModel
 
+from datapilot.constants import Extra
+
 
 class AltimateCatalogMetadata(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     dbt_schema_version: Optional[str] = "https://schemas.getdbt.com/dbt/catalog/v1.json"
     dbt_version: Optional[str] = "0.19.0"
@@ -21,8 +24,9 @@ class Config:
 
 
 class AltimateCatalogTableMetadata(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     type: str
     database: Optional[Optional[str]] = None
@@ -33,8 +37,9 @@ class Config:
 
 
 class AltimateCatalogColumnMetadata(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     type: str
     comment: Optional[Optional[str]] = None
@@ -43,8 +48,9 @@ class Config:
 
 
 class AltimateCatalogStatsItem(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     id: str
     label: str
@@ -54,8 +60,9 @@ class Config:
 
 
 class AltimateCatalogTable(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     metadata: AltimateCatalogTableMetadata
     columns: Dict[str, AltimateCatalogColumnMetadata]
@@ -64,8 +71,9 @@ class Config:
 
 
 class AltimateCatalogCatalogV1(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     metadata: AltimateCatalogMetadata
     nodes: Dict[str, AltimateCatalogTable]

src/datapilot/core/platforms/dbt/schemas/manifest.py

@@ -19,7 +19,9 @@
 from dbt_artifacts_parser.parsers.manifest.manifest_v11 import ManifestV11
 from dbt_artifacts_parser.parsers.manifest.manifest_v11 import SupportedLanguage
 from pydantic import BaseModel
-from pydantic import Extra
+from pydantic import ConfigDict
+
+from datapilot.constants import Extra
 
 
 class DBTVersion(BaseModel):
@@ -46,8 +48,9 @@ class DBTVersion(BaseModel):
 
 
 class AltimateDocs(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     show: Optional[bool] = True
     node_color: Optional[Optional[str]] = None
@@ -98,17 +101,19 @@ class AltimateAccess(Enum):
 
 
 class AltimateDBTContract(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     enforced: Optional[bool] = False
     alias_types: Optional[bool] = True
     checksum: Optional[Optional[str]] = None
 
 
 class AltimateHook(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     sql: str
     transaction: Optional[bool] = True
@@ -117,8 +122,9 @@ class Config:
 
 # TODO: Need to add the rest of the fields
 class AltimateNodeConfig(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(
+        extra=Extra.allow,
+    )
 
     _extra: Optional[Dict[str, Any]] = None
     enabled: Optional[bool] = True
@@ -203,8 +209,9 @@ class AltimateSourceConfig(BaseModel):
 
 
 class AltimateDeferRelation(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     database: Optional[str]
     schema_name: str
@@ -213,8 +220,9 @@ class Config:
 
 
 class AltimateSeedConfig(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(
+        extra=Extra.allow,
+    )
 
     _extra: Optional[Dict[str, Any]] = None
     enabled: Optional[bool] = True
@@ -314,8 +322,9 @@ class AltimateExposureType(Enum):
 
 
 class AltimateOwner(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(
+        extra=Extra.allow,
+    )
 
     _extra: Optional[Dict[str, Any]] = None
     email: Optional[Optional[str]] = None
@@ -329,17 +338,19 @@ class AltimateMaturityEnum(Enum):
 
 
 class AltimateRefArgs(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     name: str
     package: Optional[Optional[str]] = None
     version: Optional[Optional[Union[str, float]]] = None
 
 
 class AltimateExposureConfig(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(
+        extra=Extra.allow,
+    )
 
     _extra: Optional[Dict[str, Any]] = None
     enabled: Optional[bool] = True
@@ -371,17 +382,19 @@ class AltimateManifestExposureNode(BaseModel):
 
 
 class AltimateTestMetadata(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     name: str
     kwargs: Optional[Dict[str, Any]] = None
     namespace: Optional[Optional[str]] = None
 
 
 class AltimateTestConfig(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(
+        extra=Extra.allow,
+    )
 
     _extra: Optional[Dict[str, Any]] = None
     enabled: Optional[bool] = True
@@ -433,8 +446,9 @@ class AltimateManifestTestNode(BaseModel):
 
 
 class AltimateMacroArgument(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra=Extra.forbid,
+    )
 
     name: str
     type: Optional[Optional[str]] = None