@datarecce/ui 0.1.30 → 0.1.31

package/recce-source/recce/util/cloud/checks.py

@@ -0,0 +1,242 @@
+"""
+Recce Cloud API client for check operations.
+
+This module provides methods for managing checks (validation operations) in Recce Cloud,
+including CRUD operations for checks within sessions.
+"""
+
+from typing import Dict, List
+
+from recce.util.cloud.base import CloudBase
+
+
+class ChecksCloud(CloudBase):
+    """
+    Client for Recce Cloud check operations.
+
+    Provides methods to list, create, retrieve, update, and delete checks
+    within Recce Cloud sessions.
+
+    Examples:
+        >>> client = ChecksCloud(token="your-api-token")
+        >>> checks = client.list_checks(org_id="org123", project_id="proj456", session_id="sess789")
+        >>> check = client.get_check(org_id="org123", project_id="proj456",
+        ...                          session_id="sess789", check_id="check001")
+    """
+
+    def list_checks(self, org_id: str, project_id: str, session_id: str) -> List[Dict]:
+        """
+        List all checks in a session.
+
+        Args:
+            org_id: Organization ID
+            project_id: Project ID
+            session_id: Session ID
+
+        Returns:
+            List of check dictionaries
+
+        Raises:
+            RecceCloudException: If the request fails
+
+        Example:
+            >>> checks = client.list_checks("org123", "proj456", "sess789")
+            >>> print(f"Found {len(checks)} checks")
+        """
+        api_url = f"{self.base_url_v2}/organizations/{org_id}/projects/{project_id}/sessions/{session_id}/checks"
+        response = self._request("GET", api_url)
+
+        self._raise_for_status(
+            response,
+            "Failed to list checks from Recce Cloud.",
+        )
+
+        data = response.json()
+        return data.get("checks", [])
+
+    def create_check(
+        self,
+        org_id: str,
+        project_id: str,
+        session_id: str,
+        check_data: Dict,
+    ) -> Dict:
+        """
+        Create a new check in a session.
+
+        Args:
+            org_id: Organization ID
+            project_id: Project ID
+            session_id: Session ID
+            check_data: Check data to create (should include check type, params, etc.)
+
+        Returns:
+            Created check dictionary
+
+        Raises:
+            RecceCloudException: If the request fails
+
+        Example:
+            >>> check_data = {
+            ...     "name": "Schema Check",
+            ...     "type": "schema_diff",
+            ...     "params": {"model": "customers"}
+            ... }
+            >>> check = client.create_check("org123", "proj456", "sess789", check_data)
+            >>> print(f"Created check with ID: {check['id']}")
+        """
+        api_url = f"{self.base_url_v2}/organizations/{org_id}/projects/{project_id}/sessions/{session_id}/checks"
+        response = self._request("POST", api_url, json=check_data)
+
+        self._raise_for_status(
+            response,
+            "Failed to create check in Recce Cloud.",
+        )
+        data = response.json()
+        return data.get("check")
+
+    def get_check(
+        self,
+        org_id: str,
+        project_id: str,
+        session_id: str,
+        check_id: str,
+    ) -> Dict:
+        """
+        Get a specific check by ID.
+
+        Args:
+            org_id: Organization ID
+            project_id: Project ID
+            session_id: Session ID
+            check_id: Check ID
+
+        Returns:
+            Check dictionary
+
+        Raises:
+            RecceCloudException: If the request fails or check not found
+
+        Example:
+            >>> check = client.get_check("org123", "proj456", "sess789", "check001")
+            >>> print(f"Check name: {check['name']}")
+            >>> print(f"Check status: {check['status']}")
+        """
+        api_url = (
+            f"{self.base_url_v2}/organizations/{org_id}/projects/{project_id}/sessions/{session_id}/checks/{check_id}"
+        )
+        response = self._request("GET", api_url)
+
+        self._raise_for_status(
+            response,
+            f"Failed to get check {check_id} from Recce Cloud.",
+        )
+
+        data = response.json()
+        return data.get("check", {})
+
+    def update_check(
+        self,
+        org_id: str,
+        project_id: str,
+        session_id: str,
+        check_id: str,
+        check_data: Dict,
+    ) -> Dict:
+        """
+        Update an existing check.
+
+        Args:
+            org_id: Organization ID
+            project_id: Project ID
+            session_id: Session ID
+            check_id: Check ID
+            check_data: Updated check data (partial updates supported)
+
+        Returns:
+            Updated check dictionary
+
+        Raises:
+            RecceCloudException: If the request fails
+
+        Example:
+            >>> update_data = {
+            ...     "status": "approved",
+            ...     "notes": "Validated successfully"
+            ... }
+            >>> check = client.update_check("org123", "proj456", "sess789", "check001", update_data)
+            >>> print(f"Updated check status: {check['status']}")
+        """
+        api_url = (
+            f"{self.base_url_v2}/organizations/{org_id}/projects/{project_id}/sessions/{session_id}/checks/{check_id}"
+        )
+        response = self._request("PATCH", api_url, json=check_data)
+
+        self._raise_for_status(
+            response,
+            f"Failed to update check {check_id} in Recce Cloud.",
+        )
+
+        data = response.json()
+        return data.get("check", {})
+
+    def delete_check(
+        self,
+        org_id: str,
+        project_id: str,
+        session_id: str,
+        check_id: str,
+    ) -> None:
+        """
+        Delete a check.
+
+        Args:
+            org_id: Organization ID
+            project_id: Project ID
+            session_id: Session ID
+            check_id: Check ID
+
+        Raises:
+            RecceCloudException: If the request fails
+
+        Example:
+            >>> client.delete_check("org123", "proj456", "sess789", "check001")
+            >>> print("Check deleted successfully")
+        """
+        api_url = (
+            f"{self.base_url_v2}/organizations/{org_id}/projects/{project_id}/sessions/{session_id}/checks/{check_id}"
+        )
+        response = self._request("DELETE", api_url)
+
+        # DELETE typically returns 204 No Content on success
+        if response.status_code not in (200, 204):
+            self._raise_for_status(
+                response,
+                f"Failed to delete check {check_id} from Recce Cloud.",
+            )
+
+    def create_preset_check(self, org_id: str, project_id: str, check_data: Dict):
+        """
+        Create a preset check from an existing check.
+
+        Args:
+            org_id: Organization ID
+            project_id: Project ID
+            check_data: Check data including name, description, type, params, view_options, and order_index
+
+        Returns:
+            Created preset check dictionary
+
+        Raises:
+            RecceCloudException: If the request fails
+        """
+        api_url = f"{self.base_url_v2}/organizations/{org_id}/projects/{project_id}/preset-checks"
+        response = self._request("POST", api_url, json=check_data)
+
+        self._raise_for_status(
+            response,
+            "Failed to create preset check in Recce Cloud.",
+        )
+
+        data = response.json()
+        return data.get("presetCheck", {})

