Merge branch 'master' into dx-821

Author: dlawin

data_diff/__main__.py

@@ -289,6 +289,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"],
                 state=state,
             )
         else:

data_diff/dbt.py

@@ -1,3 +1,4 @@
+import json
 import os
 import re
 import time
@@ -17,6 +18,8 @@
 from . import connect_to_table, diff_tables, Algorithm
 from .cloud import DatafoldAPI, TCloudApiDataDiff, TCloudApiOrgMeta
 from .dbt_parser import DbtParser, TDatadiffConfig
+from .diff_tables import DiffResultWrapper
+from .format import jsonify, jsonify_error
 from .tracking import (
     bool_ask_for_email,
     create_email_signup_event_json,
@@ -54,13 +57,15 @@ class TDiffVars(pydantic.BaseModel):
     where_filter: Optional[str] = None
     include_columns: List[str]
     exclude_columns: List[str]
+    dbt_model: Optional[str] = None
 
 
 def dbt_diff(
     profiles_dir_override: Optional[str] = None,
     project_dir_override: Optional[str] = None,
     is_cloud: bool = False,
     dbt_selection: Optional[str] = None,
+    json_output: bool = False,
     state: Optional[str] = None,
 ) -> None:
     print_version_info()
@@ -115,12 +120,25 @@ def dbt_diff(
                 diff_thread = run_as_daemon(_cloud_diff, diff_vars, config.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))
-                + "Skipped due to unknown primary key. Add uniqueness tests, meta, or tags.\n"
-            )
+            if json_output:
+                print(
+                    json.dumps(
+                        jsonify_error(
+                            table1=diff_vars.prod_path,
+                            table2=diff_vars.dev_path,
+                            dbt_model=diff_vars.dbt_model,
+                            error="No primary key found. Add uniqueness tests, meta, or tags.",
+                        )
+                    ),
+                    flush=True,
+                )
+            else:
+                rich.print(
+                    _diff_output_base(".".join(diff_vars.dev_path), ".".join(diff_vars.prod_path))
+                    + "Skipped due to unknown primary key. Add uniqueness tests, meta, or tags.\n"
+                )
 
     # wait for all threads
     if diff_threads:
@@ -155,6 +173,7 @@ def _get_diff_vars(
     datadiff_model_config = dbt_parser.get_datadiff_model_config(model.meta)
 
     return TDiffVars(
+        dbt_model=model.unique_id,
         dev_path=dev_qualified_list,
         prod_path=prod_qualified_list,
         primary_keys=primary_keys,
@@ -205,15 +224,15 @@ def _get_prod_path_from_manifest(model, prod_manifest) -> Union[Tuple[str, str],
     return prod_database, prod_schema
 
 
-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)
 
-    table1 = connect_to_table(diff_vars.connection, dev_qualified_str, tuple(diff_vars.primary_keys), diff_vars.threads)
-    table2 = connect_to_table(
+    table1 = connect_to_table(
         diff_vars.connection, prod_qualified_str, tuple(diff_vars.primary_keys), diff_vars.threads
     )
+    table2 = connect_to_table(diff_vars.connection, dev_qualified_str, tuple(diff_vars.primary_keys), diff_vars.threads)
 
     table1_columns = table1.get_schema()
     try:
@@ -228,11 +247,11 @@ def _local_diff(diff_vars: TDiffVars) -> None:
     table1_column_names = set(table1_columns.keys())
     table2_column_names = set(table2_columns.keys())
     column_set = table1_column_names.intersection(table2_column_names)
-    columns_added = table1_column_names.difference(table2_column_names)
-    columns_removed = table2_column_names.difference(table1_column_names)
+    columns_added = table2_column_names.difference(table1_column_names)
+    columns_removed = table1_column_names.difference(table2_column_names)
     # col type is i = 1 in tuple
     columns_type_changed = {
-        k for k, v in table1_columns.items() if k in table2_columns and v[1] != table2_columns[k][1]
+        k for k, v in table2_columns.items() if k in table1_columns and v[1] != table1_columns[k][1]
     }
 
     if columns_added:
@@ -255,7 +274,7 @@ def _local_diff(diff_vars: TDiffVars) -> None:
 
     extra_columns = tuple(column_set)
 
-    diff = diff_tables(
+    diff: DiffResultWrapper = diff_tables(
         table1,
         table2,
         threaded=True,
@@ -264,6 +283,35 @@ 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
+        try:
+            list(diff)
+        except Exception as e:
+            print(
+                json.dumps(
+                    jsonify_error(list(table1.table_path), list(table2.table_path), diff_vars.dbt_model, str(e))
+                ),
+                flush=True,
+            )
+            return
+
+        print(
+            json.dumps(
+                jsonify(
+                    diff,
+                    dbt_model=diff_vars.dbt_model,
+                    with_summary=True,
+                    with_columns={
+                        "added": columns_added,
+                        "removed": columns_removed,
+                        "changed": columns_type_changed,
+                    },
+                )
+            ),
+            flush=True,
+        )
+        return
 
     if list(diff):
         diff_output_str += f"{diff.get_stats_string(is_dbt=True)} \n"