refactors; ux improvements

Author: themisvaltinos

sqlmesh/cli/main.py

@@ -894,6 +894,7 @@ def create_external_models(obj: Context, **kwargs: t.Any) -> None:
 )
 @click.option(
     "--select-model",
+    "-m",
     type=str,
     multiple=True,
     help="Select specific model that should be diffed.",

sqlmesh/core/context.py

@@ -1569,8 +1569,8 @@ def table_diff(
         self,
         source: str,
         target: str,
-        on: t.List[str] | exp.Condition | None = None,
-        skip_columns: t.List[str] | None = None,
+        on: t.Optional[t.List[str] | exp.Condition] = None,
+        skip_columns: t.Optional[t.List[str]] = None,
         select_models: t.Optional[t.Collection[str]] = None,
         where: t.Optional[str | exp.Condition] = None,
         limit: int = 20,
@@ -1579,7 +1579,7 @@ def table_diff(
         decimals: int = 3,
         skip_grain_check: bool = False,
         temp_schema: t.Optional[str] = None,
-    ) -> t.Union[TableDiff, t.List[TableDiff]]:
+    ) -> t.List[TableDiff]:
         """Show a diff between two tables.
 
         Args:
@@ -1613,26 +1613,30 @@ def table_diff(
                 raise SQLMeshError(f"Could not find environment '{target}'")
 
             selected_models = self._new_selector().expand_model_selections(select_models)
-            models_to_diff: t.List[t.Tuple[Model, EngineAdapter, str, str]] = []
+            models_to_diff: t.List[
+                t.Tuple[Model, EngineAdapter, str, str, t.Optional[t.List[str] | exp.Condition]]
+            ] = []
             models_in_source: t.List[str] = []
             models_in_target: t.List[str] = []
             models_no_diff: t.List[str] = []
+            models_without_grain: t.List[Model] = []
+            source_snapshots_to_name = {
+                snapshot.name: snapshot for snapshot in source_env.snapshots
+            }
+            target_snapshots_to_name = {
+                snapshot.name: snapshot for snapshot in target_env.snapshots
+            }
 
-            for model_or_snapshot in selected_models:
-                model = self.get_model(model_or_snapshot, raise_if_missing=True)
+            for model_fqn in selected_models:
+                model = self._models[model_fqn]
                 adapter = self._get_engine_adapter(model.gateway)
-                source_snapshot = next(
-                    (snapshot for snapshot in source_env.snapshots if snapshot.name == model.fqn),
-                    None,
-                )
-                target_snapshot = next(
-                    (snapshot for snapshot in target_env.snapshots if snapshot.name == model.fqn),
-                    None,
-                )
+                source_snapshot = source_snapshots_to_name.get(model.fqn)
+                target_snapshot = target_snapshots_to_name.get(model.fqn)
+
                 if source_snapshot is None and target_snapshot:
-                    models_in_source.append(model_or_snapshot)
+                    models_in_source.append(model_fqn)
                 elif target_snapshot is None and source_snapshot:
-                    models_in_target.append(model_or_snapshot)
+                    models_in_target.append(model_fqn)
                 elif target_snapshot and source_snapshot:
                     if source_snapshot.fingerprint != target_snapshot.fingerprint:
                         # Compare the virtual layer instead of the physical layer because the virtual layer is guaranteed to point
@@ -1644,9 +1648,27 @@ def table_diff(
                             target_env.naming_info, adapter.dialect
                         )
 
-                        models_to_diff.append((model, adapter, source, target))
+                        model_on = []
+                        if not on:
+                            for expr in [
+                                ref.expression for ref in model.all_references if ref.unique
+                            ]:
+                                if isinstance(expr, exp.Tuple):
+                                    model_on.extend(
+                                        [
+                                            key.this.sql(dialect=adapter.dialect)
+                                            for key in expr.expressions
+                                        ]
+                                    )
+                                else:
+                                    # Handle a single Column or Paren expression
+                                    model_on.append(expr.this.sql(dialect=adapter.dialect))
+
+                        models_to_diff.append((model, adapter, source, target, on or model_on))
+                        if not (on or model_on):
+                            models_without_grain.append(model)
                     else:
-                        models_no_diff.append(model_or_snapshot)
+                        models_no_diff.append(model_fqn)
 
             self.console.show_table_diff_details(
                 models_in_source,
@@ -1656,6 +1678,15 @@ def table_diff(
             )
 
             if models_to_diff:
+                if models_without_grain:
+                    model_names = "\n".join(
+                        f"─ {model.name} \n  at '{model._path}'" for model in models_without_grain
+                    )
+                    raise SQLMeshError(
+                        f"SQLMesh doesn't know how to join the tables for the following models:\n{model_names}\n"
+                        "\nPlease specify the `grains` in each model definition."
+                    )
+
                 self.console.start_table_diff_progress(len(models_to_diff))
                 tasks_num = min(len(models_to_diff), self.concurrent_tasks)
                 table_diffs = concurrent_apply_to_values(
@@ -1665,11 +1696,11 @@ def table_diff(
                         adapter=model_info[1],
                         source=model_info[2],
                         target=model_info[3],
+                        on=model_info[4],
                         source_alias=source_env.name,
                         target_alias=target_env.name,
                         limit=limit,
                         decimals=decimals,
-                        on=on,
                         skip_columns=skip_columns,
                         where=where,
                         show=show,
@@ -1698,7 +1729,7 @@ def table_diff(
         if show:
             self.console.show_table_diff(table_diffs, show_sample, skip_grain_check, temp_schema)
 
-        return table_diffs[0] if len(table_diffs) == 1 else table_diffs
+        return table_diffs
 
     def _model_diff(
         self,
@@ -1717,15 +1748,6 @@ def _model_diff(
         temp_schema: t.Optional[str] = None,
         skip_grain_check: bool = False,
     ) -> TableDiff:
-        if not on:
-            on = []
-            for expr in [ref.expression for ref in model.all_references if ref.unique]:
-                if isinstance(expr, exp.Tuple):
-                    on.extend([key.this.sql(dialect=adapter.dialect) for key in expr.expressions])
-                else:
-                    # Handle a single Column or Paren expression
-                    on.append(expr.this.sql(dialect=adapter.dialect))
-
         self.console.start_table_diff_model_progress(model.name)
 
         table_diff = self._table_diff(

tests/core/test_table_diff.py

@@ -111,7 +111,7 @@ def test_data_diff(sushi_context_fixed_date, capsys, caplog):
         target="target_dev",
         on=exp.condition("s.customer_id = t.customer_id AND s.event_date = t.event_date"),
         select_models={"sushi.customer_revenue_by_day"},
-    )
+    )[0]
 
     # verify queries were actually logged to the log file, this helps immensely with debugging
     console_output = capsys.readouterr()
@@ -169,7 +169,7 @@ def test_data_diff_decimals(sushi_context_fixed_date):
         source="table_diff_source",
         target="table_diff_target",
         on=["key"],
-    )
+    )[0]
     assert diff.row_diff().full_match_count == 3
     assert diff.row_diff().partial_match_count == 0
 
@@ -178,7 +178,7 @@ def test_data_diff_decimals(sushi_context_fixed_date):
         target="table_diff_target",
         on=["key"],
         decimals=4,
-    )
+    )[0]
 
     row_diff = diff.row_diff()
     joined_sample_columns = row_diff.joined_sample.columns
@@ -293,7 +293,7 @@ def test_grain_check(sushi_context_fixed_date):
         on=["'key_1'", "key_2"],
         select_models={"memory.sushi*"},
         skip_grain_check=False,
-    )
+    )[0]
 
     row_diff = diff.row_diff()
     assert row_diff.full_match_count == 7
@@ -421,7 +421,7 @@ def test_tables_and_grain_inferred_from_model(sushi_context_fixed_date: Context)
 
     table_diff = sushi_context_fixed_date.table_diff(
         source="unit_test", target="prod", select_models={"sushi.waiter_revenue_by_day"}
-    )
+    )[0]
     assert isinstance(table_diff, TableDiff)
     assert table_diff.source == "memory.sushi__unit_test.waiter_revenue_by_day"
     assert table_diff.target == "memory.sushi.waiter_revenue_by_day"