Feat: Greatly simplify table diff CLI invocation (#977)

* feat: greatly simplify table diff cli api

* fix: use metavar

* feat: add padding to cli and improve diff output

* fix: ensure change is exclusive to cli interface and arbitrary comp is supported

* fix: ensure model name is not required

* fix: address cli on cond usage to use multiple and actually return a seq

* chore: use row diff verbiage for brevity and semantic correctness

* ci: run linter

* ci: fix test

* chore: make aliases optional,  intent more obvious

* nit: make shorthand on flag lowercase

---------

Co-authored-by: Toby Mao <toby.mao@gmail.com>

Author: z3z1ma

sqlmesh/cli/main.py

@@ -6,8 +6,6 @@
 import typing as t
 
 import click
-from sqlglot import exp
-from sqlglot.helper import ensure_list
 
 from sqlmesh import debug_mode_enabled, enable_logging
 from sqlmesh.cli import error_handler
@@ -388,35 +386,14 @@ def create_external_models(obj: Context) -> None:
 
 
 @cli.command("table_diff")
+@click.argument("source_to_target", required=True, metavar="SOURCE:TARGET")
+@click.argument("model", required=False)
 @click.option(
-    "--source",
-    "-s",
-    type=str,
-    required=True,
-    help="The source environment or table.",
-)
-@click.option(
-    "--target",
-    "-t",
-    type=str,
-    required=True,
-    help="The target environment or table.",
-)
-@click.option(
-    "--grain",
-    type=str,
-    multiple=True,
-    help="The list of columns to use as keys.",
-)
-@click.option(
+    "-o",
     "--on",
     type=str,
-    help='The SQL join condition or list of columns to use as keys. Table aliases must be "s" and "t" for source and target.',
-)
-@click.option(
-    "--model",
-    type=str,
-    help="The model to diff against when source and target are environments and not tables.",
+    multiple=True,
+    help="The column to join on. Can be specified multiple times. The model grain will be used if not specified.",
 )
 @click.option(
     "--where",
@@ -426,20 +403,22 @@ def create_external_models(obj: Context) -> None:
 @click.option(
     "--limit",
     type=int,
+    default=20,
     help="The limit of the sample dataframe.",
 )
 @click.pass_obj
 @error_handler
-def table_diff(obj: Context, **kwargs: t.Any) -> None:
-    """Show the diff between two tables.
-
-    Can either be two tables or two environments and a model.
-    """
-    kwargs["model_or_snapshot"] = kwargs.pop("model", None)
-    on = kwargs.pop("on", None)
-    grain = ensure_list(kwargs.pop("grain", None))
-    kwargs["on"] = exp.condition(on) if on else grain
-    obj.table_diff(**kwargs)
+def table_diff(
+    obj: Context, source_to_target: str, model: t.Optional[str], **kwargs: t.Any
+) -> None:
+    """Show the diff between two tables."""
+    source, target = source_to_target.split(":")
+    obj.table_diff(
+        source=source,
+        target=target,
+        model_or_snapshot=model,
+        **kwargs,
+    )
 
 
 @cli.command("prompt")

sqlmesh/core/console.py

@@ -459,7 +459,9 @@ def loading_stop(self, id: uuid.UUID) -> None:
         del self.loading_status[id]
 
     def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
-        tree = Tree(f"[bold]Schema Diff Between '{schema_diff.source}' and '{schema_diff.target}':")
+        tree = Tree(
+            f"\n[b]Schema Diff Between '[yellow]{schema_diff.source}[/yellow]' and '[green]{schema_diff.target}[/green]':"
+        )
 
         if schema_diff.added:
             added = Tree("[green]Added Columns:")
@@ -482,10 +484,19 @@ def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
         self.console.print(tree)
 
     def show_row_diff(self, row_diff: RowDiff) -> None:
-        self.console.print(
-            f"[bold]Row Count:[/bold] {row_diff.source}: {row_diff.source_count}, {row_diff.target}: {row_diff.target_count} -- {row_diff.count_pct_change}%"
-        )
-        self.console.print(row_diff.sample.to_string(index=False))
+        source_name = row_diff.source
+        if row_diff.source_alias:
+            source_name = row_diff.source_alias.upper()
+        target_name = row_diff.target
+        if row_diff.target_alias:
+            target_name = row_diff.target_alias.upper()
+
+        self.console.print("\n[b]Row Count:[/b]")
+        self.console.print(f" [yellow]{source_name}[/yellow]: {row_diff.source_count} rows")
+        self.console.print(f" [green]{target_name}[/green]: {row_diff.target_count} rows")
+        self.console.print(f"\n[b]Row Diff[b]: {row_diff.count_pct_change:.1f}%")
+        self.console.print("\n[b]Sample Rows:[/b]")
+        self.console.print(row_diff.sample.to_string(index=False), end="\n\n")
 
     def _get_snapshot_change_category(
         self, snapshot: Snapshot, plan: Plan, auto_apply: bool

sqlmesh/core/context.py

@@ -777,6 +777,7 @@ def table_diff(
         Returns:
             The TableDiff object containing schema and summary differences.
         """
+        source_alias, target_alias = source, target
         if model_or_snapshot:
             model = self.get_model(model_or_snapshot, raise_if_missing=True)
             source_env = self.state_reader.get_environment(source)
@@ -793,6 +794,8 @@ def table_diff(
             target = next(
                 snapshot for snapshot in target_env.snapshots if snapshot.name == model.name
             ).table_name()
+            source_alias = source_env.name
+            target_alias = target_env.name
 
             if not on and model.grain:
                 on = model.grain
@@ -801,7 +804,14 @@ def table_diff(
             raise SQLMeshError("Missing join condition 'on'")
 
         table_diff = TableDiff(
-            adapter=self._engine_adapter, source=source, target=target, on=on, where=where
+            adapter=self._engine_adapter,
+            source=source,
+            target=target,
+            on=on,
+            where=where,
+            source_alias=source_alias,
+            target_alias=target_alias,
+            limit=limit,
         )
         if show:
             self.console.show_schema_diff(table_diff.schema_diff())

sqlmesh/core/table_diff.py

@@ -52,6 +52,8 @@ class RowDiff(PydanticModel, frozen=True):
     target: str
     stats: t.Dict[str, float]
     sample: pd.DataFrame
+    source_alias: t.Optional[str] = None
+    target_alias: t.Optional[str] = None
 
     @property
     def source_count(self) -> int:
@@ -83,14 +85,19 @@ def __init__(
         where: t.Optional[str | exp.Condition] = None,
         dialect: DialectType = None,
         limit: int = 20,
+        source_alias: t.Optional[str] = None,
+        target_alias: t.Optional[str] = None,
     ):
         self.adapter = adapter
         self.source = source
         self.target = target
         self.where = exp.condition(where, dialect=dialect) if where else None
         self.limit = limit
+        # Support environment aliases for diff output improvement in certain cases
+        self.source_alias = source_alias
+        self.target_alias = target_alias
 
-        if isinstance(on, list):
+        if isinstance(on, (list, tuple)):
             self.on: exp.Condition = exp.and_(
                 *(
                     exp.column(c, "s").eq(exp.column(c, "t"))
@@ -194,5 +201,7 @@ def row_diff(self) -> RowDiff:
                     target=self.target,
                     stats=self.adapter.fetchdf(summary_query).iloc[0].to_dict(),
                     sample=self.adapter.fetchdf(sample_query),
+                    source_alias=self.source_alias,
+                    target_alias=self.target_alias,
                 )
         return self._row_diff