Feat: Add support for concurrent table diff across all impacted models

Author: themisvaltinos

sqlmesh/cli/main.py

@@ -896,14 +896,21 @@ def create_external_models(obj: Context, **kwargs: t.Any) -> None:
 def table_diff(
     obj: Context, source_to_target: str, model: t.Optional[str], **kwargs: t.Any
 ) -> None:
-    """Show the diff between two tables."""
+    """Show the diff between two tables or all impacted when no model is specified."""
     source, target = source_to_target.split(":")
-    obj.table_diff(
-        source=source,
-        target=target,
-        model_or_snapshot=model,
-        **kwargs,
-    )
+    if model:
+        obj.table_diff(
+            source=source,
+            target=target,
+            model_or_snapshot=model,
+            **kwargs,
+        )
+    else:
+        obj.table_diff_impacted_models(
+            source=source,
+            target=target,
+            **kwargs,
+        )
 
 
 @cli.command("rewrite")

sqlmesh/core/console.py

@@ -197,7 +197,7 @@ def show_intervals(self, snapshot_intervals: t.Dict[Snapshot, SnapshotIntervals]
 
 
 class DifferenceConsole(abc.ABC):
-    """Console for displaying environment differences"""
+    """Console for displaying differences"""
 
     @abc.abstractmethod
     def show_environment_difference_summary(
@@ -217,6 +217,20 @@ def show_model_difference_summary(
     ) -> None:
         """Displays a summary of differences for the given models."""
 
+    @abc.abstractmethod
+    def show_table_diff_summary(self, table_diff: TableDiff) -> None:
+        """Display information about the tables being diffed and how they are being joined"""
+
+    @abc.abstractmethod
+    def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
+        """Show table schema diff."""
+
+    @abc.abstractmethod
+    def show_row_diff(
+        self, row_diff: RowDiff, show_sample: bool = True, skip_grain_check: bool = False
+    ) -> None:
+        """Show table summary diff."""
+
 
 class BaseConsole(abc.ABC):
     @abc.abstractmethod
@@ -424,20 +438,6 @@ def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
     def loading_stop(self, id: uuid.UUID) -> None:
         """Stop loading for the given id."""
 
-    @abc.abstractmethod
-    def show_table_diff_summary(self, table_diff: TableDiff) -> None:
-        """Display information about the tables being diffed and how they are being joined"""
-
-    @abc.abstractmethod
-    def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
-        """Show table schema diff."""
-
-    @abc.abstractmethod
-    def show_row_diff(
-        self, row_diff: RowDiff, show_sample: bool = True, skip_grain_check: bool = False
-    ) -> None:
-        """Show table summary diff."""
-
 
 class NoopConsole(Console):
     def start_plan_evaluation(self, plan: EvaluatablePlan) -> None:
@@ -697,6 +697,7 @@ class TerminalConsole(Console):
     """A rich based implementation of the console."""
 
     TABLE_DIFF_SOURCE_BLUE = "#0248ff"
+    TABLE_DIFF_TARGET_GREEN = "green"
 
     def __init__(
         self,
@@ -1544,39 +1545,58 @@ def _show_summary_tree_for(
                 )
             tree.add(self._limit_model_names(removed_tree, self.verbosity))
         if modified_snapshot_ids:
-            direct = Tree("[bold][direct]Directly Modified:")
-            indirect = Tree("[bold][indirect]Indirectly Modified:")
-            metadata = Tree("[bold][metadata]Metadata Updated:")
-            for s_id in modified_snapshot_ids:
-                name = s_id.name
-                display_name = context_diff.snapshots[s_id].display_name(
-                    environment_naming_info,
-                    default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
-                    dialect=self.dialect,
-                )
-                if context_diff.directly_modified(name):
-                    direct.add(
-                        f"[direct]{display_name}"
-                        if no_diff
-                        else Syntax(f"{display_name}\n{context_diff.text_diff(name)}", "sql")
-                    )
-                elif context_diff.indirectly_modified(name):
-                    indirect.add(f"[indirect]{display_name}")
-                elif context_diff.metadata_updated(name):
-                    metadata.add(
-                        f"[metadata]{display_name}"
-                        if no_diff
-                        else Syntax(f"{display_name}\n{context_diff.text_diff(name)}", "sql")
-                    )
+            tree = self._add_modified_models(
+                context_diff,
+                modified_snapshot_ids,
+                tree,
+                environment_naming_info,
+                default_catalog,
+                no_diff,
+            )
 
-            if direct.children:
-                tree.add(direct)
-            if indirect.children:
-                tree.add(self._limit_model_names(indirect, self.verbosity))
-            if metadata.children:
-                tree.add(metadata)
         self._print(tree)
 
+    def _add_modified_models(
+        self,
+        context_diff: ContextDiff,
+        modified_snapshot_ids: t.Set[SnapshotId],
+        tree: Tree,
+        environment_naming_info: EnvironmentNamingInfo,
+        default_catalog: t.Optional[str] = None,
+        no_diff: bool = True,
+    ) -> Tree:
+        direct = Tree("[bold][direct]Directly Modified:")
+        indirect = Tree("[bold][indirect]Indirectly Modified:")
+        metadata = Tree("[bold][metadata]Metadata Updated:")
+        for s_id in modified_snapshot_ids:
+            name = s_id.name
+            display_name = context_diff.snapshots[s_id].display_name(
+                environment_naming_info,
+                default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
+                dialect=self.dialect,
+            )
+            if context_diff.directly_modified(name):
+                direct.add(
+                    f"[direct]{display_name}"
+                    if no_diff
+                    else Syntax(f"{display_name}\n{context_diff.text_diff(name)}", "sql")
+                )
+            elif context_diff.indirectly_modified(name):
+                indirect.add(f"[indirect]{display_name}")
+            elif context_diff.metadata_updated(name):
+                metadata.add(
+                    f"[metadata]{display_name}"
+                    if no_diff
+                    else Syntax(f"{display_name}\n{context_diff.text_diff(name)}", "sql")
+                )
+        if direct.children:
+            tree.add(direct)
+        if indirect.children:
+            tree.add(self._limit_model_names(indirect, self.verbosity))
+        if metadata.children:
+            tree.add(metadata)
+        return tree
+
     def _show_options_after_categorization(
         self,
         plan_builder: PlanBuilder,
@@ -1897,7 +1917,9 @@ def show_table_diff_summary(self, table_diff: TableDiff) -> None:
             )
             envs.add(source)
 
-            target = Tree(f"Target: [green]{table_diff.target_alias}[/green]")
+            target = Tree(
+                f"Target: [{self.TABLE_DIFF_TARGET_GREEN}]{table_diff.target_alias}[/{self.TABLE_DIFF_TARGET_GREEN}]"
+            )
             envs.add(target)
 
             tree.add(envs)
@@ -1907,7 +1929,9 @@ def show_table_diff_summary(self, table_diff: TableDiff) -> None:
         tables.add(
             f"Source: [{self.TABLE_DIFF_SOURCE_BLUE}]{table_diff.source}[/{self.TABLE_DIFF_SOURCE_BLUE}]"
         )
-        tables.add(f"Target: [green]{table_diff.target}[/green]")
+        tables.add(
+            f"Target: [{self.TABLE_DIFF_TARGET_GREEN}]{table_diff.target}[/{self.TABLE_DIFF_TARGET_GREEN}]"
+        )
 
         tree.add(tables)
 
@@ -1928,7 +1952,7 @@ def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
         if schema_diff.target_alias:
             target_name = schema_diff.target_alias.upper()
 
-        first_line = f"\n[b]Schema Diff Between '[{self.TABLE_DIFF_SOURCE_BLUE}]{source_name}[/{self.TABLE_DIFF_SOURCE_BLUE}]' and '[green]{target_name}[/green]'"
+        first_line = f"\n[b]Schema Diff Between '[{self.TABLE_DIFF_SOURCE_BLUE}]{source_name}[/{self.TABLE_DIFF_SOURCE_BLUE}]' and '[{self.TABLE_DIFF_TARGET_GREEN}]{target_name}[/{self.TABLE_DIFF_TARGET_GREEN}]'"
         if schema_diff.model_name:
             first_line = (
                 first_line + f" environments for model '[blue]{schema_diff.model_name}[/blue]'"
@@ -2032,7 +2056,7 @@ def show_row_diff(
 
                 column_styles = {
                     source_name: self.TABLE_DIFF_SOURCE_BLUE,
-                    target_name: "green",
+                    target_name: self.TABLE_DIFF_TARGET_GREEN,
                 }
 
                 for column, [source_column, target_column] in columns.items():
@@ -2637,23 +2661,11 @@ def show_model_difference_summary(
             no_diff: Hide the actual SQL differences.
         """
         added_snapshots = {context_diff.snapshots[s_id] for s_id in context_diff.added}
-        added_snapshot_models = {s for s in added_snapshots if s.is_model}
-        if added_snapshot_models:
+        if added_snapshots:
             self._print("\n**Added Models:**")
-            added_models = sorted(added_snapshot_models)
-            list_length = len(added_models)
-            if (
-                self.verbosity < Verbosity.VERY_VERBOSE
-                and list_length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD
-            ):
-                self._print(added_models[0])
-                self._print(f"- `.... {list_length - 2} more ....`\n")
-                self._print(added_models[-1])
-            else:
-                for snapshot in added_models:
-                    self._print(
-                        f"- `{snapshot.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
-                    )
+            self._print_models_with_threshold(
+                environment_naming_info, {s for s in added_snapshots if s.is_model}, default_catalog
+            )
 
         added_snapshot_audits = {s for s in added_snapshots if s.is_audit}
         if added_snapshot_audits:
@@ -2664,23 +2676,13 @@ def show_model_difference_summary(
                 )
 
         removed_snapshot_table_infos = set(context_diff.removed_snapshots.values())
-        removed_model_snapshot_table_infos = {s for s in removed_snapshot_table_infos if s.is_model}
-        if removed_model_snapshot_table_infos:
+        if removed_snapshot_table_infos:
             self._print("\n**Removed Models:**")
-            removed_models = sorted(removed_model_snapshot_table_infos)
-            list_length = len(removed_models)
-            if (
-                self.verbosity < Verbosity.VERY_VERBOSE
-                and list_length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD
-            ):
-                self._print(removed_models[0])
-                self._print(f"- `.... {list_length - 2} more ....`\n")
-                self._print(removed_models[-1])
-            else:
-                for snapshot_table_info in removed_models:
-                    self._print(
-                        f"- `{snapshot_table_info.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
-                    )
+            self._print_models_with_threshold(
+                environment_naming_info,
+                {s for s in removed_snapshot_table_infos if s.is_model},
+                default_catalog,
+            )
 
         removed_audit_snapshot_table_infos = {s for s in removed_snapshot_table_infos if s.is_audit}
         if removed_audit_snapshot_table_infos:
@@ -2694,48 +2696,72 @@ def show_model_difference_summary(
             current_snapshot for current_snapshot, _ in context_diff.modified_snapshots.values()
         }
         if modified_snapshots:
-            directly_modified = []
-            indirectly_modified = []
-            metadata_modified = []
-            for snapshot in modified_snapshots:
-                if context_diff.directly_modified(snapshot.name):
-                    directly_modified.append(snapshot)
-                elif context_diff.indirectly_modified(snapshot.name):
-                    indirectly_modified.append(snapshot)
-                elif context_diff.metadata_updated(snapshot.name):
-                    metadata_modified.append(snapshot)
-            if directly_modified:
-                self._print("\n**Directly Modified:**")
-                for snapshot in sorted(directly_modified):
-                    self._print(
-                        f"- `{snapshot.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
-                    )
-                    if not no_diff:
-                        self._print(f"```diff\n{context_diff.text_diff(snapshot.name)}\n```")
-            if indirectly_modified:
-                self._print("\n**Indirectly Modified:**")
-                indirectly_modified = sorted(indirectly_modified)
-                modified_length = len(indirectly_modified)
-                if (
-                    self.verbosity < Verbosity.VERY_VERBOSE
-                    and modified_length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD
-                ):
-                    self._print(
-                        f"- `{indirectly_modified[0].display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`\n"
-                        f"- `.... {modified_length - 2} more ....`\n"
-                        f"- `{indirectly_modified[-1].display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
-                    )
-                else:
-                    for snapshot in indirectly_modified:
-                        self._print(
-                            f"- `{snapshot.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
-                        )
-            if metadata_modified:
-                self._print("\n**Metadata Updated:**")
-                for snapshot in sorted(metadata_modified):
-                    self._print(
-                        f"- `{snapshot.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
-                    )
+            self._print_modified_models(
+                context_diff, modified_snapshots, environment_naming_info, default_catalog, no_diff
+            )
+
+    def _print_models_with_threshold(
+        self,
+        environment_naming_info: EnvironmentNamingInfo,
+        snapshot_table_infos: t.Set[SnapshotInfoLike],
+        default_catalog: t.Optional[str] = None,
+    ) -> None:
+        models = sorted(snapshot_table_infos)
+        list_length = len(models)
+        if (
+            self.verbosity < Verbosity.VERY_VERBOSE
+            and list_length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD
+        ):
+            self._print(
+                f"- `{models[0].display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
+            )
+            self._print(f"- `.... {list_length - 2} more ....`\n")
+            self._print(
+                f"- `{models[-1].display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
+            )
+        else:
+            for snapshot_table_info in models:
+                self._print(
+                    f"- `{snapshot_table_info.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
+                )
+
+    def _print_modified_models(
+        self,
+        context_diff: ContextDiff,
+        modified_snapshots: t.Set[Snapshot],
+        environment_naming_info: EnvironmentNamingInfo,
+        default_catalog: t.Optional[str] = None,
+        no_diff: bool = True,
+    ) -> None:
+        directly_modified = []
+        indirectly_modified = []
+        metadata_modified = []
+        for snapshot in modified_snapshots:
+            if context_diff.directly_modified(snapshot.name):
+                directly_modified.append(snapshot)
+            elif context_diff.indirectly_modified(snapshot.name):
+                indirectly_modified.append(snapshot)
+            elif context_diff.metadata_updated(snapshot.name):
+                metadata_modified.append(snapshot)
+        if directly_modified:
+            self._print("\n**Directly Modified:**")
+            for snapshot in sorted(directly_modified):
+                self._print(
+                    f"- `{snapshot.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
+                )
+                if not no_diff:
+                    self._print(f"```diff\n{context_diff.text_diff(snapshot.name)}\n```")
+        if indirectly_modified:
+            self._print("\n**Indirectly Modified:**")
+            self._print_models_with_threshold(
+                environment_naming_info, set(indirectly_modified), default_catalog
+            )
+        if metadata_modified:
+            self._print("\n**Metadata Updated:**")
+            for snapshot in sorted(metadata_modified):
+                self._print(
+                    f"- `{snapshot.display_name(environment_naming_info, default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None, dialect=self.dialect)}`"
+                )
 
     def _show_missing_dates(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
         """Displays the models with missing dates."""