recce-nightly 1.2.0.20250506__py3-none-any.whl → 1.26.0.20251124__py3-none-any.whl

recce/tasks/query.py

@@ -1,14 +1,14 @@
 import typing
-from typing import Optional, Tuple, List
+from typing import List, Optional, Tuple
 
 from pydantic import BaseModel
 
-from .core import Task, TaskResultDiffer, CheckValidator
-from .dataframe import DataFrame
-from .valuediff import ValueDiffMixin
 from ..core import default_context
 from ..exceptions import RecceException
 from ..models import Check
+from .core import CheckValidator, Task, TaskResultDiffer
+from .dataframe import DataFrame
+from .valuediff import ValueDiffMixin
 
 QUERY_LIMIT = 2000
 
@@ -19,11 +19,8 @@ if typing.TYPE_CHECKING:
 class QueryMixin:
     @classmethod
     def execute_sql_with_limit(
-        cls,
-        sql_template,
-        base: bool = False,
-        limit: Optional[int] = None
-    ) -> Tuple['agate.Table', bool]:
+        cls, sql_template, base: bool = False, limit: Optional[int] = None
+    ) -> Tuple["agate.Table", bool]:
         """
         Execute a SQL template and return the result as an agate table.
         :param sql_template: SQL template to execute
@@ -32,8 +29,10 @@ class QueryMixin:
         :return: Tuple of agate table and whether there are more rows to fetch
         """
         from jinja2.exceptions import TemplateSyntaxError
+
         dbt_adapter = default_context().adapter
         from dbt.exceptions import TargetNotFoundError
+
         try:
             sql = dbt_adapter.generate_sql(sql_template, base)
 
@@ -51,7 +50,7 @@ class QueryMixin:
             raise RecceException(f"Jinja template error: line {e.lineno}: {str(e)}")
 
     @classmethod
-    def execute_sql(cls, sql_template, base: bool = False) -> 'agate.Table':
+    def execute_sql(cls, sql_template, base: bool = False) -> "agate.Table":
         result, _ = cls.execute_sql_with_limit(sql_template, base)
         return result
 
@@ -87,6 +86,7 @@ class QueryTask(Task, QueryMixin):
 
     def execute_dbt(self):
         from recce.adapter.dbt_adapter import DbtAdapter
+
         dbt_adapter: DbtAdapter = default_context().adapter
 
         limit = QUERY_LIMIT
@@ -101,9 +101,10 @@ class QueryTask(Task, QueryMixin):
 
     def execute_sqlmesh(self):
         from ..adapter.sqlmesh_adapter import SqlmeshAdapter
+
         sqlmesh_adapter: SqlmeshAdapter = default_context().adapter
 
-        sql = self.params.get('sql_template')
+        sql = self.params.get("sql_template")
         limit = QUERY_LIMIT
         df, more = sqlmesh_adapter.fetchdf_with_limit(sql, base=self.is_base, limit=limit)
         return DataFrame.from_pandas(df, limit=limit, more=more)
@@ -111,7 +112,7 @@ class QueryTask(Task, QueryMixin):
     def execute(self):
         context = default_context()
 
-        if context.adapter_type == 'sqlmesh':
+        if context.adapter_type == "sqlmesh":
             return self.execute_sqlmesh()
         else:
             return self.execute_dbt()
@@ -139,8 +140,13 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
         self.connection = None
         self.legacy_surrogate_key = True
 
-    def _query_diff(self, dbt_adapter, sql_template: str, base_sql_template: Optional[str] = None,
-                    preview_change: bool = False):
+    def _query_diff(
+        self,
+        dbt_adapter,
+        sql_template: str,
+        base_sql_template: Optional[str] = None,
+        preview_change: bool = False,
+    ):
         limit = QUERY_LIMIT
 
         self.connection = dbt_adapter.get_thread_connection()
@@ -155,40 +161,76 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
 
         return QueryDiffResult(
             base=DataFrame.from_agate(base, limit=limit, more=base_more),
-            current=DataFrame.from_agate(current, limit=limit, more=current_more)
+            current=DataFrame.from_agate(current, limit=limit, more=current_more),
         )
 