package/recce-source/recce/util/io.py

@@ -0,0 +1,120 @@
+import gzip
+import os
+import tempfile
+from abc import ABC, ABCMeta, abstractmethod
+from enum import Enum
+
+
+class SupportedFileTypes(Enum):
+    FILE = "file"
+    GZIP = "gzip"
+    ZIP = "zip"
+
+
+def file_io_factory(file_type: SupportedFileTypes):
+    if file_type == SupportedFileTypes.FILE:
+        return FileIO
+    elif file_type == SupportedFileTypes.GZIP:
+        return GzipFileIO
+    elif file_type == SupportedFileTypes.ZIP:
+        return ZipFileIO
+    else:
+        raise ValueError(f"Unsupported file type: {file_type}")
+
+
+class AbstractFileIO(metaclass=ABCMeta):
+    @staticmethod
+    @abstractmethod
+    def write(path: str, data: str, **kwargs):
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def read(path: str, **kwargs) -> str:
+        raise NotImplementedError
+
+
+class FileIO(AbstractFileIO, ABC):
+    @staticmethod
+    def write(path: str, data: str, **kwargs):
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(data)
+
+    @staticmethod
+    def read(path: str, **kwargs) -> str:
+        with open(path, "r", encoding="utf-8") as f:
+            return f.read()
+
+
+class GzipFileIO(AbstractFileIO, ABC):
+    @staticmethod
+    def write(path: str, data: str, **kwargs):
+        with gzip.open(path, "wt") as f:
+            f.write(data)
+
+    @staticmethod
+    def read(path: str, **kwargs) -> str:
+        with gzip.open(path, "rt") as f:
+            return f.read()
+
+    @staticmethod
+    def read_fileobj(fileobj) -> bytes:
+        with gzip.GzipFile(fileobj=fileobj) as f:
+            return f.read()
+
+
+class ZipFileIO(AbstractFileIO, ABC):
+    @staticmethod
+    def _is_pyminizip_installed():
+        try:
+            import pyminizip
+
+            # Use the module to avoid F401
+            return pyminizip is not None
+        except ImportError:
+            raise ImportError("pyminizip is not installed. Please install it using `pip install pyminizip`")
+
+    @staticmethod
+    def read(path: str, **kwargs) -> str:
+        ZipFileIO._is_pyminizip_installed()
+        import pyminizip
+
+        cwd = os.getcwd()
+        password = kwargs.get("password")
+        zip_dir_name = kwargs.get("zip_dir_name")
+        if zip_dir_name is None:
+            raise ValueError("zip_dir_name is required for zipping")
+
+        try:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, zip_dir_name)
+                pyminizip.uncompress(path, password, tmp_dir, 0)
+                content = FileIO.read(tmp_file)
+        except Exception as e:
+            error_msg = str(e)
+            if "-3" in error_msg:
+                raise Exception("Invalid password to uncompress state file.")
+            raise Exception(f"Failed to uncompress state file: {error_msg}")
+        finally:
+            os.chdir(cwd)
+
+        return content
+
+    @staticmethod
+    def write(path: str, data: str, **kwargs):
+        ZipFileIO._is_pyminizip_installed()
+        import pyminizip
+
+        cwd = os.getcwd()
+        password = kwargs.get("password")
+        zip_dir_name = kwargs.get("zip_dir_name")
+        if zip_dir_name is None:
+            raise ValueError("zip_dir_name is required for zipping")
+
+        try:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, zip_dir_name)
+                FileIO.write(tmp_file, data)
+                pyminizip.compress(tmp_file, None, path, password, 9)
+        finally:
+            os.chdir(cwd)

