refactor: replace CatalogV1 import with custom implementation and update catalog loading logic

Author: gaurpulkit

src/datapilot/core/platforms/dbt/factory.py

@@ -1,9 +1,10 @@
-from dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1
+# Remove the import of CatalogV1 from dbt_artifacts_parser since we use our custom version
 from dbt_artifacts_parser.parsers.manifest.manifest_v10 import ManifestV10
 from dbt_artifacts_parser.parsers.manifest.manifest_v11 import ManifestV11
 from dbt_artifacts_parser.parsers.manifest.manifest_v12 import ManifestV12
 
 from datapilot.core.platforms.dbt.schemas.manifest import Catalog
+from datapilot.core.platforms.dbt.schemas.manifest import CatalogV1
 from datapilot.core.platforms.dbt.schemas.manifest import Manifest
 from datapilot.core.platforms.dbt.wrappers.catalog.v1.wrapper import CatalogV1Wrapper
 from datapilot.core.platforms.dbt.wrappers.manifest.v10.wrapper import ManifestV10Wrapper

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

@@ -5,7 +5,8 @@
 from typing import Optional
 from typing import Union
 
-from dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1
+from dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1 as BaseCatalogV1
+from dbt_artifacts_parser.parsers.catalog.catalog_v1 import Metadata as BaseMetadata
 from dbt_artifacts_parser.parsers.manifest.manifest_v1 import ManifestV1
 from dbt_artifacts_parser.parsers.manifest.manifest_v2 import ManifestV2
 from dbt_artifacts_parser.parsers.manifest.manifest_v3 import ManifestV3
@@ -43,6 +44,19 @@ class DBTVersion(BaseModel):
     ManifestV1,
 ]
 
+
+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
 
 

src/datapilot/core/platforms/dbt/utils.py

@@ -6,7 +6,6 @@
 from typing import Tuple
 from typing import Union
 
-from dbt_artifacts_parser.parser import parse_catalog
 from dbt_artifacts_parser.parser import parse_manifest
 
 from datapilot.core.platforms.dbt.constants import BASE
@@ -23,6 +22,7 @@
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateManifestSourceNode
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateManifestTestNode
 from datapilot.core.platforms.dbt.schemas.manifest import Catalog
+from datapilot.core.platforms.dbt.schemas.manifest import CatalogV1
 from datapilot.core.platforms.dbt.schemas.manifest import Manifest
 from datapilot.exceptions.exceptions import AltimateFileNotFoundError
 from datapilot.exceptions.exceptions import AltimateInvalidJSONError
@@ -83,14 +83,14 @@ def load_catalog(catalog_path: str) -> Catalog:
     try:
         catalog_dict = load_json(catalog_path)
     except FileNotFoundError as e:
-        raise AltimateFileNotFoundError(f"Manifest file not found: {catalog_path}. Error: {e}") from e
+        raise AltimateFileNotFoundError(f"Catalog file not found: {catalog_path}. Error: {e}") from e
     except ValueError as e:
         raise AltimateInvalidJSONError(f"Invalid JSON file: {catalog_path}. Error: {e}") from e
 
     try:
-        catalog: Catalog = parse_catalog(catalog_dict)
+        catalog: Catalog = CatalogV1(**catalog_dict)
     except ValueError as e:
-        raise AltimateInvalidManifestError(f"Invalid manifest file: {catalog_path}. Error: {e}") from e
+        raise AltimateInvalidManifestError(f"Invalid catalog file: {catalog_path}. Error: {e}") from e
 
     return catalog
 

src/datapilot/core/platforms/dbt/wrappers/catalog/v1/wrapper.py

@@ -1,5 +1,4 @@
-from dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1
-
+from datapilot.core.platforms.dbt.schemas.manifest import CatalogV1
 from datapilot.core.platforms.dbt.wrappers.catalog.wrapper import BaseCatalogWrapper
 
 

src/datapilot/utils/utils.py

@@ -9,10 +9,10 @@
 from typing import List
 from typing import Union
 
-from dbt_artifacts_parser.parser import parse_catalog
 from dbt_artifacts_parser.parser import parse_manifest
 
 from datapilot.config.config import load_config
+from datapilot.core.platforms.dbt.schemas.manifest import CatalogV1
 from datapilot.schemas.nodes import ModelNode
 from datapilot.schemas.nodes import SourceNode
 
@@ -306,7 +306,7 @@ def generate_partial_manifest_catalog(changed_files, base_path: str = "./"):
         catalog = fill_catalog(table_columns_map, manifest, catalog, sources, "sources")
 
         selected_models = [node.unique_id for node in nodes + sources]
-        return selected_models, parse_manifest(manifest), parse_catalog(catalog)
+        return selected_models, parse_manifest(manifest), CatalogV1(**catalog)
     except Exception as e:
         raise Exception("Unable to generate partial manifest and catalog") from e