add json format for dbt diff values

Author: vvkh

data_diff/__main__.py

@@ -279,6 +279,7 @@ def main(conf, run, **kw):
                 project_dir_override=project_dir_override,
                 is_cloud=kw["cloud"],
                 dbt_selection=kw["select"],
+                json_output=kw["json_output"]
             )
         else:
             return _data_diff(dbt_project_dir=project_dir_override,

data_diff/dbt.py

@@ -1,3 +1,4 @@
+import json
 import os
 import time
 import webbrowser
@@ -11,6 +12,8 @@
 from . import connect_to_table, diff_tables, Algorithm
 from .cloud import DatafoldAPI, TCloudApiDataDiff, TCloudApiOrgMeta, get_or_create_data_source
 from .dbt_parser import DbtParser, PROJECT_FILE
+from .diff_tables import DiffResultWrapper
+from .format import jsonify
 from .tracking import (
     set_entrypoint_name,
     set_dbt_user_id,
@@ -52,6 +55,7 @@ def dbt_diff(
     project_dir_override: Optional[str] = None,
     is_cloud: bool = False,
     dbt_selection: Optional[str] = None,
+    json_output: bool = False,
 ) -> None:
     print_version_info()
     diff_threads = []
@@ -113,7 +117,7 @@ def dbt_diff(
                 diff_thread = run_as_daemon(_cloud_diff, diff_vars, datasource_id, api, org_meta)
                 diff_threads.append(diff_thread)
             else:
-                _local_diff(diff_vars)
+                _local_diff(diff_vars, json_output)
         else:
             rich.print(
                 _diff_output_base(".".join(diff_vars.dev_path), ".".join(diff_vars.prod_path))
@@ -186,7 +190,7 @@ def _get_diff_vars(
     )
 
 
-def _local_diff(diff_vars: TDiffVars) -> None:
+def _local_diff(diff_vars: TDiffVars, json_output: bool = False) -> None:
     dev_qualified_str = ".".join(diff_vars.dev_path)
     prod_qualified_str = ".".join(diff_vars.prod_path)
     diff_output_str = _diff_output_base(dev_qualified_str, prod_qualified_str)
@@ -236,7 +240,7 @@ def _local_diff(diff_vars: TDiffVars) -> None:
 
     extra_columns = tuple(column_set)
 
-    diff = diff_tables(
+    diff: DiffResultWrapper = diff_tables(
         table1,
         table2,
         threaded=True,
@@ -245,6 +249,16 @@ def _local_diff(diff_vars: TDiffVars) -> None:
         where=diff_vars.where_filter,
         skip_null_keys=True,
     )
+    if json_output:
+        # drain the iterator to get accumulated stats in diff.info_tree
+        list(diff)
+
+        print(json.dumps(jsonify(diff, with_summary=True, with_columns={
+            "added": columns_added,
+            "removed": columns_removed,
+            "changed": columns_type_changed,
+        })))
+        return
 
     if list(diff):
         diff_output_str += f"{diff.get_stats_string(is_dbt=True)} \n"

data_diff/format.py

@@ -0,0 +1,236 @@
+import collections
+
+from data_diff.diff_tables import DiffResultWrapper
+from typing import TypedDict, Any, Optional
+
+
+
+class ColumnsDiff(TypedDict):
+    removed: list[str]
+    added: list[str]
+    changed: list[str]
+
+
+def jsonify(diff: DiffResultWrapper,
+            with_summary: bool = False,
+            with_columns: ColumnsDiff | None = None) -> 'JsonDiff':
+    """
+    Converts the diff result into a JSON-serializable format.
+    Optionally add stats summary and schema diff.
+    """
+    diff_info = diff.info_tree.info
+    table1 = diff_info.tables[0]
+    table2 = diff_info.tables[1]
+    key_columns = table1.key_columns
+
+    t1_exclusive_rows = []
+    t2_exclusive_rows = []
+    diff_rows = []
+    schema = [field for field, _ in diff_info.diff_schema]
+
+    t1_exclusive_rows, t2_exclusive_rows, diff_rows = _group_rows(diff_info, schema)
+
+
+    diff_rows_jsonified = []
+    for row in diff_rows:
+        diff_rows_jsonified.append(_jsonify_diff(row, key_columns))
+
+    t1_exclusive_rows_jsonified = []
+    for row in t1_exclusive_rows:
+        t1_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
+
+    t2_exclusive_rows_jsonified = []
+    for row in t2_exclusive_rows:
+        t2_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
+    
+    summary = None
+    if with_summary:
+        summary = _jsonify_diff_summary(diff.get_stats_dict())
+    
+    columns = None
+    if with_columns:
+        added, removed, changed = with_columns['added'], with_columns['removed'], with_columns['changed']
+        columns = _jsonify_columns_diff(added, removed, changed)
+
+    is_different = bool(
+        t1_exclusive_rows
+        or t2_exclusive_rows
+        or diff_rows
+        or with_columns and (
+            with_columns['added']
+            or with_columns['removed']
+            or with_columns['changed']
+        )
+    )
+    return {
+        'isDifferent': is_different,
+        'table1': list(table1.table_path),
+        'table2': list(table2.table_path),
+        'rows': {
+            'exclusive': {
+                'table1': t1_exclusive_rows_jsonified,
+                'table2': t2_exclusive_rows_jsonified,
+            },
+            'diff': diff_rows_jsonified,
+        },
+        'summary': summary,
+        'columns': columns,
+    }
+
+class JsonDiff(TypedDict):
+    table1: list[str]
+    table2: list[str]
+    rows: TypedDict('Rows', {
+        'exclusive': TypedDict('Exclusive', {
+            'table1': list['JsonExclusiveRow'],
+            'table2': list['JsonExclusiveRow'],
+        }),
+        'diff': list['JsonDiffRow'],
+    })
+    summary: Optional['JsonDiffSummary' ]
+    columns: Optional['JsonColumnsSummary']
+
+
+
+class JsonExclusiveRowValue(TypedDict):
+    """
+    Value of a single column in a row
+    """
+    isPK: bool
+    value: Any
+
+class JsonDiffRowValue(TypedDict):
+    """
+    Pair of diffed values for 2 rows with equal PKs
+    """
+    table1: Any
+    table2: Any
+    isDiff: bool
+    isPK: bool
+
+
+JsonDiffRow = dict[str, JsonDiffRowValue]
+JsonExclusiveRow = dict[str, JsonExclusiveRowValue]
+
+
+class JsonDiffSummary(TypedDict):
+    rows: TypedDict('Rows', {
+        'total': TypedDict('Total', {
+            'table1': int,
+            'table2': int,
+        }),
+        'exclusive': TypedDict('Exclusive', {
+            'table1': int,
+            'table2': int,
+        }),
+        'updated': int,
+        'unchanged': int,
+    })
+    stats: TypedDict('Stats', {
+        'diffCounts': dict[str, int],
+    })
+
+class JsonColumnsSummary(TypedDict):
+    exclusive: TypedDict('Exclusive', {
+        'table1': list[str],
+        'table2': list[str],
+    })
+    typeChanged: list[str]
+
+
+
+def _group_rows(diff_info: DiffResultWrapper, 
+                schema: list[str]) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]:
+    t1_exclusive_rows = []
+    t2_exclusive_rows = []
+    diff_rows = []
+
+    for row in diff_info.diff:
+        row_w_schema = dict(zip(schema, row))
+        is_t1_exclusive = row_w_schema['is_exclusive_a']
+        is_t2_exclusive = row_w_schema['is_exclusive_b']
+
+        if is_t1_exclusive:
+            t1_exclusive_rows.append(row_w_schema)
+
+        elif is_t2_exclusive:
+            t2_exclusive_rows.append(row_w_schema)
+
+        else:
+            diff_rows.append(row_w_schema)
+    
+    return t1_exclusive_rows, t2_exclusive_rows, diff_rows
+
+
+def _jsonify_diff(row: dict[str, Any], key_columns: list[str]) -> JsonDiffRowValue:
+    columns = collections.defaultdict(dict)
+    for field, value in row.items():
+        if field in ('is_exclusive_a', 'is_exclusive_b'):
+            continue
+
+        if field.startswith('is_diff_'):
+            column_name = field.replace('is_diff_', '')
+            columns[column_name]['isDiff'] = bool(value)
+
+        elif field.endswith('_a'):
+            column_name = field.replace('_a', '')
+            columns[column_name]['table1'] = value
+            columns[column_name]['isPK'] = column_name in key_columns
+
+        elif field.endswith('_b'):
+            column_name = field.replace('_b', '')
+            columns[column_name]['table2'] = value
+            columns[column_name]['isPK'] = column_name in key_columns
+    
+    return columns
+
+
+def _jsonify_exclusive(row: dict[str, Any], key_columns: list[str]) -> JsonExclusiveRow:
+    columns = collections.defaultdict(dict)
+    for field, value in row.items():
+        if field in ('is_exclusive_a', 'is_exclusive_b'):
+            continue
+        if field.startswith('is_diff_'):
+            continue
+        if field.endswith('_b') and row['is_exclusive_b']:
+            column_name = field.replace('_b', '')
+            columns[column_name]['isPK'] = column_name in key_columns
+            columns[column_name]['value'] = value
+        elif field.endswith('_a') and row['is_exclusive_a']:
+            column_name = field.replace('_a', '')
+            columns[column_name]['isPK'] = column_name in key_columns
+            columns[column_name]['value'] = value
+    return columns
+
+
+def _jsonify_diff_summary(stats_dict: dict) -> JsonDiffSummary:
+    return {
+        'rows': {
+            'total': {
+               'table1': stats_dict["rows_A"],
+               'table2': stats_dict["rows_B"]
+            },
+            'exclusive': {
+                'table1': stats_dict["exclusive_A"],
+                'table2': stats_dict["exclusive_B"],
+            },
+            'updated': stats_dict["updated"],
+            'unchanged': stats_dict["unchanged"]
+        },
+        'stats': {
+            'diffCounts': stats_dict["stats"]['diff_counts']
+        }
+    }
+
+
+def _jsonify_columns_diff(added_columns: list[str], 
+                          removed_columns: list[str],
+                          changed_columns: list[str]) -> JsonColumnsSummary:
+    columns = {
+        'exclusive': {
+            'table2': list(added_columns),
+            'table1': list(removed_columns),
+        },
+        'typeChanged': list(changed_columns),
+    }
+    return columns

data_diff/info_tree.py

@@ -1,3 +1,4 @@
+from typing import Any
 from typing import List, Dict
 
 from runtype import dataclass
@@ -9,14 +10,16 @@
 class SegmentInfo:
     tables: List[TableSegment]
 
-    diff: list = None
+    diff: list[tuple[Any, ...]] = None
+    diff_schema: tuple[tuple[str, type], ...] = None
     is_diff: bool = None
     diff_count: int = None
 
     rowcounts: Dict[int, int] = {}
     max_rows: int = None
 
-    def set_diff(self, diff):
+    def set_diff(self, schema: tuple[tuple[str, type]], diff: list[tuple[Any, ...]]):
+        self.diff_schema = schema
         self.diff = diff
         self.diff_count = len(diff)
         self.is_diff = self.diff_count > 0

data_diff/joindiff_tables.py

@@ -206,7 +206,7 @@ def _diff_segments(
             assert len(a_cols) == len(b_cols)
             logger.debug("Querying for different rows")
             diff = db.query(diff_rows, list)
-            info_tree.info.set_diff(diff)
+            info_tree.info.set_diff(tuple(diff_rows.schema.items()), diff)
             for is_xa, is_xb, *x in diff:
                 if is_xa and is_xb:
                     # Can't both be exclusive, meaning a pk is NULL