@datarecce/ui 0.1.30 → 0.1.31

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

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

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

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

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