feat: fix project governance exposures issue and add integration name option enhancement (#65)

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

* fix: update exposure config handling in ManifestV12Wrapper to use AltimateExposureConfig

* feat: add functionality to resolve integration names to IDs and enhance onboarding options in CLI

* Bump version: 0.0.22 → 0.0.23 and update relevant files to reflect the new version.

* refactor: update import paths for Catalog and CatalogV1 to reflect new schema structure

Author: gaurpulkit

.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