@datarecce/ui 0.1.40 → 0.2.0

package/recce-source/recce/tasks/rowcount.py

@@ -1,277 +0,0 @@
-from typing import List, Literal, Optional, Union
-
-from pydantic import BaseModel
-
-from recce.core import default_context
-from recce.models import Check
-from recce.tasks import Task
-from recce.tasks.core import CheckValidator, TaskResultDiffer
-from recce.tasks.query import QueryMixin
-
-
-class RowCountParams(BaseModel):
-    node_names: Optional[list[str]] = None
-    node_ids: Optional[list[str]] = None
-
-
-class RowCountTask(Task, QueryMixin):
-    def __init__(self, params: dict):
-        super().__init__()
-        self.params = RowCountParams(**params) if params is not None else RowCountParams()
-        self.connection = None
-
-    def _query_row_count(self, dbt_adapter, model_name, base=False):
-        node = dbt_adapter.find_node_by_name(model_name, base=base)
-        if node is None:
-            return None
-
-        if node.resource_type != "model" and node.resource_type != "snapshot":
-            return None
-
-        if node.config and node.config.materialized not in ["table", "view", "incremental", "snapshot"]:
-            return None
-
-        relation = dbt_adapter.create_relation(model_name, base=base)
-        if relation is None:
-            return None
-
-        sql_template = r"select count(*) from {{ relation }}"
-        sql = dbt_adapter.generate_sql(sql_template, context=dict(relation=relation))
-        _, table = dbt_adapter.execute(sql, fetch=True)
-        return int(table[0][0]) if table[0][0] is not None else 0
-
-    def execute(self):
-        result = {}
-
-        dbt_adapter = default_context().adapter
-
-        query_candidates = []
-        if self.params.node_ids or self.params.node_names:
-            for node_id in self.params.node_ids or []:
-                name = dbt_adapter.get_node_name_by_id(node_id)
-                if name:
-                    query_candidates.append(name)
-            for node in self.params.node_names or []:
-                query_candidates.append(node)
-        else:
-
-            def countable(unique_id):
-                return unique_id.startswith("model") or unique_id.startswith("snapshot") or unique_id.startswith("seed")
-
-            node_ids = dbt_adapter.select_nodes(
-                select=self.params.select,
-                exclude=self.params.exclude,
-                packages=self.params.packages,
-                view_mode=self.params.view_mode,
-            )
-            node_ids = list(filter(countable, node_ids))
-            for node_id in node_ids:
-                name = dbt_adapter.get_node_name_by_id(node_id)
-                if name:
-                    query_candidates.append(name)
-
-        # Query row count for nodes that are not cached
-        with dbt_adapter.connection_named("query"):
-            self.connection = dbt_adapter.get_thread_connection()
-            completed = 0
-            total = len(query_candidates)
-            for node in query_candidates:
-                self.update_progress(message=f"Query: {node} [{completed}/{total}]", percentage=completed / total)
-
-                row_count = self._query_row_count(dbt_adapter, node, base=False)
-                self.check_cancel()
-                result[node] = {
-                    "curr": row_count,
-                }
-                completed += 1
-
-        return result
-
-    def cancel(self):
-        super().cancel()
-        if self.connection:
-            self.close_connection(self.connection)
-
-
-class RowCountDiffParams(BaseModel):
-    node_names: Optional[list[str]] = None
-    node_ids: Optional[list[str]] = None
-    select: Optional[str] = None
-    exclude: Optional[str] = None
-    packages: Optional[list[str]] = None
-    view_mode: Optional[Literal["all", "changed_models"]] = None
-
-
-class RowCountDiffTask(Task, QueryMixin):
-    def __init__(self, params: dict):
-        super().__init__()
-        self.params = RowCountDiffParams(**params) if params is not None else RowCountDiffParams()
-        self.connection = None
-
-    def _query_row_count(self, dbt_adapter, model_name, base=False):
-        node = dbt_adapter.find_node_by_name(model_name, base=base)
-        if node is None:
-            return None
-
-        if node.resource_type != "model" and node.resource_type != "snapshot":
-            return None
-
-        if node.config and node.config.materialized not in ["table", "view", "incremental", "snapshot"]:
-            return None
-
-        relation = dbt_adapter.create_relation(model_name, base=base)
-        if relation is None:
-            return None
-
-        sql_template = r"select count(*) from {{ relation }}"
-        sql = dbt_adapter.generate_sql(sql_template, context=dict(relation=relation))
-        _, table = dbt_adapter.execute(sql, fetch=True)
-        return int(table[0][0]) if table[0][0] is not None else 0
-
-    def execute_dbt(self):
-        result = {}
-
-        dbt_adapter = default_context().adapter
-
-        query_candidates = []
-        if self.params.node_ids or self.params.node_names:
-            for node_id in self.params.node_ids or []:
-                name = dbt_adapter.get_node_name_by_id(node_id)
-                if name:
-                    query_candidates.append(name)
-            for node in self.params.node_names or []:
-                query_candidates.append(node)
-        else:
-
-            def countable(unique_id):
-                return unique_id.startswith("model") or unique_id.startswith("snapshot") or unique_id.startswith("seed")
-
-            node_ids = dbt_adapter.select_nodes(
-                select=self.params.select,
-                exclude=self.params.exclude,
-                packages=self.params.packages,
-                view_mode=self.params.view_mode,
-            )
-            node_ids = list(filter(countable, node_ids))
-            for node_id in node_ids:
-                name = dbt_adapter.get_node_name_by_id(node_id)
-                if name:
-                    query_candidates.append(name)
-
-        # Query row count for nodes that are not cached
-        with dbt_adapter.connection_named("query"):
-            self.connection = dbt_adapter.get_thread_connection()
-            completed = 0
-            total = len(query_candidates)
-            for node in query_candidates:
-                self.update_progress(message=f"Diff: {node} [{completed}/{total}]", percentage=completed / total)
-
-                base_row_count = self._query_row_count(dbt_adapter, node, base=True)
-                self.check_cancel()
-                curr_row_count = self._query_row_count(dbt_adapter, node, base=False)
-                self.check_cancel()
-                result[node] = {
-                    "base": base_row_count,
-                    "curr": curr_row_count,
-                }
-                completed += 1
-
-        return result
-
-    def execute_sqlmesh(self):
-        result = {}
-
-        query_candidates = []
-
-        for node_id in self.node_ids or []:
-            query_candidates.append(node_id)
-        for node_name in self.params.node_names or []:
-            query_candidates.append(node_name)
-
-        from recce.adapter.sqlmesh_adapter import SqlmeshAdapter
-
-        sqlmesh_adapter: SqlmeshAdapter = default_context().adapter
-
-        for name in query_candidates:
-            base_row_count = None
-            curr_row_count = None
-
-            try:
-                df, _ = sqlmesh_adapter.fetchdf_with_limit(f"select count(*) from {name}", base=True)
-                base_row_count = int(df.iloc[0, 0])
-            except Exception:
-                pass
-            self.check_cancel()
-
-            try:
-                df, _ = sqlmesh_adapter.fetchdf_with_limit(f"select count(*) from {name}", base=False)
-                curr_row_count = int(df.iloc[0, 0])
-            except Exception:
-                pass
-            self.check_cancel()
-            result[name] = {
-                "base": base_row_count,
-                "curr": curr_row_count,
-            }
-
-        return result
-
-    def execute(self):
-        context = default_context()
-        if context.adapter_type == "dbt":
-            return self.execute_dbt()
-        else:
-            return self.execute_sqlmesh()
-
-    def cancel(self):
-        super().cancel()
-        if self.connection:
-            self.close_connection(self.connection)
-
-
-class RowCountDiffResultDiffer(TaskResultDiffer):
-    def _check_result_changed_fn(self, result):
-        base = {}
-        current = {}
-
-        for node, row_counts in result.items():
-            base[node] = row_counts["base"]
-            current[node] = row_counts["curr"]
-
-        return TaskResultDiffer.diff(base, current)
-
-    def _get_related_node_ids(self) -> Union[List[str], None]:
-        """
-        Get the related node ids.
-        Should be implemented by subclass.
-        """
-        params = self.run.params
-        if params.get("model"):
-            return [TaskResultDiffer.get_node_id_by_name(params.get("model"))]
-        elif params.get("node_names"):
-            names = params.get("node_names", [])
-            return [TaskResultDiffer.get_node_id_by_name(name) for name in names]
-        elif params.get("node_ids"):
-            return params.get("node_ids", [])
-        else:
-            return TaskResultDiffer.get_node_ids_by_selector(
-                select=params.get("select"),
-                exclude=params.get("exclude"),
-                packages=params.get("packages"),
-                view_mode=params.get("view_mode"),
-            )
-
-    def _get_changed_nodes(self) -> Union[List[str], None]:
-        if self.changes:
-            # Both affected_root_keys of deepdiff v7 (OrderedSet) and v8 (SetOrdered) are iterable
-            # Convert to list directly
-            return list(self.changes.affected_root_keys)
-        return None
-
-
-class RowCountDiffCheckValidator(CheckValidator):
-    def validate_check(self, check: Check):
-        try:
-            RowCountDiffParams(**check.params)
-        except Exception as e:
-            raise ValueError(f"Invalid params: str{e}")

