Make local dbt data diffs concurrent

data_diff/databases/base.py

@@ -931,7 +931,7 @@ def name(self):
     def compile(self, sql_ast):
         return self.dialect.compile(Compiler(self), sql_ast)
 
-    def query(self, sql_ast: Union[Expr, Generator], res_type: type = None):
+    def query(self, sql_ast: Union[Expr, Generator], res_type: type = None, log_message: Optional[str] = None):
         """Query the given SQL code/AST, and attempt to convert the result to type 'res_type'
 
         If given a generator, it will execute all the yielded sql queries with the same thread and cursor.
@@ -956,7 +956,10 @@ def query(self, sql_ast: Union[Expr, Generator], res_type: type = None):
                 if sql_code is SKIP:
                     return SKIP
 
-            logger.debug("Running SQL (%s):\n%s", self.name, sql_code)
+            if log_message:
+                logger.debug("Running SQL (%s): %s \n%s", self.name, log_message, sql_code)
+            else:
+                logger.debug("Running SQL (%s):\n%s", self.name, sql_code)
 
         if self._interactive and isinstance(sql_ast, Select):
             explained_sql = self.compile(Explain(sql_ast))
@@ -1022,7 +1025,7 @@ def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
         Note: This method exists instead of select_table_schema(), just because not all databases support
               accessing the schema using a SQL query.
         """
-        rows = self.query(self.select_table_schema(path), list)
+        rows = self.query(self.select_table_schema(path), list, log_message=path)
         if not rows:
             raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns")
 
@@ -1044,7 +1047,7 @@ def query_table_unique_columns(self, path: DbPath) -> List[str]:
         """Query the table for its unique columns for table in 'path', and return {column}"""
         if not self.SUPPORTS_UNIQUE_CONSTAINT:
             raise NotImplementedError("This database doesn't support 'unique' constraints")