package/recce-source/recce/util/lineage.py

@@ -0,0 +1,83 @@
+from typing import Dict, Iterable, Set, Tuple
+
+from recce.models.types import CllColumn, CllNode
+
+
+def find_upstream(node_ids: Iterable, parent_map):
+    visited = set()
+    upstream = set()
+
+    stack = list(node_ids)
+    while stack:
+        current = stack.pop()
+        if current in visited:
+            continue
+        visited.add(current)
+        parents = parent_map.get(current, [])
+        for parent in parents:
+            if parent not in upstream:
+                upstream.add(parent)
+                stack.append(parent)
+    return upstream
+
+
+def find_downstream(node_ids: Iterable, child_map):
+    visited = set()
+    downstream = set()
+
+    stack = list(node_ids)
+    while stack:
+        current = stack.pop()
+        if current in visited:
+            continue
+        visited.add(current)
+        children = child_map.get(current, [])
+        for child in children:
+            if child not in downstream:
+                downstream.add(child)
+                stack.append(child)
+    return downstream
+
+
+def find_column_dependencies(node_column_id: str, parent_map: Dict, child_map: Dict) -> Tuple[Set, Set]:
+    upstream_cols = find_upstream([node_column_id], parent_map)
+    downstream_cols = find_downstream([node_column_id], child_map)
+    return upstream_cols, downstream_cols
+
+
+def filter_lineage_vertices(
+    lineage_nodes: Dict[str, CllNode], lineage_columns: Dict[str, CllColumn], relevant_columns: Set[str]
+) -> Tuple[Dict[str, CllNode], Dict[str, CllColumn]]:
+    nodes = {}
+    columns = {}
+
+    for node_id, node in lineage_nodes.items():
+        if node_id in relevant_columns:
+            nodes[node_id] = node
+
+    for col_id, column in lineage_columns.items():
+        if col_id in relevant_columns:
+            columns[col_id] = column
+
+    return nodes, columns
+
+
+def filter_dependency_maps(
+    parent_map: Dict, child_map: Dict, relevant_ids: Set
+) -> Tuple[Dict[str, Set], Dict[str, Set]]:
+    p_map = {}
+    c_map = {}
+    for node_id, parents in parent_map.items():
+        if node_id in relevant_ids:
+            p_map[node_id] = {p for p in parents if p in relevant_ids}
+
+    for node_id, children in child_map.items():
+        if node_id in relevant_ids:
+            c_map[node_id] = {c for c in children if c in relevant_ids}
+
+    return p_map, c_map
+
+
+def build_column_key(node_id: str, column_name: str) -> str:
+    """Build a unique column key from node name and column name."""
+    return f"{node_id}_{column_name}"