package/recce-source/recce/tasks/schema.py

@@ -1,65 +0,0 @@
-from typing import List, Literal, Optional, Union
-
-from pydantic import BaseModel
-
-from recce.models import Check
-from recce.tasks.core import CheckValidator, TaskResultDiffer
-
-
-class SchemaDiffResultDiffer:
-    related_node_ids: List[str] = None
-
-    def __init__(self, check, base_lineage, curr_lineage):
-        self.check = check
-        self.related_node_ids = self._get_related_node_ids()
-        self.changes = self._check_result_changed_fn(base_lineage, curr_lineage)
-        self.changed_nodes = self._get_changed_nodes()
-
-    def _get_related_node_ids(self) -> Union[List[str], None]:
-        params = self.check.params
-        if params.get("node_id"):
-            return params.get("node_id") if isinstance(params.get("node_id"), list) else [params.get("node_id")]
-        else:
-            return TaskResultDiffer.get_node_ids_by_selector(
-                select=params.get("select"),
-                exclude=params.get("exclude"),
-                packages=params.get("packages"),
-                view_mode=params.get("view_mode"),
-            )
-
-    def _check_result_changed_fn(self, base_lineage, curr_lineage):
-        base = {}
-        current = {}
-        base_nodes = base_lineage.get("nodes", {})
-        curr_nodes = curr_lineage.get("nodes", {})
-        for node_id in self.related_node_ids:
-            node = curr_nodes.get(node_id) or base_nodes.get(node_id)
-            if not node:
-                continue
-
-            node_name = node.get("name")
-            base[node_name] = base_nodes.get(node_id, {}).get("columns", {})
-            current[node_name] = curr_nodes.get(node_id, {}).get("columns", {})
-
-        return TaskResultDiffer.diff(base, current)
-
-    def _get_changed_nodes(self) -> Union[List[str], None]:
-        if self.changes:
-            # Both affected_root_keys of deepdiff v7 (OrderedSet) and v8 (SetOrdered) are iterable
-            # Convert to list directly
-            return list(self.changes.affected_root_keys)
-        return None
-
-
-class SchemaDiffParams(BaseModel):
-    node_id: Optional[Union[str, List[str]]] = None
-    select: Optional[str] = None
-    exclude: Optional[str] = None
-    packages: Optional[list[str]] = None
-    view_mode: Optional[Literal["all", "changed_models"]] = None
-
-
-class SchemaDiffCheckValidator(CheckValidator):
-    def validate_check(self, check: Check):
-        SchemaDiffParams(**check.params)
-        SchemaDiffParams(**check.view_options)

