recce-nightly 0.62.0.20250417__py3-none-any.whl → 1.30.0.20251221__py3-none-any.whl

recce/artifact.py

@@ -9,9 +9,9 @@ from urllib.parse import urlencode
 import requests
 from rich.console import Console
 
-from recce.git import hosting_repo, commit_hash_from_branch, current_branch
+from recce.git import commit_hash_from_branch, current_branch, hosting_repo
 from recce.state import s3_sse_c_headers
-from recce.util.recce_cloud import RecceCloud, PresignedUrlMethod
+from recce.util.recce_cloud import PresignedUrlMethod, RecceCloud
 
 
 def verify_artifacts_path(target_path: str) -> bool:
@@ -30,10 +30,7 @@ def verify_artifacts_path(target_path: str) -> bool:
     if not os.path.isdir(target_path):
         return False
 
-    required_artifacts_files = [
-        'manifest.json',
-        'catalog.json'
-    ]
+    required_artifacts_files = ["manifest.json", "catalog.json"]
 
     if all(f in os.listdir(target_path) for f in required_artifacts_files):
         # Check if the required files are present in the target path
@@ -43,35 +40,35 @@ def verify_artifacts_path(target_path: str) -> bool:
 
 
 def parse_dbt_version(file_path: str) -> str:
-    with open(file_path, 'r') as f:
+    with open(file_path, "r", encoding="utf-8") as f:
         data = json.load(f)
 
-    dbt_version = data.get('metadata', {}).get('dbt_version', None)
+    dbt_version = data.get("metadata", {}).get("dbt_version", None)
     return dbt_version
 
 
 def archive_artifacts(target_path: str) -> (str, str):
     if verify_artifacts_path(target_path) is False:
-        raise Exception(f'Invalid target path: {target_path}')
+        raise Exception(f"Invalid target path: {target_path}")
 
-    manifest_path = os.path.join(target_path, 'manifest.json')
-    catalog_path = os.path.join(target_path, 'catalog.json')
+    manifest_path = os.path.join(target_path, "manifest.json")
+    catalog_path = os.path.join(target_path, "catalog.json")
 
     dbt_version = parse_dbt_version(manifest_path)
     if dbt_version is None:
-        raise Exception('Failed to parse dbt version from manifest.json')
+        raise Exception("Failed to parse dbt version from manifest.json")
 
     # prepare the temporary artifacts path
     tmp_dir = tempfile.mkdtemp()
-    artifacts_tar_path = os.path.join(tmp_dir, 'dbt_artifacts.tar')
-    artifacts_tar_gz_path = artifacts_tar_path + '.gz'
+    artifacts_tar_path = os.path.join(tmp_dir, "dbt_artifacts.tar")
+    artifacts_tar_gz_path = artifacts_tar_path + ".gz"
 
-    with tarfile.open(artifacts_tar_path, 'w') as tar:
-        tar.add(manifest_path, arcname='manifest.json')
-        tar.add(catalog_path, arcname='catalog.json')
+    with tarfile.open(artifacts_tar_path, "w") as tar:
+        tar.add(manifest_path, arcname="manifest.json")
+        tar.add(catalog_path, arcname="catalog.json")
 
     # Compress the tar file
-    with open(artifacts_tar_path, 'rb') as f_in, gzip.open(artifacts_tar_gz_path, 'wb') as f_out:
+    with open(artifacts_tar_path, "rb") as f_in, gzip.open(artifacts_tar_gz_path, "wb") as f_out:
         f_out.writelines(f_in)
 
     # Clean up the temporary directory
@@ -83,6 +80,64 @@ def archive_artifacts(target_path: str) -> (str, str):
     return artifacts_tar_gz_path, dbt_version
 
 
