@datarecce/ui 0.1.30 → 0.1.31

package/recce-source/recce/apis/run_func.py

@@ -0,0 +1,258 @@
+import asyncio
+import logging
+from typing import List, Optional
+
+from recce.core import default_context
+from recce.exceptions import RecceException
+from recce.models import Run, RunDAO, RunType
+from recce.models.types import RunStatus
+
+running_tasks = {}
+logger = logging.getLogger("uvicorn")
+
+
+def _get_ref_model(sql_template: str) -> Optional[str]:
+    import re
+
+    pattern = r'\bref\(["\']?(\w+)["\']?\)\s*}}'
+    matches = re.findall(pattern, sql_template)
+    if len(matches) == 1:
+        ref = matches[0]
+        return ref
+
+    return None
+
+
+def generate_run_name(run):
+    # parse utc time with timezone
+
+    import dateutil
+
+    run_type = run.type
+    params = run.params
+    now = dateutil.parser.parse(run.run_at)
+
+    if run_type == RunType.QUERY:
+        ref = _get_ref_model(params.get("sql_template"))
+        if ref:
+            return f"query of {ref}".capitalize()
+        return f"{'query'.capitalize()} - {now}"
+    elif run_type == RunType.QUERY_DIFF:
+        ref = _get_ref_model(params.get("sql_template"))
+        if ref:
+            return f"query diff of {ref}".capitalize()
+        return f"{'query diff'.capitalize()} - {now}"
+    elif run_type == RunType.VALUE_DIFF:
+        model = params.get("model")
+        return f"value diff of {model}".capitalize()
+    elif run_type == RunType.VALUE_DIFF_DETAIL:
+        model = params.get("model")
+        return f"value diff detail of {model}".capitalize()
+    elif run_type == RunType.PROFILE_DIFF:
+        model = params.get("model")
+        return f"profile diff of {model}".capitalize()
+    elif run_type == RunType.ROW_COUNT_DIFF:
+        nodes = params.get("node_names")
+        if nodes:
+            if len(nodes) == 1:
+                node = nodes[0]
+                return f"row count diff of {node}".capitalize()
+            else:
+                return f"row count of {len(nodes)} nodes".capitalize()
+        else:
+            return "row count of multiple nodes".capitalize()
+    elif run_type == RunType.TOP_K_DIFF:
+        model = params.get("model")
+        column = params.get("column_name")
+        return f"top-k diff of {model}.{column} ".capitalize()
+    elif run_type == RunType.HISTOGRAM_DIFF:
+        model = params.get("model")
+        column = params.get("column_name")
+        return f"histogram diff of {model}.{column} ".capitalize()
+    else:
+        return f"{'run'.capitalize()} - {now}"
+
+
+def create_task(run_type: RunType, params: dict):
+    context = default_context()
+    if context is not None and context.adapter_type == "sqlmesh":
+        from recce.adapter.sqlmesh_adapter import (
+            sqlmesh_supported_registry as sqlmesh_registry,
+        )
+
+        registry = sqlmesh_registry
+    else:
+        from recce.adapter.dbt_adapter import dbt_supported_registry as dbt_registry
+
+        registry = dbt_registry
+
+    taskClz = registry.get(run_type)
+    if not taskClz:
+        raise NotImplementedError()
+    return taskClz(params)
+
+
+def submit_run(type, params, check_id=None):
+    try:
+        run_type = RunType(type)
+    except ValueError:
+        raise RecceException(f"Run type '{type}' not supported")
+
+    try:
+        task = create_task(run_type, params)
+    except NotImplementedError:
+        raise RecceException(f"Run type '{type}' not supported")
+
+    context = default_context()
+    if context.review_mode is True:
+        from recce.adapter.dbt_adapter import DbtAdapter
+
+        dbt_adaptor: DbtAdapter = context.adapter
+        if dbt_adaptor.adapter is None:
+            raise RecceException("Recce Server is not launched under DBT project folder.")
+
+    run = Run(type=run_type, params=params, check_id=check_id, status=RunStatus.RUNNING)
+    run.name = generate_run_name(run)
+    RunDAO().create(run)
+
+    loop = asyncio.get_running_loop()
+    running_tasks[run.run_id] = task
+
+    def progress_listener(message=None, percentage=None):
+        run.progress = {"message": message, "percentage": percentage}
+
+    task.progress_listener = progress_listener
+
+    async def update_run_result(run, result, error, updated_params=None):
+        """Update run with result, error, and optionally updated params."""
+        if run is None:
+            return
+        if result is not None:
+            run.result = result
+            run.status = RunStatus.FINISHED
+        if updated_params is not None:
+            # Merge updated params (preserves any fields not in updated_params)
+            run.params.update(updated_params)
+        if error is not None:
+            failed_reason = str(error) if str(error) != "None" else repr(error)
+            run.error = failed_reason
+            if run.status != RunStatus.CANCELLED:
+                run.status = RunStatus.FAILED
+        run.progress = None
+
+    def fn():
+        try:
+            result = task.execute()
+
+            # Extract updated params from task after execution
+            updated_params = None
+            if hasattr(task, "params") and task.params is not None:
+                # Serialization logic:
+                # - Most tasks use Pydantic models (v2: model_dump, v1: dict)
+                # - Some tasks may use plain dicts
+                # - If params is an unexpected type, log a warning for debugging
+                # - Handle the case where model_dump() or dict() raises an exception.
+                try:
+                    if hasattr(task.params, "model_dump"):
+                        updated_params = task.params.model_dump()
+                    elif hasattr(task.params, "dict"):
+                        updated_params = task.params.dict()
+                    elif isinstance(task.params, dict):
+                        updated_params = task.params
+                    else:
+                        logger.warning(
+                            f"Could not serialize task.params for run_id={run.run_id}: "
+                            f"unexpected type {type(task.params)} with value {repr(task.params)}"
+                        )
+                except Exception as e:
+                    logger.warning(f"Failed to serialize task.params: {e}")
+                    updated_params = None
+
+            asyncio.run_coroutine_threadsafe(update_run_result(run, result, None, updated_params), loop)
+            return result
+        except BaseException as e:
+            asyncio.run_coroutine_threadsafe(update_run_result(run, None, e, None), loop)
+            if isinstance(e, RecceException) and e.is_raise is False:
+                return None
+            import sentry_sdk
+
+            sentry_sdk.capture_exception(e)
+            failed_reason = str(e) if str(e) != "None" else repr(e)
+            failed_reason = failed_reason.replace(". ", ".\n")
+            logger.error(f"Failed to execute {run_type} task: {failed_reason}")
+            return None
+
+    future = loop.run_in_executor(None, fn)
+    return run, future
+
+
+def cancel_run(run_id):
+    run = RunDAO().find_run_by_id(run_id)
+    if run is None:
+        raise RecceException(f"Run ID '{run_id}' not found")
+
+    task = running_tasks.get(run_id)
+    if task is None:
+        raise RecceException(f"Run task for Run ID '{run_id}' not found")
+
+    task.cancel()
+    run.status = RunStatus.CANCELLED
+
+
+def materialize_run_results(runs: List[Run], nodes: List[str] = None):
+    """
+    Materialize the run results for nodes. It walks through all runs and get the last results for primary run types.
+
+    The result format
+    {
+       'node_id': {
+          'row_count_diff': {
+            'run_id': '<run_id>',
+            'result': '<result>'
+          },
+          'value_diff': {
+            'run_id': '<run_id>',
+            'result': '<result>'
+          },
+       },
+    }
+    """
+
+    context = default_context()
+    if context:
+        mame_to_unique_id = context.build_name_to_unique_id_index(excluded_types={"semantic_model", "metric"})
+    else:
+        mame_to_unique_id = {}
+
+    result = {}
+    for run in runs:
+        if not run.result:
+            continue
+
+        if run.type == RunType.ROW_COUNT_DIFF:
+            for model_name, node_run_result in run.result.items():
+                key = mame_to_unique_id.get(model_name, model_name)
+
+                if nodes:
+                    if key not in nodes:
+                        continue
+
+                if model_name not in result:
+                    node_result = result[key] = {}
+                else:
+                    node_result = result.get(key)
+                node_result["row_count_diff"] = {"run_id": run.run_id, "result": node_run_result}
+        elif run.type == RunType.ROW_COUNT:
+            for model_name, node_run_result in run.result.items():
+                key = mame_to_unique_id.get(model_name, model_name)
+
+                if nodes:
+                    if key not in nodes:
+                        continue
+
+                if model_name not in result:
+                    node_result = result[key] = {}
+                else:
+                    node_result = result.get(key)
+                node_result["row_count"] = {"run_id": run.run_id, "result": node_run_result}
+    return result