package/recce-source/recce/tasks/top_k.py

@@ -1,172 +0,0 @@
-from typing import Optional
-
-from pydantic import BaseModel
-
-from recce.core import default_context
-from recce.models import Check
-from recce.tasks import Task
-from recce.tasks.core import CheckValidator, TaskResultDiffer
-from recce.tasks.query import QueryMixin
-
-
-class TopKDiffParams(BaseModel):
-    model: str
-    column_name: str
-    k: Optional[int] = 10
-
-
-class TopKDiffTask(Task, QueryMixin):
-    def __init__(self, params):
-        super().__init__()
-        self.params = TopKDiffParams(**params)
-        self.connection = None
-
-    def _query_row_count_diff(self, dbt_adapter, base_relation, curr_relation, column):
-        """
-        Query the row count of the base and current relations
-
-        :return: [base_total, base_valids, curr_total, curr_valids]
-        """
-
-        sql_template = r"""
-        select count(*), count({{column}}) from {{ base_relation }}
-        UNION ALL
-        select count(*), count({{column}}) from {{ curr_relation }}
-        """
-        sql = dbt_adapter.generate_sql(
-            sql_template,
-            context=dict(
-                base_relation=base_relation,
-                curr_relation=curr_relation,
-                column=column,
-            ),
-        )
-        _, table = dbt_adapter.execute(sql, fetch=True)
-
-        result = (table[0][0], table[0][1], table[1][0], table[1][1])
-
-        return (int(v) if v is not None else 0 for v in result)
-
-    def _query_top_k(self, dbt_adapter, base_relation, curr_relation, column, k):
-        sql_template = r"""
-        WITH
-        BASE_CAT as (
-            select
-                coalesce(cast({{column}} as {{ dbt.type_string() }}), '__null__') as category,
-                count(*) as c
-            from {{base_relation}}
-            {% if not include_null %}
-            where {{column}} is not null
-            {% endif %}
-            group by 1
-        ),
-        CURR_CAT as (
-            select
-                coalesce(cast({{column}} as {{ dbt.type_string() }}), '__null__') as category,
-                count(*) as c
-            from {{curr_relation}}
-            {% if not include_null %}
-            where {{column}} is not null
-            {% endif %}
-            group by 1
-        )
-        select
-            coalesce(CURR_CAT.category, BASE_CAT.category) as category,
-            coalesce(BASE_CAT.c, 0) as base_count,
-            coalesce(CURR_CAT.c, 0) as curr_count
-        from CURR_CAT
-        full outer join BASE_CAT
-        on CURR_CAT.category = BASE_CAT.category
-        order by curr_count desc, base_count desc
-        limit {{k}}
-        """
-        sql = dbt_adapter.generate_sql(
-            sql_template,
-            context=dict(
-                base_relation=base_relation,
-                curr_relation=curr_relation,
-                column=column,
-                k=k,
-                include_null=False,
-            ),
-        )
-        _, table = dbt_adapter.execute(sql, fetch=True)
-
-        categories = []
-        base_counts = []
-        curr_counts = []
-
-        for row in table:
-            categories.append(row[0] if row[0] != "__null__" else None)
-            base_counts.append(int(row[1] if row[1] else 0))
-            curr_counts.append(int(row[2] if row[2] else 0))
-
-        return categories, base_counts, curr_counts
-
-    def execute(self):
-
-        from recce.adapter.dbt_adapter import DbtAdapter
-
-        dbt_adapter: DbtAdapter = default_context().adapter
-
-        with dbt_adapter.connection_named("query"):
-            self.connection = dbt_adapter.get_thread_connection()
-            model = self.params.model
-            column = self.params.column_name
-            k = self.params.k or 10
-
-            base_relation = dbt_adapter.create_relation(model, base=True)
-            if base_relation is None:
-                raise ValueError(f"Model '{model}' not found in the manifest")
-
-            curr_relation = dbt_adapter.create_relation(model, base=False)
-            if curr_relation is None:
-                raise ValueError(f"Model '{model}' not found in the manifest")
-
-            self.check_cancel()
-            categories, base_counts, curr_counts = self._query_top_k(
-                dbt_adapter, base_relation, curr_relation, column, k
-            )
-            self.check_cancel()
-
-            base_total, base_valids, curr_total, curr_valids = self._query_row_count_diff(
-                dbt_adapter, base_relation, curr_relation, column
-            )
-
-            result = {
-                "base": {
-                    "values": categories,
-                    "counts": base_counts,
-                    "valids": base_valids,
-                    "total": base_total,
-                },
-                "current": {
-                    "values": categories,
-                    "counts": curr_counts,
-                    "valids": curr_valids,
-                    "total": curr_total,
-                },
-            }
-            return result
-
-    def cancel(self):
-        super().cancel()
-        if self.connection:
-            self.close_connection(self.connection)
-
-
-class TopKDiffTaskResultDiffer(TaskResultDiffer):
-    def _check_result_changed_fn(self, result):
-        base = result.get("base")
-        current = result.get("current")
-
-        return TaskResultDiffer.diff(base, current)
-
-
-class TopKDiffCheckValidator(CheckValidator):
-
-    def validate_check(self, check: Check):
-        try:
-            TopKDiffParams(**check.params)
-        except Exception as e:
-            raise ValueError(f"Invalid check: {str(e)}")