package/recce-source/recce/util/logger.py

@@ -0,0 +1,25 @@
+import logging
+
+
+class CustomFormatter(logging.Formatter):
+    grey = "\x1b[38;20m"
+    blue = "\x1b[34;20m"
+    green = "\x1b[32;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format = "[%(levelname)s] %(name)s: %(message)s (%(filename)s:%(lineno)d)"
+
+    FORMATS = {
+        logging.DEBUG: blue + format + reset,
+        logging.INFO: green + format + reset,
+        logging.WARNING: yellow + format + reset,
+        logging.ERROR: red + format + reset,
+        logging.CRITICAL: bold_red + format + reset,
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)

package/recce-source/recce/util/onboarding_state.py

@@ -0,0 +1,45 @@
+from enum import Enum
+from typing import Union
+
+from recce.util.recce_cloud import (
+    get_recce_cloud_onboarding_state,
+    set_recce_cloud_onboarding_state,
+)
+
+
+class OnboardingState(Enum):
+    SIGNUP_SUCCESSFUL = "signup_successful"
+    LAUNCHED_WITH_TOKEN = "launched_with_token"
+    CONFIGURE_TWO_ENV = "configure_two_env"
+
+
+def update_onboarding_state(api_token: Union[str, bool, None], is_single_env: bool) -> OnboardingState:
+    if api_token:
+        # existing onboarding_state values -> new, launched, launched_with_two_envs
+        #   new -> launched
+        #   new -> launched_with_two_envs
+        #   launched -> launched_with_two_envs
+        cloud_onboarding_state = get_recce_cloud_onboarding_state(api_token)
+
+        if cloud_onboarding_state == OnboardingState.SIGNUP_SUCCESSFUL.value:
+            # User has an API Token and is a "new" user
+            if is_single_env:
+                # Mark the onboarding state as "launched" if the user is new
+                set_recce_cloud_onboarding_state(api_token, OnboardingState.LAUNCHED_WITH_TOKEN.value)
+                return OnboardingState.LAUNCHED_WITH_TOKEN
+            else:
+                set_recce_cloud_onboarding_state(api_token, OnboardingState.CONFIGURE_TWO_ENV.value)
+                return OnboardingState.CONFIGURE_TWO_ENV
+        elif cloud_onboarding_state == OnboardingState.LAUNCHED_WITH_TOKEN.value:
+            # User has an API Token and has Two Environments
+            if is_single_env:
+                # Just return the current state
+                return OnboardingState.LAUNCHED_WITH_TOKEN
+            else:
+                set_recce_cloud_onboarding_state(api_token, OnboardingState.CONFIGURE_TWO_ENV.value)
+                return OnboardingState.CONFIGURE_TWO_ENV
+        elif cloud_onboarding_state == OnboardingState.CONFIGURE_TWO_ENV.value:
+            # Just return the current state
+            return OnboardingState.CONFIGURE_TWO_ENV
+
+    return OnboardingState.SIGNUP_SUCCESSFUL