+def upload_artifacts_to_session(target_path: str, session_id: str, token: str, debug: bool = False):
+    """Upload dbt artifacts to a specific session ID in Recce Cloud."""
+    console = Console()
+    if verify_artifacts_path(target_path) is False:
+        console.print(f"[[red]Error[/red]] Invalid target path: {target_path}")
+        console.print("Please provide a valid target path containing manifest.json and catalog.json.")
+        return 1
+
+    manifest_path = os.path.join(target_path, "manifest.json")
+    catalog_path = os.path.join(target_path, "catalog.json")
+
+    # get the adapter type from the manifest file
+    with open(manifest_path, "r", encoding="utf-8") as f:
+        manifest_data = json.load(f)
+        adapter_type = manifest_data.get("metadata", {}).get("adapter_type")
+        if adapter_type is None:
+            raise Exception("Failed to parse adapter type from manifest.json")
+
+    recce_cloud = RecceCloud(token)
+
+    session = recce_cloud.get_session(session_id)
+
+    org_id = session.get("org_id")
+    if org_id is None:
+        raise Exception(f"Session ID {session_id} does not belong to any organization.")
+
+    project_id = session.get("project_id")
+    if project_id is None:
+        raise Exception(f"Session ID {session_id} does not belong to any project.")
+
+    # Get the presigned URL for uploading the artifacts using session ID
+    console.print(f'Uploading artifacts for session ID "{session_id}"')
+    presigned_urls = recce_cloud.get_upload_urls_by_session_id(org_id, project_id, session_id)
+    if debug:
+        console.rule("Debug information", style="blue")
+        console.print(f"Org ID: {org_id}")
+        console.print(f"Project ID: {project_id}")
+        console.print(f"Session ID: {session_id}")
+        console.print(f"Manifest path: {presigned_urls['manifest_url']}")
+        console.print(f"Catalog path: {presigned_urls['catalog_url']}")
+        console.print(f"Adapter type: {adapter_type}")
+
+    # Upload the compressed artifacts (no password needed for session uploads)
+    console.print(f'Uploading manifest from path "{manifest_path}"')
+    response = requests.put(presigned_urls["manifest_url"], data=open(manifest_path, "rb").read())
+    if response.status_code != 200 and response.status_code != 204:
+        raise Exception(response.text)
+    console.print(f'Uploading catalog from path "{catalog_path}"')
+    response = requests.put(presigned_urls["catalog_url"], data=open(catalog_path, "rb").read())
+    if response.status_code != 200 and response.status_code != 204:
+        raise Exception(response.text)
+
+    # Update the session metadata
+    recce_cloud.update_session(org_id, project_id, session_id, adapter_type)
+
+    return 0
+
+
 def upload_dbt_artifacts(target_path: str, branch: str, token: str, password: str, debug: bool = False):
     console = Console()
     if verify_artifacts_path(target_path) is False:
