final changes and lint

Author: dlawin

data_diff/dbt.py

@@ -12,7 +12,13 @@
 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 .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
 
@@ -23,7 +29,6 @@
 LOWER_DBT_V = "1.0.0"
 UPPER_DBT_V = "1.5.0"
 
-set_entrypoint_name("CLI-dbt")
 
 @dataclass
 class DiffVars:
@@ -37,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()
@@ -200,6 +206,7 @@ def _cloud_diff(diff_vars: DiffVars) -> None:
 
     start = time.monotonic()
     error = None
+    diff_id = None
     try:
         response = requests.request("POST", url, headers=headers, json=payload, timeout=30)
         response.raise_for_status()
@@ -217,29 +224,29 @@ def _cloud_diff(diff_vars: DiffVars) -> None:
             + "\n"
         )
     except BaseException as ex:  # Catch KeyboardInterrupt too
-            error = ex
+        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
+                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
+        if error:
+            raise error
 
 
 class DbtParser:

data_diff/diff_tables.py

@@ -14,6 +14,7 @@
 
 from data_diff.info_tree import InfoTree, SegmentInfo
 
+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
@@ -31,8 +32,6 @@ class Algorithm(Enum):
 DiffResult = Iterator[Tuple[str, tuple]]  # Iterator[Tuple[Literal["+", "-"], tuple]]
 
 
-
-
 @dataclass
 class ThreadBase:
     "Provides utility methods for optional threading"
@@ -120,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"
@@ -139,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,
@@ -186,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

@@ -92,7 +92,7 @@ def create_end_event_json(
     diff_count: int,
     error: Optional[str],
     diff_id: Optional[int] = None,
-    is_cloud: bool = False
+    is_cloud: bool = False,
 ):
     return {
         "event": "os_diff_run_end",
@@ -111,7 +111,7 @@ def create_end_event_json(
             "data_diff_version:": __version__,
             "entrypoint_name": entrypoint_name,
             "is_cloud": is_cloud,
-            "diff_id": diff_id
+            "diff_id": diff_id,
         },
     }
 

data_diff/utils.py

@@ -73,6 +73,7 @@ def get_timestamp(_match):
 
     return re.sub("%t", get_timestamp, name)
 
+
 def truncate_error(error: str):
     first_line = error.split("\n", 1)[0]
     return re.sub("'(.*?)'", "'***'", first_line)