-    def _query_diff_join(self, dbt_adapter, sql_template: str, primary_keys: List[str],
-                         base_sql_template: Optional[str] = None, preview_change: bool = False):
+    def _query_diff_join(
+        self,
+        dbt_adapter,
+        sql_template: str,
+        primary_keys: List[str],
+        base_sql_template: Optional[str] = None,
+        preview_change: bool = False,
+    ):
 
         query_template = r"""
-            {% set a_query %}
-                {{ base_query }}
-            {% endset %}
-
-            {% set b_query %}
-                {{ current_query }}
-            {% endset %}
-
-            {{ audit_helper.compare_queries(
-                a_query=a_query,
-                b_query=b_query,
-                primary_key=__PRIMARY_KEY__,
-                summarize=False,
-            ) }} limit {{ limit }}
-            """
-
-        if len(primary_keys) > 1:
-            self._verify_dbt_packages_deps(dbt_adapter)
-            self.check_cancel()
+        with a_query as (
+            {{ base_query }}
+        ),
+
+        b_query as (
+            {{ current_query }}
+        ),
+
+        a_intersect_b as (
+            select * from a_query
+            {{ dbt.intersect() }}
+            select * from b_query
+        ),
+
+        a_except_b as (
+            select * from a_query
+            {{ dbt.except() }}
+            select * from b_query
+        ),
+
+        b_except_a as (
+            select * from b_query
+            {{ dbt.except() }}
+            select * from a_query
+        ),
+
+        all_records as (
+            select
+                *,
+                true as in_a,
+                true as in_b
+            from a_intersect_b
+
+            union all
+
+            select
+                *,
+                true as in_a,
+                false as in_b
+            from a_except_b
+
+            union all
+
+            select
+                *,
+                false as in_a,
+                true as in_b
+            from b_except_a
+        )
 
-            if self.legacy_surrogate_key:
-                new_primary_key = 'dbt_utils.surrogate_key(primary_key)'
-            else:
-                new_primary_key = 'dbt_utils.generate_surrogate_key(primary_key)'
-        else:
-            new_primary_key = 'primary_key'
-        query_template = query_template.replace('__PRIMARY_KEY__', new_primary_key)
+        select * from all_records
+        where not (in_a and in_b)
+        order by {{ primary_keys | join(',\n') }}, in_a desc, in_b desc
+        limit {{ limit }}
+        """
+
+        self.check_cancel()
 
         if preview_change:
             base_query = dbt_adapter.generate_sql(base_sql_template, base=False)
@@ -196,19 +238,20 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
             base_query = dbt_adapter.generate_sql(base_sql_template or sql_template, base=True)
         current_query = dbt_adapter.generate_sql(sql_template, base=False)
 
-        sql = dbt_adapter.generate_sql(query_template, context=dict(
-            base_query=base_query,
-            current_query=current_query,
-            primary_key=primary_keys if len(primary_keys) != 1 else primary_keys[0],
-            limit=QUERY_LIMIT,
-        ))
+        sql = dbt_adapter.generate_sql(
+            query_template,
+            context=dict(
+                base_query=base_query,
+                current_query=current_query,
+                primary_keys=primary_keys,
+                limit=QUERY_LIMIT,
+            ),
+        )
 
         _, table = dbt_adapter.execute(sql, fetch=True)
         self.check_cancel()
 
-        return QueryDiffResult(
-            diff=DataFrame.from_agate(table)
-        )
+        return QueryDiffResult(diff=DataFrame.from_agate(table))
 
     @staticmethod
     def _select_single_model(model_name):
@@ -216,6 +259,7 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
 
     def execute_dbt(self):
         from recce.adapter.dbt_adapter import DbtAdapter
+
         dbt_adapter: DbtAdapter = default_context().adapter
 
         with dbt_adapter.connection_named("query"):
@@ -228,11 +272,20 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
                 preview_change = True
 
             if primary_keys:
-                return self._query_diff_join(dbt_adapter, sql_template, primary_keys,
-                                             base_sql_template=base_sql_template, preview_change=preview_change)
-
-            return self._query_diff(dbt_adapter, sql_template, base_sql_template=base_sql_template,
-                                    preview_change=preview_change)
+                return self._query_diff_join(
+                    dbt_adapter,
+                    sql_template,
+                    primary_keys,
+                    base_sql_template=base_sql_template,
+                    preview_change=preview_change,
+                )
+
+            return self._query_diff(
+                dbt_adapter,
+                sql_template,
+                base_sql_template=base_sql_template,
+                preview_change=preview_change,
+            )
 
     def _sqlmesh_query_diff(self, sql, base_sql=None):
         from ..adapter.sqlmesh_adapter import SqlmeshAdapter
@@ -244,7 +297,7 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
         curr, curr_more = sqlmesh_adapter.fetchdf_with_limit(sql, base=False, limit=limit)
         return QueryDiffResult(
             base=DataFrame.from_pandas(base, limit=limit, more=base_more),
-            current=DataFrame.from_pandas(curr, limit=limit, more=curr_more)
+            current=DataFrame.from_pandas(curr, limit=limit, more=curr_more),
         )
 
     def _sqlmesh_query_diff_join(self, sql, primary_keys, base_sql=None):
@@ -257,21 +310,18 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
         expr_curr = sqlmesh_adapter.replace_virtual_tables(sql, base=False)
         import sqlglot as g
 
-        expr = g.select(
-            '*',
-        ).with_(
-            'a', as_=expr_base
-        ).with_(
-            'b', as_=expr_curr
-        ).with_(
-            'a_interset_b', as_='select * from a intersect select * from b'
-        ).with_(
-            'a_except_b', as_='select * from a except select * from b'
-        ).with_(
-            'b_except_a', as_='select * from b except select * from a'
-        ).with_(
-            'all_records',
-            as_='''
+        expr = (
+            g.select(
+                "*",
+            )
+            .with_("a", as_=expr_base)
+            .with_("b", as_=expr_curr)
+            .with_("a_interset_b", as_="select * from a intersect select * from b")
+            .with_("a_except_b", as_="select * from a except select * from b")
+            .with_("b_except_a", as_="select * from b except select * from a")
+            .with_(
+                "all_records",
+                as_="""
             SELECT
               *,
               TRUE AS in_a,