package/recce-source/recce/util/perf_tracking.py

@@ -0,0 +1,85 @@
+import time
+from dataclasses import dataclass
+
+
+@dataclass
+class LineagePerfTracker:
+    lineage_start = None
+    lineage_elapsed = None
+    column_lineage_start = None
+    column_lineage_elapsed = None
+
+    total_nodes = None
+    init_nodes = None
+    cll_nodes = 0
+    change_analysis_nodes = 0
+    anchor_nodes = None
+
+    params = None
+
+    def start_lineage(self):
+        self.lineage_start = time.perf_counter_ns()
+
+    def end_lineage(self):
+        if self.lineage_start is None:
+            return
+        self.lineage_elapsed = (time.perf_counter_ns() - self.lineage_start) / 1000000
+
+    def start_column_lineage(self):
+        self.column_lineage_start = time.perf_counter_ns()
+
+    def end_column_lineage(self):
+        if self.column_lineage_start is None:
+            return
+        self.column_lineage_elapsed = (time.perf_counter_ns() - self.column_lineage_start) / 1000000
+
+    def set_total_nodes(self, total_nodes):
+        self.total_nodes = total_nodes
+
+    def set_init_nodes(self, init_nodes):
+        self.init_nodes = init_nodes
+
+    def set_anchor_nodes(self, anchor_nodes):
+        self.anchor_nodes = anchor_nodes
+
+    def increment_cll_nodes(self):
+        self.cll_nodes += 1
+
+    def increment_change_analysis_nodes(self):
+        self.change_analysis_nodes += 1
+
+    def set_params(self, has_node, has_column, change_analysis, no_cll, no_upstream, no_downstream):
+        self.params = {
+            "has_node": has_node,
+            "has_column": has_column,
+            "change_analysis": change_analysis,
+            "no_cll": no_cll,
+            "no_upstream": no_upstream,
+            "no_downstream": no_downstream,
+        }
+
+    def to_dict(self):
+        return {
+            "lineage_elapsed_ms": self.lineage_elapsed,
+            "column_lineage_elapsed_ms": self.column_lineage_elapsed,
+            "total_nodes": self.total_nodes,
+            "init_nodes": self.init_nodes,
+            "cll_nodes": self.cll_nodes,
+            "change_analysis_nodes": self.change_analysis_nodes,
+            "anchor_nodes": self.anchor_nodes,
+            "params": self.params,
+        }
+
+    def reset(self):
+        self.lineage_start = None
+        self.lineage_elapsed = None
+        self.column_lineage_start = None
+        self.column_lineage_elapsed = None
+
+        self.total_nodes = None
+        self.init_nodes = None
+        self.change_analysis_nodes = 0
+        self.cll_nodes = 0
+        self.anchor_nodes = 0
+
+        self.params = None

package/recce-source/recce/util/pydantic_model.py

@@ -0,0 +1,22 @@
+import pydantic
+from pydantic import BaseModel
+
+
+def pydantic_model_json_dump(model: BaseModel):
+    pydantic_version = pydantic.version.VERSION
+    pydantic_major = pydantic_version.split(".")[0]
+
+    if pydantic_major == "1":
+        return model.json(exclude_none=True)
+    else:
+        return model.model_dump_json(exclude_none=True)
+
+
+def pydantic_model_dump(model: BaseModel):
+    pydantic_version = pydantic.version.VERSION
+    pydantic_major = pydantic_version.split(".")[0]
+
+    if pydantic_major == "1":
+        return model.dict()
+    else:
+        return model.model_dump()