feat: fix project governance exposures issue and add integration name option enhancement

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.21
+current_version = 0.0.23
 commit = True
 tag = True
 

docs/conf.py

@@ -15,7 +15,7 @@
 year = "2024"
 author = "Altimate Inc."
 copyright = f"{year}, {author}"
-version = release = "0.0.21"
+version = release = "0.0.23"
 
 pygments_style = "trac"
 templates_path = ["."]

setup.py

@@ -13,7 +13,7 @@ def read(*names, **kwargs):
 
 setup(
     name="altimate-datapilot-cli",
-    version="0.0.22",
+    version="0.0.23",
     license="MIT",
     description="Assistant for Data Teams",
     long_description="{}\n{}".format(

src/datapilot/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.21"
+__version__ = "0.0.23"

src/datapilot/clients/altimate/client.py

@@ -110,3 +110,8 @@ def get_all_dbt_configs(self):
         endpoint = "/dbtconfig/"
         params = {"size": 100}
         return self.get(endpoint, params=params)
+
+    def get_all_integrations(self):
+        """Get all integrations"""
+        endpoint = "/dbt/v1/project_integrations/"
+        return self.get(endpoint)

src/datapilot/clients/altimate/utils.py

@@ -134,3 +134,28 @@ def get_all_dbt_configs(
     """Get all DBT configs from the API."""
     api_client = APIClient(api_token=api_token, base_url=backend_url, tenant=tenant)
     return api_client.get_all_dbt_configs()
+
+
+def get_all_integrations(
+    api_token,
+    tenant,
+    backend_url,
+):
+    """Get all integrations from the API."""
+    api_client = APIClient(api_token=api_token, base_url=backend_url, tenant=tenant)
+    return api_client.get_all_integrations()
+
+
+def resolve_integration_name_to_id(
+    integration_name,
+    api_token,
+    tenant,
+    backend_url,
+):
+    """Resolve integration name to ID."""
+    integrations = get_all_integrations(api_token, tenant, backend_url)
+    if integrations:
+        matching_integrations = [i for i in integrations if i.get("name") == integration_name]
+        if matching_integrations:
+            return matching_integrations[0].get("id")
+    return None

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

@@ -6,6 +6,7 @@
 from datapilot.clients.altimate.utils import check_token_and_instance
 from datapilot.clients.altimate.utils import get_all_dbt_configs
 from datapilot.clients.altimate.utils import onboard_file
+from datapilot.clients.altimate.utils import resolve_integration_name_to_id
 from datapilot.clients.altimate.utils import start_dbt_ingestion
 from datapilot.clients.altimate.utils import validate_credentials
 from datapilot.clients.altimate.utils import validate_permissions
@@ -136,9 +137,14 @@ def project_health(
     "--dbt_core_integration_id",
     "--dbt_integration_id",
     "dbt_integration_id",  # This is the parameter name that will be passed to the function
-    prompt="DBT Integration ID",
     help="DBT Core Integration ID or DBT Integration ID",
 )
+@click.option(
+    "--dbt_core_integration_name",
+    "--dbt_integration_name",
+    "dbt_integration_name",  # This is the parameter name that will be passed to the function
+    help="DBT Core Integration Name or DBT Integration Name (alternative to ID)",
+)
 @click.option(
     "--dbt_core_integration_environment",
     "--dbt_integration_environment",
@@ -154,11 +160,12 @@ def onboard(
     instance_name,
     backend_url,
     dbt_integration_id,
+    dbt_integration_name,
     dbt_integration_environment,
     manifest_path,
     catalog_path,
 ):
-    """Onboard a manifest file to DBT."""
+    """Onboard a manifest file to DBT. You can specify either --dbt_integration_id or --dbt_integration_name."""
 
     # For onboard command, token and instance_name are required
     if not token:
@@ -176,6 +183,21 @@ def onboard(
         click.echo("Error: You don't have permission to perform this action.")
         return
 
+    # Resolve integration name to ID if name is provided instead of ID
+    if not dbt_integration_id and not dbt_integration_name:
+        dbt_integration_id = click.prompt("DBT Integration ID")
+    elif dbt_integration_name and not dbt_integration_id:
+        click.echo(f"Resolving integration name '{dbt_integration_name}' to ID...")
+        resolved_id = resolve_integration_name_to_id(dbt_integration_name, token, instance_name, backend_url)
+        if resolved_id:
+            dbt_integration_id = resolved_id
+            click.echo(f"Found integration ID: {dbt_integration_id}")
+        else:
+            click.echo(f"Error: Integration with name '{dbt_integration_name}' not found.")
+            return
+    elif dbt_integration_name and dbt_integration_id:
+        click.echo("Warning: Both integration ID and name provided. Using ID and ignoring name.")
+
     try:
         load_manifest(manifest_path)
     except Exception as e:

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

@@ -15,7 +15,7 @@
 from datapilot.core.platforms.dbt.insights import INSIGHTS
 from datapilot.core.platforms.dbt.insights.schema import DBTInsightResult
 from datapilot.core.platforms.dbt.insights.schema import DBTModelInsightResponse
-from datapilot.core.platforms.dbt.schemas.manifest import Catalog
+from datapilot.core.platforms.dbt.schemas.catalog import Catalog
 from datapilot.core.platforms.dbt.schemas.manifest import Manifest
 from datapilot.core.platforms.dbt.utils import get_models
 from datapilot.utils.formatting.utils import RED

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.catalog import Catalog
+from datapilot.core.platforms.dbt.schemas.catalog 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/catalog.py

@@ -5,6 +5,8 @@
 from typing import Optional
 from typing import Union
 
+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 pydantic.main import BaseModel
 
 
@@ -52,3 +54,19 @@ class AltimateCatalogCatalogV1(BaseModel):
     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

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

@@ -5,7 +5,6 @@
 from typing import Optional
 from typing import Union
 
-from dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1
 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,8 +42,6 @@ class DBTVersion(BaseModel):
     ManifestV1,
 ]
 
-Catalog = CatalogV1
-
 
 class AltimateDocs(BaseModel):
     show: Optional[bool] = True

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
@@ -18,11 +17,12 @@
 from datapilot.core.platforms.dbt.constants import STAGING
 from datapilot.core.platforms.dbt.exceptions import AltimateInvalidManifestError
 from datapilot.core.platforms.dbt.factory import DBTFactory
+from datapilot.core.platforms.dbt.schemas.catalog import Catalog
+from datapilot.core.platforms.dbt.schemas.catalog import CatalogV1
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateManifestExposureNode
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateManifestNode
 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 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.catalog import CatalogV1
 from datapilot.core.platforms.dbt.wrappers.catalog.wrapper import BaseCatalogWrapper
 
 

src/datapilot/core/platforms/dbt/wrappers/manifest/v12/wrapper.py

@@ -12,6 +12,7 @@
 from datapilot.core.platforms.dbt.constants import SINGULAR
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateDBTContract
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateDependsOn
+from datapilot.core.platforms.dbt.schemas.manifest import AltimateExposureConfig
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateExposureType
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateExternalTable
 from datapilot.core.platforms.dbt.schemas.manifest import AltimateFileHash
@@ -204,7 +205,7 @@ def _get_exposure(self, exposure: ExposureNode) -> AltimateManifestExposureNode:
             maturity=AltimateMaturityEnum(exposure.maturity.value) if exposure.maturity else None,
             meta=exposure.meta,
             tags=exposure.tags,
-            config=AltimateSourceConfig(**exposure.config.model_dump()) if exposure.config else None,
+            config=AltimateExposureConfig(**exposure.config.model_dump()) if exposure.config else None,
             unrendered_config=exposure.unrendered_config,
             url=exposure.url,
             depends_on=(

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.catalog 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