feat: added catalog and env support (#36)

* feat: added catalog support

* feat: added env

* WIP

* WIP

* WIP

* feat: remove file ids

* WIP

* WIP

* WIP

* fix: imports

---------

Co-authored-by: surya <[email protected]>

Author: AdiGajbhiye

src/datapilot/clients/altimate/client.py

@@ -62,7 +62,7 @@ def post(self, endpoint, data=None, timeout=None):
         response = requests.post(url, headers=headers, json=data, timeout=timeout)
         self.logger.debug(f"Received POST response with status: {response.status_code }")
 
-        return response
+        return response.json()
 
     def put(self, endpoint, data, timeout=None):
         url = f"{self.base_url}{endpoint}"
@@ -83,3 +83,7 @@ def get_signed_url(self, params=None):
     def validate_credentials(self):
         endpoint = "/dbt/v3/validate-credentials"
         return self.get(endpoint)
+
+    def start_dbt_ingestion(self, params=None):
+        endpoint = "/dbt/v1/start_dbt_ingestion"
+        return self.post(endpoint, data=params)

src/datapilot/clients/altimate/utils.py

@@ -46,18 +46,22 @@ def validate_credentials(
     return api_client.validate_credentials()
 
 
-def onboard_manifest(api_token, tenant, dbt_core_integration_id, manifest_path, backend_url) -> Dict:
+def onboard_file(api_token, tenant, dbt_core_integration_id, dbt_core_integration_environment, file_type, file_path, backend_url) -> Dict:
     api_client = APIClient(api_token, base_url=backend_url, tenant=tenant)
 
-    params = {"dbt_core_integration_id": dbt_core_integration_id, "file_type": "manifest"}
+    params = {
+        "dbt_core_integration_id": dbt_core_integration_id,
+        "dbt_core_integration_environment_type": dbt_core_integration_environment,
+        "file_type": file_type,
+    }
     signed_url_data = api_client.get_signed_url(params)
     if signed_url_data:
         signed_url = signed_url_data.get("url")
         file_id = signed_url_data.get("dbt_core_integration_file_id")
         api_client.log(f"Received signed URL: {signed_url}")
         api_client.log(f"Received File ID: {file_id}")
 
-        upload_response = upload_content_to_signed_url(manifest_path, signed_url)
+        upload_response = upload_content_to_signed_url(file_path, signed_url)
 
         if upload_response:
             verify_params = {"dbt_core_integration_file_id": file_id}
@@ -73,3 +77,20 @@ def onboard_manifest(api_token, tenant, dbt_core_integration_id, manifest_path,
             "ok": False,
             "message": "Error in uploading the manifest.                                                                                                                              ",
         }
+
+
+def start_dbt_ingestion(api_token, tenant, dbt_core_integration_id, dbt_core_integration_environment, backend_url):
+    api_client = APIClient(api_token, base_url=backend_url, tenant=tenant)
+    params = {
+        "dbt_core_integration_id": dbt_core_integration_id,
+        "dbt_core_integration_environment_type": dbt_core_integration_environment,
+    }
+    data = api_client.start_dbt_ingestion(params)
+    if data and data.get("ok"):
+        return {"ok": True}
+    else:
+        api_client.log("Error starting dbt ingestion worker")
+        return {
+            "ok": False,
+            "message": "Error starting dbt ingestion worker.                                                                                                                              ",
+        }

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

@@ -3,7 +3,8 @@
 import click
 
 from datapilot.clients.altimate.utils import check_token_and_instance
-from datapilot.clients.altimate.utils import onboard_manifest
+from datapilot.clients.altimate.utils import onboard_file
+from datapilot.clients.altimate.utils import start_dbt_ingestion
 from datapilot.clients.altimate.utils import validate_credentials
 from datapilot.config.config import load_config
 from datapilot.core.platforms.dbt.constants import MODEL
@@ -87,9 +88,21 @@ def project_health(manifest_path, catalog_path, config_path=None, select=None):
 @click.option("--token", prompt="API Token", help="Your API token for authentication.")
 @click.option("--instance-name", prompt="Instance Name", help="Your tenant ID.")
 @click.option("--dbt_core_integration_id", prompt="DBT Core Integration ID", help="DBT Core Integration ID")
+@click.option(
+    "--dbt_core_integration_environment", default="PROD", prompt="DBT Core Integration Environment", help="DBT Core Integration Environment"
+)
 @click.option("--manifest-path", required=True, prompt="Manifest Path", help="Path to the manifest file.")
+@click.option("--catalog-path", required=False, prompt=False, help="Path to the catalog file.")
 @click.option("--backend-url", required=False, help="Altimate's Backend URL", default="https://api.myaltimate.com")
-def onboard(token, instance_name, dbt_core_integration_id, manifest_path, backend_url="https://api.myaltimate.com", env=None):
+def onboard(
+    token,
+    instance_name,
+    dbt_core_integration_id,
+    dbt_core_integration_environment,
+    manifest_path,
+    catalog_path,
+    backend_url="https://api.myaltimate.com",
+):
     """Onboard a manifest file to DBT."""
     check_token_and_instance(token, instance_name)
 
@@ -104,9 +117,27 @@ def onboard(token, instance_name, dbt_core_integration_id, manifest_path, backen
         click.echo(f"Error: {e}")
         return
 
-    response = onboard_manifest(token, instance_name, dbt_core_integration_id, manifest_path, backend_url)
-
+    response = onboard_file(
+        token, instance_name, dbt_core_integration_id, dbt_core_integration_environment, "manifest", manifest_path, backend_url
+    )
     if response["ok"]:
         click.echo("Manifest onboarded successfully!")
     else:
         click.echo(f"{response['message']}")
+
+    if not catalog_path:
+        return
+
+    response = onboard_file(
+        token, instance_name, dbt_core_integration_id, dbt_core_integration_environment, "catalog", catalog_path, backend_url
+    )
+    if response["ok"]:
+        click.echo("Catalog onboarded successfully!")
+    else:
+        click.echo(f"{response['message']}")
+
+    response = start_dbt_ingestion(token, instance_name, dbt_core_integration_id, dbt_core_integration_environment, backend_url)
+    if response["ok"]:
+        click.echo("Onboarding completed successfully!")
+    else:
+        click.echo(f"{response['message']}")