package/recce-source/recce/tasks/utils.py

@@ -1,147 +0,0 @@
-"""Utility functions for task operations."""
-
-from typing import List, Optional
-
-from recce.tasks.dataframe import DataFrame
-
-
-def strip_identifier_quotes(identifier: str) -> str:
-    """
-    Strip SQL identifier quotes from a column name.
-
-    Different databases use different quoting styles:
-    - Double quotes: "column" (PostgreSQL, Snowflake, etc.)
-    - Backticks: `column` (MySQL, BigQuery)
-    - Square brackets: [column] (SQL Server)
-
-    Args:
-        identifier: Column name that may be quoted
-
-    Returns:
-        Column name with quotes stripped
-
-    Examples:
-        >>> strip_identifier_quotes('"myColumn"')
-        'myColumn'
-        >>> strip_identifier_quotes('`my_column`')
-        'my_column'
-        >>> strip_identifier_quotes('[Column Name]')
-        'Column Name'
-        >>> strip_identifier_quotes('regular_column')
-        'regular_column'
-    """
-    if not identifier or len(identifier) < 2:
-        return identifier
-
-    # Check for double quotes
-    if identifier.startswith('"') and identifier.endswith('"'):
-        return identifier[1:-1]
-
-    # Check for backticks
-    if identifier.startswith("`") and identifier.endswith("`"):
-        return identifier[1:-1]
-
-    # Check for square brackets
-    if identifier.startswith("[") and identifier.endswith("]"):
-        return identifier[1:-1]
-
-    return identifier
-
-
-def normalize_keys_to_columns(
-    keys: Optional[List[str]],
-    column_keys: List[str],
-) -> Optional[List[str]]:
-    """
-    Normalize user-provided keys to match actual column keys from the warehouse.
-
-    Different warehouses return column names in different cases:
-    - Snowflake: UPPERCASE (unless quoted)
-    - PostgreSQL/Redshift: lowercase (unless quoted)
-    - BigQuery: preserves original case
-
-    This function first attempts an exact match (for quoted columns that preserve
-    case), then falls back to case-insensitive matching to align user input
-    with the actual column keys returned by the warehouse.
-
-    Args:
-        keys: User-provided keys (e.g., primary_keys from params)
-        column_keys: Actual column keys from the query result
-
-    Returns:
-        List of keys normalized to match column_keys casing,
-        or None if keys is None.
-        If a key doesn't match any column, it's preserved as-is.
-
-    Examples:
-        >>> normalize_keys_to_columns(["payment_id"], ["PAYMENT_ID", "ORDER_ID"])
-        ["PAYMENT_ID"]
-
-        >>> normalize_keys_to_columns(["ID", "NAME"], ["id", "name", "value"])
-        ["id", "name"]
-
-        >>> normalize_keys_to_columns(["preCommitID"], ["preCommitID", "order_id"])
-        ["preCommitID"]  # Exact match preserved for quoted columns
-
-        >>> normalize_keys_to_columns(['"customerID"'], ["customerID", "amount"])
-        ["customerID"]  # Quotes stripped, then matched
-
-        >>> normalize_keys_to_columns(['`my_column`'], ["MY_COLUMN"])
-        ["MY_COLUMN"]  # Backticks stripped, then case-insensitive match
-    """
-    if keys is None:
-        return None
-
-    # Strip quotes from all keys first - quotes are for SQL execution,
-    # but the frontend should receive unquoted column names
-    unquoted_keys = [strip_identifier_quotes(key) for key in keys]
-
-    if not column_keys:
-        return unquoted_keys
-
-    # Build both exact and case-insensitive lookup maps
-    exact_key_set = set(column_keys)
-    case_insensitive_map = {col.lower(): col for col in column_keys}
-
-    normalized = []
-    for key in unquoted_keys:
-        if key in exact_key_set:
-            # Exact match - use as-is (handles quoted columns that preserved case)
-            normalized.append(key)
-        else:
-            # Case-insensitive fallback
-            actual_key = case_insensitive_map.get(key.lower())
-            normalized.append(actual_key if actual_key is not None else key)
-
-    return normalized
-
-
-def normalize_boolean_flag_columns(df: "DataFrame") -> "DataFrame":
-    """
-    Normalize boolean flag columns (in_a, in_b) to lowercase for cross-warehouse consistency.
-
-    Different warehouses return column names in different cases:
-    - Snowflake: IN_A, IN_B (UPPERCASE)
-    - PostgreSQL/Redshift: in_a, in_b (lowercase)
-    - BigQuery: preserves original case
-
-    This function ensures these columns are always lowercase in the DataFrame
-    sent to the frontend, enabling exact string matching.
-
-    Args:
-        df: DataFrame that may contain IN_A/IN_B columns
-
-    Returns:
-        DataFrame with in_a/in_b columns normalized to lowercase
-    """
-    from .dataframe import DataFrame, DataFrameColumn
-
-    normalized_columns = []
-    for col in df.columns:
-        key_upper = col.key.upper() if col.key else ""
-        if key_upper in ("IN_A", "IN_B"):
-            normalized_columns.append(DataFrameColumn(key=col.key.lower(), name=col.name.lower(), type=col.type))
-        else:
-            normalized_columns.append(col)
-
-    return DataFrame(columns=normalized_columns, data=df.data, limit=df.limit, more=df.more)