refactor: replace Pydantic BaseModel with dataclasses

src/datapilot/core/insights/schema.py

@@ -1,16 +1,16 @@
+from dataclasses import dataclass
 from enum import Enum
 from typing import Dict
 
-from pydantic import BaseModel
-
 
 class Severity(Enum):
     INFO = "INFO"
     WARNING = "WARNING"
     ERROR = "ERROR"
 
 
-class InsightResult(BaseModel):
+@dataclass
+class InsightResult:
     name: str
     type: str
     message: str
@@ -19,6 +19,7 @@ class InsightResult(BaseModel):
     metadata: Dict
 
 
-class InsightResponse(BaseModel):
+@dataclass
+class InsightResponse:
     insight: InsightResult
     severity: Severity = Severity.ERROR

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

@@ -1,73 +1,64 @@
+from dataclasses import dataclass
 from datetime import datetime
 from typing import ClassVar
 from typing import Dict
 from typing import List
 from typing import Optional
 from typing import Union
 
-from pydantic.main import BaseModel
+from vendor.dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1
 
-from vendor.dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1 as BaseCatalogV1
-from vendor.dbt_artifacts_parser.parsers.catalog.catalog_v1 import Metadata as BaseMetadata
 
-
-class AltimateCatalogMetadata(BaseModel):
+@dataclass
+class AltimateCatalogMetadata:
     dbt_schema_version: Optional[str] = "https://schemas.getdbt.com/dbt/catalog/v1.json"
     dbt_version: Optional[str] = "0.19.0"
     generated_at: Optional[datetime] = "2021-02-10T04:42:33.680487Z"
     invocation_id: Optional[Optional[str]] = None
     env: ClassVar[Optional[Dict[str, str]]] = {}
 
 
-class AltimateCatalogTableMetadata(BaseModel):
+@dataclass
+class AltimateCatalogTableMetadata:
     type: str
-    database: Optional[Optional[str]] = None
     schema_name: str
     name: str
+    database: Optional[Optional[str]] = None
     comment: Optional[Optional[str]] = None
     owner: Optional[Optional[str]] = None
 
 
-class AltimateCatalogColumnMetadata(BaseModel):
+@dataclass
+class AltimateCatalogColumnMetadata:
     type: str
-    comment: Optional[Optional[str]] = None
     index: int
     name: str
+    comment: Optional[Optional[str]] = None
 
 
-class AltimateCatalogStatsItem(BaseModel):
+@dataclass
+class AltimateCatalogStatsItem:
     id: str
     label: str
+    include: bool
     value: Optional[Optional[Union[bool, str, float]]] = None
     description: Optional[Optional[str]] = None
-    include: bool
 
 
-class AltimateCatalogTable(BaseModel):
+@dataclass
+class AltimateCatalogTable:
     metadata: AltimateCatalogTableMetadata
     columns: Dict[str, AltimateCatalogColumnMetadata]
     stats: Dict[str, AltimateCatalogStatsItem]
     unique_id: Optional[Optional[str]] = None
 
 
-class AltimateCatalogCatalogV1(BaseModel):
+@dataclass
+class AltimateCatalogCatalogV1:
     metadata: AltimateCatalogMetadata
     nodes: Dict[str, AltimateCatalogTable]
     sources: Dict[str, AltimateCatalogTable]
     errors: Optional[Optional[List[str]]] = None
 
 
-# Custom classes to handle extra fields in newer dbt versions
-class Metadata(BaseMetadata):
-    class Config:
-        extra = "allow"  # Allow extra fields in metadata
-
-
-class CatalogV1(BaseCatalogV1):
-    metadata: Metadata  # Use our custom metadata class
-
-    class Config:
-        extra = "allow"  # Allow extra fields
-
-
 Catalog = CatalogV1