@datarecce/ui 0.1.41 → 0.2.0

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

@@ -1,258 +0,0 @@
-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

@@ -1,266 +0,0 @@
-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)