@@ -100,36 +155,33 @@ def upload_dbt_artifacts(target_path: str, branch: str, token: str, password: st
     compress_file_path, dbt_version = archive_artifacts(target_path)
     repo = hosting_repo()
     sha = commit_hash_from_branch(branch)
-    metadata = {
-        'commit': sha,
-        'dbt_version': dbt_version
-    }
+    metadata = {"commit": sha, "dbt_version": dbt_version}
 
     # Get the presigned URL for uploading the artifacts
-    presigned_url = RecceCloud(token).get_presigned_url(
+    presigned_url = RecceCloud(token).get_presigned_url_by_github_repo(
         method=PresignedUrlMethod.UPLOAD,
         repository=repo,
-        artifact_name='dbt_artifacts.tar.gz',
+        artifact_name="dbt_artifacts.tar.gz",
         branch=branch,
-        metadata=metadata
+        metadata=metadata,
     )
 
     if debug:
-        console.rule('Debug information', style='blue')
-        console.print(f'Branch: {branch}')
-        console.print(f'Commit hash: {sha}')
-        console.print(f'GitHub repository: {repo}')
-        console.print(f'Artifact path: {compress_file_path}')
-        console.print(f'DBT version: {dbt_version}')
-        console.print(f'Presigned URL: {presigned_url}')
+        console.rule("Debug information", style="blue")
+        console.print(f"Branch: {branch}")
+        console.print(f"Commit hash: {sha}")
+        console.print(f"GitHub repository: {repo}")
+        console.print(f"Artifact path: {compress_file_path}")
+        console.print(f"DBT version: {dbt_version}")
+        console.print(f"Presigned URL: {presigned_url}")
     console.print(f'Uploading the dbt artifacts from path "{target_path}" to branch "{branch}"')
 
     # Upload the compressed artifacts
 
     headers = s3_sse_c_headers(password)
     if metadata:
-        headers['x-amz-tagging'] = urlencode(metadata)
-    response = requests.put(presigned_url, data=open(compress_file_path, 'rb').read(), headers=headers)
+        headers["x-amz-tagging"] = urlencode(metadata)
+    response = requests.put(presigned_url, data=open(compress_file_path, "rb").read(), headers=headers)
     if response.status_code != 200:
         raise Exception({response.text})
 
@@ -143,29 +195,29 @@ def upload_dbt_artifacts(target_path: str, branch: str, token: str, password: st
         pass
 
 
-def download_dbt_artifacts(target_path: str, branch: str, token: str, password: str,
-                           force: bool = False,
-                           debug: bool = False):
+def download_dbt_artifacts(
+    target_path: str, branch: str, token: str, password: str, force: bool = False, debug: bool = False
+):
     console = Console()
     repo = hosting_repo()
     sha = None
     dbt_version = None
 
-    presigned_url, tags = RecceCloud(token).get_download_presigned_url_with_tags(
+    presigned_url, tags = RecceCloud(token).get_download_presigned_url_by_github_repo_with_tags(
         repository=repo,
-        artifact_name='dbt_artifacts.tar.gz',
+        artifact_name="dbt_artifacts.tar.gz",
         branch=branch,
     )
     if tags:
-        sha = tags.get('commit')
-        dbt_version = tags.get('dbt_version')
+        sha = tags.get("commit")
+        dbt_version = tags.get("dbt_version")
 
     if debug:
-        console.rule('Debug information', style='blue')
-        console.print(f'Git Branch: {branch}')
-        console.print(f'Git Commit: {sha}')
-        console.print(f'GitHub repository: {repo}')
-        console.print(f'DBT version: {dbt_version}')
+        console.rule("Debug information", style="blue")
+        console.print(f"Git Branch: {branch}")
+        console.print(f"Git Commit: {sha}")
+        console.print(f"GitHub repository: {repo}")
+        console.print(f"DBT version: {dbt_version}")
     console.print(f'Downloading from branch: "{branch}" and extracting to "{target_path}"')
 
     headers = s3_sse_c_headers(password)
@@ -177,16 +229,17 @@ def download_dbt_artifacts(target_path: str, branch: str, token: str, password:
     if os.path.exists(target_path):
         if not force:
             raise Exception(
-                f'Path {target_path} already exists. Please provide a new path or use \'--force\' option to overwrite the existing folder.')
-        console.print(f'[[yellow]Warning[/yellow]] Overwrite existing path: {target_path}')
+                f"Path {target_path} already exists. Please provide a new path or use '--force' option to overwrite the existing folder."
+            )
+        console.print(f"[[yellow]Warning[/yellow]] Overwrite existing path: {target_path}")
         shutil.rmtree(target_path)
     os.mkdir(target_path)
 
-    tar_gz_file = os.path.join(target_path, 'dbt_artifacts.tar.gz')
-    with open(tar_gz_file, 'wb') as f:
+    tar_gz_file = os.path.join(target_path, "dbt_artifacts.tar.gz")
+    with open(tar_gz_file, "wb") as f:
         f.write(response.content)
 
-    with tarfile.open(tar_gz_file, 'r') as tar:
+    with tarfile.open(tar_gz_file, "r") as tar:
         tar.extractall(path=target_path)
 
     # Clean up the compressed artifacts
@@ -196,3 +249,18 @@ def download_dbt_artifacts(target_path: str, branch: str, token: str, password:
     except FileNotFoundError:
         pass
     return 0
+
+
+def delete_dbt_artifacts(branch: str, token: str, debug: bool = False):
+    """Delete dbt artifacts from a specific branch in Recce Cloud."""
+    console = Console()
+    repo = hosting_repo()
+
+    if debug:
+        console.rule("Debug information", style="blue")
+        console.print(f"Git Branch: {branch}")
+        console.print(f"GitHub repository: {repo}")
+
+    console.print(f'Deleting dbt artifacts from branch: "{branch}"')
+
+    RecceCloud(token).purge_artifacts(repo, branch=branch)