-        res = self.query(self.select_table_unique_columns(path), List[str])
+        res = self.query(self.select_table_unique_columns(path), List[str], log_message=path)
         return list(res)
 
     def _process_table_schema(
@@ -1086,7 +1089,9 @@ def _refine_coltypes(
         fields = [Code(self.dialect.normalize_uuid(self.dialect.quote(c), String_UUID())) for c in text_columns]
 
         samples_by_row = self.query(
-            table(*table_path).select(*fields).where(Code(where) if where else SKIP).limit(sample_size), list
+            table(*table_path).select(*fields).where(Code(where) if where else SKIP).limit(sample_size),
+            list,
+            log_message=table_path,
         )
         if not samples_by_row:
             raise ValueError(f"Table {table_path} is empty.")

data_diff/dbt.py

@@ -8,6 +8,7 @@
 import pydantic
 import rich
 from rich.prompt import Prompt
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from data_diff.errors import (
     DataDiffCustomSchemaNoConfigError,
@@ -80,7 +81,6 @@ def dbt_diff(
     production_schema_flag: Optional[str] = None,
 ) -> None:
     print_version_info()
-    diff_threads = []
     set_entrypoint_name(os.getenv("DATAFOLD_TRIGGERED_BY", "CLI-dbt"))
     dbt_parser = DbtParser(profiles_dir_override, project_dir_override, state)
     models = dbt_parser.get_models(dbt_selection)
@@ -112,7 +112,11 @@ def dbt_diff(
     else:
         dbt_parser.set_connection()
 
-    with log_status_handler.status if log_status_handler else nullcontext():
+    futures = {}
+
+    with log_status_handler.status if log_status_handler else nullcontext(), ThreadPoolExecutor(
+        max_workers=dbt_parser.threads
+    ) as executor:
         for model in models:
             if log_status_handler:
                 log_status_handler.set_prefix(f"Diffing {model.alias} \n")
@@ -140,12 +144,12 @@ def dbt_diff(
 
             if diff_vars.primary_keys:
                 if is_cloud:
-                    diff_thread = run_as_daemon(
+                    future = executor.submit(
                         _cloud_diff, diff_vars, config.datasource_id, api, org_meta, log_status_handler
                     )
-                    diff_threads.append(diff_thread)
                 else:
-                    _local_diff(diff_vars, json_output)
+                    future = executor.submit(_local_diff, diff_vars, json_output, log_status_handler)
+                futures[future] = model
             else:
                 if json_output:
                     print(
@@ -165,10 +169,12 @@ def dbt_diff(
                         + "Skipped due to unknown primary key. Add uniqueness tests, meta, or tags.\n"
                     )
 
-        # wait for all threads
-        if diff_threads:
-            for thread in diff_threads:
-                thread.join()
+    for future in as_completed(futures):
+        model = futures[future]
+        try:
+            future.result()  # if error occurred, it will be raised here
+        except Exception as e:
+            logger.error(f"An error occurred during the execution of a diff task: {model.unique_id} - {e}")
 
     _extension_notification()
 
@@ -265,15 +271,17 @@ def _get_prod_path_from_manifest(model, prod_manifest) -> Union[Tuple[str, str,
     return prod_database, prod_schema, prod_alias
 
 
-def _local_diff(diff_vars: TDiffVars, json_output: bool = False) -> None:
+def _local_diff(
+    diff_vars: TDiffVars, json_output: bool = False, log_status_handler: Optional[LogStatusHandler] = None
+) -> None:
+    if log_status_handler:
+        log_status_handler.diff_started(diff_vars.dev_path[-1])
     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, 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 = connect_to_table(diff_vars.connection, prod_qualified_str, tuple(diff_vars.primary_keys))
+    table2 = connect_to_table(diff_vars.connection, dev_qualified_str, tuple(diff_vars.primary_keys))
 
     try:
         table1_columns = table1.get_schema()
@@ -373,6 +381,9 @@ def _local_diff(diff_vars: TDiffVars, json_output: bool = False) -> None:
         diff_output_str += no_differences_template()
         rich.print(diff_output_str)
 
+    if log_status_handler:
+        log_status_handler.diff_finished(diff_vars.dev_path[-1])
+
 
 def _initialize_api() -> Optional[DatafoldAPI]:
     datafold_host = os.environ.get("DATAFOLD_HOST")
@@ -406,7 +417,7 @@ def _cloud_diff(
     log_status_handler: Optional[LogStatusHandler] = None,
 ) -> None:
     if log_status_handler:
-        log_status_handler.cloud_diff_started(diff_vars.dev_path[-1])
+        log_status_handler.diff_started(diff_vars.dev_path[-1])
     diff_output_str = _diff_output_base(".".join(diff_vars.dev_path), ".".join(diff_vars.prod_path))
     payload = TCloudApiDataDiff(
         data_source1_id=datasource_id,
@@ -476,7 +487,7 @@ def _cloud_diff(
             rich.print(diff_output_str)
 
         if log_status_handler:
-            log_status_handler.cloud_diff_finished(diff_vars.dev_path[-1])
+            log_status_handler.diff_finished(diff_vars.dev_path[-1])
     except BaseException as ex:  # Catch KeyboardInterrupt too
         error = ex
     finally:

data_diff/dbt_parser.py

@@ -446,17 +446,17 @@ def get_pk_from_model(self, node, unique_columns: dict, pk_tag: str) -> List[str
 
             from_meta = [name for name, params in node.columns.items() if pk_tag in params.meta] or None
             if from_meta:
-                logger.debug("Found PKs via META: " + str(from_meta))
+                logger.debug(f"Found PKs via META [{node.name}]: " + str(from_meta))
                 return from_meta
 
             from_tags = [name for name, params in node.columns.items() if pk_tag in params.tags] or None
             if from_tags:
-                logger.debug("Found PKs via Tags: " + str(from_tags))
+                logger.debug(f"Found PKs via Tags [{node.name}]: " + str(from_tags))
                 return from_tags
             if node.unique_id in unique_columns:
                 from_uniq = unique_columns.get(node.unique_id)
                 if from_uniq is not None:
-                    logger.debug("Found PKs via Uniqueness tests: " + str(from_uniq))
+                    logger.debug(f"Found PKs via Uniqueness tests [{node.name}]: {str(from_uniq)}")
                     return list(from_uniq)
 
         except (KeyError, IndexError, TypeError) as e:

data_diff/joindiff_tables.py

@@ -162,7 +162,7 @@ def _diff_tables_root(self, table1: TableSegment, table2: TableSegment, info_tre
                 yield from self._diff_segments(None, table1, table2, info_tree, None)
             else:
                 yield from self._bisect_and_diff_tables(table1, table2, info_tree)
-            logger.info("Diffing complete")
+            logger.info(f"Diffing complete: {table1.table_path} <> {table2.table_path}")
             if self.materialize_to_table:
                 logger.info("Materialized diff to table '%s'.", ".".join(self.materialize_to_table))
 
@@ -193,8 +193,8 @@ def _diff_segments(
             partial(self._collect_stats, 1, table1, info_tree),
             partial(self._collect_stats, 2, table2, info_tree),
             partial(self._test_null_keys, table1, table2),
-            partial(self._sample_and_count_exclusive, db, diff_rows, a_cols, b_cols),
-            partial(self._count_diff_per_column, db, diff_rows, list(a_cols), is_diff_cols),
+            partial(self._sample_and_count_exclusive, db, diff_rows, a_cols, b_cols, table1, table2),
+            partial(self._count_diff_per_column, db, diff_rows, list(a_cols), is_diff_cols, table1, table2),
             partial(
                 self._materialize_diff,
                 db,
@@ -205,8 +205,8 @@ def _diff_segments(
             else None,
         ):
             assert len(a_cols) == len(b_cols)
-            logger.debug("Querying for different rows")
-            diff = db.query(diff_rows, list)
+            logger.debug(f"Querying for different rows: {table1.table_path}")
+            diff = db.query(diff_rows, list, log_message=table1.table_path)
             info_tree.info.set_diff(diff, schema=tuple(diff_rows.schema.items()))
             for is_xa, is_xb, *x in diff:
                 if is_xa and is_xb:
@@ -227,7 +227,7 @@ def _diff_segments(
                     yield "+", tuple(b_row)
 
     def _test_duplicate_keys(self, table1: TableSegment, table2: TableSegment):
-        logger.debug("Testing for duplicate keys")
+        logger.debug(f"Testing for duplicate keys: {table1.table_path} <> {table2.table_path}")
 
         # Test duplicate keys
         for ts in [table1, table2]:
@@ -240,24 +240,24 @@ def _test_duplicate_keys(self, table1: TableSegment, table2: TableSegment):
 
             unvalidated = list(set(key_columns) - set(unique))
             if unvalidated:
-                logger.info(f"Validating that the are no duplicate keys in columns: {unvalidated}")
+                logger.info(f"Validating that the are no duplicate keys in columns: {unvalidated} for {ts.table_path}")
                 # Validate that there are no duplicate keys
                 self.stats["validated_unique_keys"] = self.stats.get("validated_unique_keys", []) + [unvalidated]
                 q = t.select(total=Count(), total_distinct=Count(Concat(this[unvalidated]), distinct=True))
-                total, total_distinct = ts.database.query(q, tuple)
+                total, total_distinct = ts.database.query(q, tuple, log_message=ts.table_path)
                 if total != total_distinct:
                     raise ValueError("Duplicate primary keys")
 
     def _test_null_keys(self, table1, table2):
-        logger.debug("Testing for null keys")
+        logger.debug(f"Testing for null keys: {table1.table_path} <> {table2.table_path}")
 
         # Test null keys
         for ts in [table1, table2]:
             t = ts.make_select()
             key_columns = ts.key_columns
 
             q = t.select(*this[key_columns]).where(or_(this[k] == None for k in key_columns))
-            nulls = ts.database.query(q, list)
+            nulls = ts.database.query(q, list, log_message=ts.table_path)
             if nulls:
                 if self.skip_null_keys:
                     logger.warning(
@@ -267,7 +267,7 @@ def _test_null_keys(self, table1, table2):
                     raise ValueError(f"NULL values in one or more primary keys of {ts.table_path}")
 
     def _collect_stats(self, i, table_seg: TableSegment, info_tree: InfoTree):
-        logger.debug(f"Collecting stats for table #{i}")
+        logger.debug(f"Collecting stats for table #{i}: {table_seg.table_path}")
         db = table_seg.database
 
         # Metrics
@@ -288,7 +288,7 @@ def _collect_stats(self, i, table_seg: TableSegment, info_tree: InfoTree):
         )
         col_exprs["count"] = Count()
 
-        res = db.query(table_seg.make_select().select(**col_exprs), tuple)
+        res = db.query(table_seg.make_select().select(**col_exprs), tuple, log_message=table_seg.table_path)
 
         for col_name, value in safezip(col_exprs, res):
             if value is not None:
@@ -303,7 +303,7 @@ def _collect_stats(self, i, table_seg: TableSegment, info_tree: InfoTree):
                 else:
                     self.stats[stat_name] = value
 
-        logger.debug("Done collecting stats for table #%s", i)
+        logger.debug("Done collecting stats for table #%s: %s", i, table_seg.table_path)
 
     def _create_outer_join(self, table1, table2):
         db = table1.database
@@ -334,23 +334,46 @@ def _create_outer_join(self, table1, table2):
         diff_rows = all_rows.where(or_(this[c] == 1 for c in is_diff_cols))
         return diff_rows, a_cols, b_cols, is_diff_cols, all_rows
 
-    def _count_diff_per_column(self, db, diff_rows, cols, is_diff_cols):
-        logger.debug("Counting differences per column")
-        is_diff_cols_counts = db.query(diff_rows.select(sum_(this[c]) for c in is_diff_cols), tuple)
+    def _count_diff_per_column(
+        self,
+        db,
+        diff_rows,
+        cols,
+        is_diff_cols,
+        table1: Optional[TableSegment] = None,
+        table2: Optional[TableSegment] = None,
+    ):
+        logger.info(type(table1))
+        logger.debug(f"Counting differences per column: {table1.table_path} <> {table2.table_path}")
+        is_diff_cols_counts = db.query(
+            diff_rows.select(sum_(this[c]) for c in is_diff_cols),
+            tuple,
+            log_message=f"{table1.table_path} <> {table2.table_path}",
+        )
         diff_counts = {}
         for name, count in safezip(cols, is_diff_cols_counts):
             diff_counts[name] = diff_counts.get(name, 0) + (count or 0)
         self.stats["diff_counts"] = diff_counts
 
-    def _sample_and_count_exclusive(self, db, diff_rows, a_cols, b_cols):
+    def _sample_and_count_exclusive(
+        self,
+        db,
+        diff_rows,
+        a_cols,
+        b_cols,
+        table1: Optional[TableSegment] = None,
+        table2: Optional[TableSegment] = None,
+    ):
         if isinstance(db, (Oracle, MsSQL)):
             exclusive_rows_query = diff_rows.where((this.is_exclusive_a == 1) | (this.is_exclusive_b == 1))
         else:
             exclusive_rows_query = diff_rows.where(this.is_exclusive_a | this.is_exclusive_b)
 
         if not self.sample_exclusive_rows:
-            logger.debug("Counting exclusive rows")
-            self.stats["exclusive_count"] = db.query(exclusive_rows_query.count(), int)
+            logger.debug(f"Counting exclusive rows: {table1.table_path} <> {table2.table_path}")
+            self.stats["exclusive_count"] = db.query(
+                exclusive_rows_query.count(), int, log_message=f"{table1.table_path} <> {table2.table_path}"
+            )
             return
 
         logger.info("Counting and sampling exclusive rows")

data_diff/utils.py

@@ -485,31 +485,31 @@ def __init__(self):
         super().__init__()
         self.status = Status("")
         self.prefix = ""
-        self.cloud_diff_status = {}
+        self.diff_status = {}
 
     def emit(self, record):
         log_entry = self.format(record)
-        if self.cloud_diff_status:
-            self._update_cloud_status(log_entry)
+        if self.diff_status:
+            self._update_diff_status(log_entry)
         else:
             self.status.update(self.prefix + log_entry)
 
     def set_prefix(self, prefix_string):
         self.prefix = prefix_string
 
-    def cloud_diff_started(self, model_name):
-        self.cloud_diff_status[model_name] = "[yellow]In Progress[/]"
-        self._update_cloud_status()
+    def diff_started(self, model_name):
+        self.diff_status[model_name] = "[yellow]In Progress[/]"
+        self._update_diff_status()
 
-    def cloud_diff_finished(self, model_name):
-        self.cloud_diff_status[model_name] = "[green]Finished   [/]"
-        self._update_cloud_status()
+    def diff_finished(self, model_name):
+        self.diff_status[model_name] = "[green]Finished   [/]"
+        self._update_diff_status()
 
-    def _update_cloud_status(self, log=None):
-        cloud_status_string = "\n"
-        for model_name, status in self.cloud_diff_status.items():
-            cloud_status_string += f"{status} {model_name}\n"
-        self.status.update(f"{cloud_status_string}{log or ''}")
+    def _update_diff_status(self, log=None):
+        status_string = "\n"
+        for model_name, status in self.diff_status.items():
+            status_string += f"{status} {model_name}\n"
+        self.status.update(f"{status_string}{log or ''}")
 
 
 class UnknownMeta(type):