Merge branch 'dbt' of https://github.com/dlawin/data-diff into dlawin-dbt

Author: erezsh

data_diff/dbt.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import time
 import rich
 import yaml
 from dataclasses import dataclass
@@ -11,6 +12,14 @@
 from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
 from dbt.config.renderer import ProfileRenderer
 
+from .tracking import (
+    set_entrypoint_name,
+    create_end_event_json,
+    create_start_event_json,
+    send_event_json,
+    is_tracking_enabled,
+)
+from .utils import run_as_daemon, truncate_error
 from . import connect_to_table, diff_tables, Algorithm
 
 RUN_RESULTS_PATH = "/target/run_results.json"
@@ -33,6 +42,7 @@ class DiffVars:
 def dbt_diff(
     profiles_dir_override: Optional[str] = None, project_dir_override: Optional[str] = None, is_cloud: bool = False
 ) -> None:
+    set_entrypoint_name("CLI-dbt")
     dbt_parser = DbtParser(profiles_dir_override, project_dir_override, is_cloud)
     models = dbt_parser.get_models()
     dbt_parser.set_project_dict()
@@ -190,22 +200,53 @@ def _cloud_diff(diff_vars: DiffVars) -> None:
         "Authorization": f"Key {api_key}",
         "Content-Type": "application/json",
     }
+    if is_tracking_enabled():
+        event_json = create_start_event_json({"is_cloud": True, "datasource_id": diff_vars.datasource_id})
+        run_as_daemon(send_event_json, event_json)
 
-    response = requests.request("POST", url, headers=headers, json=payload, timeout=30)
-    response.raise_for_status()
-    data = response.json()
-    diff_id = data["id"]
-    # TODO in future we should support self hosted datafold
-    diff_url = f"https://app.datafold.com/datadiffs/{diff_id}/overview"
-    rich.print(
-        "[red]"
-        + ".".join(diff_vars.dev_path)
-        + " <> "
-        + ".".join(diff_vars.prod_path)
-        + "[/] \n    Diff in progress: \n    "
-        + diff_url
-        + "\n"
-    )
+    start = time.monotonic()
+    error = None
+    diff_id = None
+    try:
+        response = requests.request("POST", url, headers=headers, json=payload, timeout=30)
+        response.raise_for_status()
+        data = response.json()
+        diff_id = data["id"]
+        # TODO in future we should support self hosted datafold
+        diff_url = f"https://app.datafold.com/datadiffs/{diff_id}/overview"
+        rich.print(
+            "[red]"
+            + ".".join(diff_vars.dev_path)
+            + " <> "
+            + ".".join(diff_vars.prod_path)
+            + "[/] \n    Diff in progress: \n    "
+            + diff_url
+            + "\n"
+        )
+    except BaseException as ex:  # Catch KeyboardInterrupt too
+        error = ex
+    finally:
+        # we don't currently have much of this information
+        # but I imagine a future iteration of this _cloud method
+        # will poll for results
+        if is_tracking_enabled():
+            err_message = truncate_error(repr(error))
+            event_json = create_end_event_json(
+                is_success=error is None,
+                runtime_seconds=time.monotonic() - start,
+                data_source_1_type="",
+                data_source_2_type="",
+                table1_count=0,
+                table2_count=0,
+                diff_count=0,
+                error=err_message,
+                diff_id=diff_id,
+                is_cloud=True,
+            )
+            send_event_json(event_json)
+
+        if error:
+            raise error
 
 
 class DbtParser:
@@ -230,7 +271,6 @@ def get_models(self):
 
         dbt_version = parse_version(run_results_obj.metadata.dbt_version)
 
-        # TODO 1.4 support
         if dbt_version < parse_version(LOWER_DBT_V) or dbt_version >= parse_version(UPPER_DBT_V):
             raise Exception(
                 f"Found dbt: v{dbt_version} Expected the dbt project's version to be >= {LOWER_DBT_V} and < {UPPER_DBT_V}"

data_diff/diff_tables.py

@@ -14,7 +14,7 @@
 
 from data_diff.info_tree import InfoTree, SegmentInfo
 
-from .utils import run_as_daemon, safezip, getLogger
+from .utils import run_as_daemon, safezip, getLogger, truncate_error
 from .thread_utils import ThreadedYielder
 from .table_segment import TableSegment
 from .tracking import create_end_event_json, create_start_event_json, send_event_json, is_tracking_enabled
@@ -32,11 +32,6 @@ class Algorithm(Enum):
 DiffResult = Iterator[Tuple[str, tuple]]  # Iterator[Tuple[Literal["+", "-"], tuple]]
 
 
-def truncate_error(error: str):
-    first_line = error.split("\n", 1)[0]
-    return re.sub("'(.*?)'", "'***'", first_line)
-
-
 @dataclass
 class ThreadBase:
     "Provides utility methods for optional threading"
@@ -124,7 +119,6 @@ def _get_stats(self) -> DiffStats:
         return DiffStats(diff_by_sign, table1_count, table2_count, unchanged, diff_percent)
 
     def get_stats_string(self):
-
         diff_stats = self._get_stats()
         string_output = ""
         string_output += f"{diff_stats.table1_count} rows in table A\n"
@@ -143,7 +137,6 @@ def get_stats_string(self):
         return string_output
 
     def get_stats_dict(self):
-
         diff_stats = self._get_stats()
         json_output = {
             "rows_A": diff_stats.table1_count,
@@ -190,7 +183,6 @@ def _diff_tables_wrapper(self, table1: TableSegment, table2: TableSegment, info_
         start = time.monotonic()
         error = None
         try:
-
             # Query and validate schema
             table1, table2 = self._threaded_call("with_schema", [table1, table2])
             self._validate_and_adjust_columns(table1, table2)

data_diff/tracking.py

@@ -85,12 +85,14 @@ def create_start_event_json(diff_options: Dict[str, Any]):
 def create_end_event_json(
     is_success: bool,
     runtime_seconds: float,
-    db1: str,
-    db2: str,
+    data_source_1_type: str,
+    data_source_2_type: str,
     table1_count: int,
     table2_count: int,
     diff_count: int,
     error: Optional[str],
+    diff_id: Optional[int] = None,
+    is_cloud: bool = False,
 ):
     return {
         "event": "os_diff_run_end",
@@ -100,14 +102,16 @@ def create_end_event_json(
             "time": time(),
             "is_success": is_success,
             "runtime_seconds": runtime_seconds,
-            "data_source_1_type": db1,
-            "data_source_2_type": db2,
+            "data_source_1_type": data_source_1_type,
+            "data_source_2_type": data_source_2_type,
             "table_1_rows_cnt": table1_count,
             "table_2_rows_cnt": table2_count,
             "diff_rows_cnt": diff_count,
             "error_message": error,
             "data_diff_version:": __version__,
             "entrypoint_name": entrypoint_name,
+            "is_cloud": is_cloud,
+            "diff_id": diff_id,
         },
     }
 

data_diff/utils.py

@@ -72,3 +72,8 @@ def get_timestamp(_match):
         return datetime.now().isoformat("_", "seconds").replace(":", "_")
 
     return re.sub("%t", get_timestamp, name)
+
+
+def truncate_error(error: str):
+    first_line = error.split("\n", 1)[0]
+    return re.sub("'(.*?)'", "'***'", first_line)