add events for dbt

Author: dlawin

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,8 @@
 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"
@@ -20,6 +23,7 @@
 LOWER_DBT_V = "1.0.0"
 UPPER_DBT_V = "1.5.0"
 
+set_entrypoint_name("CLI-dbt")
 
 @dataclass
 class DiffVars:
@@ -190,22 +194,52 @@ 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
+    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 +264,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,6 @@
 
 from data_diff.info_tree import InfoTree, SegmentInfo
 
-from .utils import run_as_daemon, safezip, getLogger
 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,9 +31,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

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,7 @@ 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)