@@ -289,19 +339,21 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
               FALSE AS in_a,
               TRUE AS in_b
             FROM b_except_a
-            '''
-        ).with_(
-            'final',
-            as_=f'''
+            """,
+            )
+            .with_(
+                "final",
+                as_=f"""
                     select * from all_records
                     where not (in_a and in_b)
                     order by {", ".join(primary_keys)}, in_a desc, in_b desc
-                    '''
-        ).from_('final').limit(1000)
-        diff, diff_more = sqlmesh_adapter.fetchdf_with_limit(expr, limit=limit)
-        return QueryDiffResult(
-            diff=DataFrame.from_pandas(diff, limit=limit, more=diff_more)
+                    """,
+            )
+            .from_("final")
+            .limit(1000)
         )
+        diff, diff_more = sqlmesh_adapter.fetchdf_with_limit(expr, limit=limit)
+        return QueryDiffResult(diff=DataFrame.from_pandas(diff, limit=limit, more=diff_more))
 
     def execute_sqlmesh(self):
         sql = self.params.sql_template
@@ -316,7 +368,7 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
     def execute(self):
         context = default_context()
 
-        if context.adapter_type == 'sqlmesh':
+        if context.adapter_type == "sqlmesh":
             return self.execute_sqlmesh()
         else:
             return self.execute_dbt()
@@ -329,14 +381,14 @@ class QueryDiffTask(Task, QueryMixin, ValueDiffMixin):
 
 class QueryDiffResultDiffer(TaskResultDiffer):
     def _check_result_changed_fn(self, result):
-        base = result.get('base')
-        current = result.get('current')
-        diff = result.get('diff')
+        base = result.get("base")
+        current = result.get("current")
+        diff = result.get("diff")
 
         if diff is None:
             return TaskResultDiffer.diff(base, current)
         else:
-            diff_data = diff.get('data')
+            diff_data = diff.get("data")
             if diff_data is None or len(diff_data) == 0:
                 return None
 

recce/tasks/rowcount.py

@@ -1,11 +1,11 @@
-from typing import Optional, Union, List, Literal
+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 TaskResultDiffer, CheckValidator
+from recce.tasks.core import CheckValidator, TaskResultDiffer
 from recce.tasks.query import QueryMixin
 
 
@@ -25,10 +25,10 @@ class RowCountTask(Task, QueryMixin):
         if node is None:
             return None
 
-        if node.resource_type != 'model' and node.resource_type != 'snapshot':
+        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']:
+        if node.config and node.config.materialized not in ["table", "view", "incremental", "snapshot"]:
             return None
 
         relation = dbt_adapter.create_relation(model_name, base=base)
@@ -54,8 +54,9 @@ class RowCountTask(Task, QueryMixin):
             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')
+                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,
@@ -80,7 +81,7 @@ class RowCountTask(Task, QueryMixin):
                 row_count = self._query_row_count(dbt_adapter, node, base=False)
                 self.check_cancel()
                 result[node] = {
-                    'curr': row_count,
+                    "curr": row_count,
                 }
                 completed += 1
 
@@ -98,7 +99,7 @@ class RowCountDiffParams(BaseModel):
     select: Optional[str] = None
     exclude: Optional[str] = None
     packages: Optional[list[str]] = None
-    view_mode: Optional[Literal['all', 'changed_models']] = None
+    view_mode: Optional[Literal["all", "changed_models"]] = None
 
 
 class RowCountDiffTask(Task, QueryMixin):
@@ -112,10 +113,10 @@ class RowCountDiffTask(Task, QueryMixin):
         if node is None:
             return None
 
-        if node.resource_type != 'model' and node.resource_type != 'snapshot':
+        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']:
+        if node.config and node.config.materialized not in ["table", "view", "incremental", "snapshot"]:
             return None
 
         relation = dbt_adapter.create_relation(model_name, base=base)
@@ -141,8 +142,9 @@ class RowCountDiffTask(Task, QueryMixin):
             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')
+                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,
@@ -169,8 +171,8 @@ class RowCountDiffTask(Task, QueryMixin):
                 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,
+                    "base": base_row_count,
+                    "curr": curr_row_count,
                 }
                 completed += 1
 
@@ -187,6 +189,7 @@ class RowCountDiffTask(Task, QueryMixin):
             query_candidates.append(node_name)
 
         from recce.adapter.sqlmesh_adapter import SqlmeshAdapter
+
         sqlmesh_adapter: SqlmeshAdapter = default_context().adapter
 
         for name in query_candidates:
@@ -194,28 +197,28 @@ class RowCountDiffTask(Task, QueryMixin):
             curr_row_count = None
 
             try:
-                df, _ = sqlmesh_adapter.fetchdf_with_limit(f'select count(*) from {name}', base=True)
+                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)
+                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,
+                "base": base_row_count,
+                "curr": curr_row_count,
             }
 
         return result
 
     def execute(self):
         context = default_context()
-        if context.adapter_type == 'dbt':
+        if context.adapter_type == "dbt":
             return self.execute_dbt()
         else:
             return self.execute_sqlmesh()
@@ -232,8 +235,8 @@ class RowCountDiffResultDiffer(TaskResultDiffer):
         current = {}
 
         for node, row_counts in result.items():
-            base[node] = row_counts['base']
-            current[node] = row_counts['curr']
+            base[node] = row_counts["base"]
+            current[node] = row_counts["curr"]
 
         return TaskResultDiffer.diff(base, current)
 
@@ -243,24 +246,27 @@ class RowCountDiffResultDiffer(TaskResultDiffer):
         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', [])
+        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', [])
+        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'),
+                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:
-            return self.changes.affected_root_keys.items
+            # 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):
@@ -268,4 +274,4 @@ class RowCountDiffCheckValidator(CheckValidator):
         try:
             RowCountDiffParams(**check.params)
         except Exception as e:
-            raise ValueError(f'Invalid params: str{e}')
+            raise ValueError(f"Invalid params: str{e}")

recce/tasks/schema.py

@@ -1,9 +1,9 @@
-from typing import Union, List, Optional, Literal
+from typing import List, Literal, Optional, Union
 
 from pydantic import BaseModel
 
 from recce.models import Check
-from recce.tasks.core import TaskResultDiffer, CheckValidator
+from recce.tasks.core import CheckValidator, TaskResultDiffer
 
 
 class SchemaDiffResultDiffer:
@@ -17,35 +17,38 @@ class SchemaDiffResultDiffer:
 
     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')]
+        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'),
+                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', {})
+        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', {})
+            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:
-            return self.changes.affected_root_keys.items
+            # 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):
@@ -53,7 +56,7 @@ class SchemaDiffParams(BaseModel):
     select: Optional[str] = None
     exclude: Optional[str] = None
     packages: Optional[list[str]] = None
-    view_mode: Optional[Literal['all', 'changed_models']] = None
+    view_mode: Optional[Literal["all", "changed_models"]] = None
 
 
 class SchemaDiffCheckValidator(CheckValidator):