package/recce-source/recce/artifact.py

@@ -0,0 +1,266 @@
+import gzip
+import json
+import os
+import shutil
+import tarfile
+import tempfile
+from urllib.parse import urlencode
+
+import requests
+from rich.console import Console
+
+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 PresignedUrlMethod, RecceCloud
+
+
+def verify_artifacts_path(target_path: str) -> bool:
+    """
+    Verify if the target path is a valid artifacts path.
+
+    :param target_path: the target path to check
+    :return: True if the target path is a valid artifacts path, False otherwise
+    """
+    if not target_path:
+        return False
+
+    if not os.path.exists(target_path):
+        return False
+
+    if not os.path.isdir(target_path):
+        return False
+
+    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
+        return True
+
+    return False
+
+
+def parse_dbt_version(file_path: str) -> str:
+    with open(file_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    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}")
+
+    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")
+
+    # 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"
+
+    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:
+        f_out.writelines(f_in)
+
+    # Clean up the temporary directory
+    try:
+        os.remove(artifacts_tar_path)
+    except FileNotFoundError:
+        pass
+
+    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:
+        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
+
+    if branch != current_branch():
+        console.print(
+            f"[[yellow]Warning[/yellow]] You are uploading the dbt artifacts as branch '{branch}'. "
+            f"However, the current branch is '{current_branch()}'."
+        )
+        console.print("Please make sure you are uploading the dbt artifacts to the correct branch.")
+
+    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}
+
+    # Get the presigned URL for uploading the artifacts
+    presigned_url = RecceCloud(token).get_presigned_url_by_github_repo(
+        method=PresignedUrlMethod.UPLOAD,
+        repository=repo,
+        artifact_name="dbt_artifacts.tar.gz",
+        branch=branch,
+        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.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)
+    if response.status_code != 200:
+        raise Exception({response.text})
+
+    # Clean up the compressed artifacts
+    try:
+        # Remove the compressed artifacts
+        os.remove(compress_file_path)
+        # Clean up the temporary directory
+        os.rmdir(os.path.dirname(compress_file_path))
+    except FileNotFoundError:
+        pass
+
+
+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_by_github_repo_with_tags(
+        repository=repo,
+        artifact_name="dbt_artifacts.tar.gz",
+        branch=branch,
+    )
+    if tags:
+        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.print(f'Downloading from branch: "{branch}" and extracting to "{target_path}"')
+
+    headers = s3_sse_c_headers(password)
+    response = requests.get(presigned_url, headers=headers)
+
+    if response.status_code != 200:
+        raise Exception(response.text)
+
+    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}")
+        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:
+        f.write(response.content)
+
+    with tarfile.open(tar_gz_file, "r") as tar:
+        tar.extractall(path=target_path)
+
+    # Clean up the compressed artifacts
+    try:
+        # Remove the compressed artifacts
+        os.remove(tar_gz_file)
+    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)