Feat: table-diff, easily diff two tables. (#907)

* Feat: table-diff, easily diff two tables.

* pr feedback

* Fix: test no longer creates duckdb file.

* add cli and magics

Author: tobymao

sqlmesh/cli/main.py

@@ -392,6 +392,54 @@ def create_external_models(obj: Context) -> None:
     obj.create_external_models()
 
 
+@cli.command("table_diff")
+@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(
+    "--on",
+    type=str,
+    nargs="+",
+    required=True,
+    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.",
+)
+@click.option(
+    "--where",
+    type=str,
+    help="An optional where statement to filter results.",
+)
+@click.option(
+    "--limit",
+    type=int,
+    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)
+    obj.table_diff(**kwargs)
+
+
 @cli.command("prompt")
 @click.argument("prompt")
 @click.option(

sqlmesh/core/config/loader.py

@@ -32,6 +32,10 @@ def load_config_from_paths(
 
         extension = path.name.split(".")[-1].lower()
         if extension in ("yml", "yaml"):
+            if config_name != "config":
+                raise ConfigError(
+                    f"YAML configs do not support multiple configs. Use Python instead."
+                )
             next_config = load_config_from_yaml(path)
         elif extension == "py":
             next_config = load_config_from_python_module(path, config_name=config_name)

sqlmesh/core/console.py

@@ -24,6 +24,7 @@
 
     from sqlmesh.core.context_diff import ContextDiff
     from sqlmesh.core.plan import Plan
+    from sqlmesh.core.table_diff import RowDiff, SchemaDiff
 
     LayoutWidget = t.TypeVar("LayoutWidget", bound=t.Union[widgets.VBox, widgets.HBox])
 
@@ -122,6 +123,14 @@ 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_schema_diff(self, schema_diff: SchemaDiff) -> None:
+        """Show table schema diff"""
+
+    @abc.abstractmethod
+    def show_row_diff(self, row_diff: RowDiff) -> None:
+        """Show table summary diff"""
+
 
 class TerminalConsole(Console):
     """A rich based implementation of the console"""
@@ -449,6 +458,35 @@ def loading_stop(self, id: uuid.UUID) -> None:
         self.loading_status[id].stop()
         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}':")
+
+        if schema_diff.added:
+            added = Tree("[green]Added Columns:")
+            for c, t in schema_diff.added:
+                added.add(f"[green]{c} ({t})")
+            tree.add(added)
+
+        if schema_diff.removed:
+            removed = Tree("[red]Removed Columns:")
+            for c, t in schema_diff.removed:
+                removed.add(f"[red]{c} ({t})")
+            tree.add(removed)
+
+        if schema_diff.modified:
+            modified = Tree("[magenta]Modified Columns:")
+            for c, (ft, tt) in schema_diff.modified.items():
+                modified.add(f"[magenta]{c} ({ft} -> {tt})")
+            tree.add(modified)
+
+        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))
+
     def _get_snapshot_change_category(
         self, snapshot: Snapshot, plan: Plan, auto_apply: bool
     ) -> None:

sqlmesh/core/context.py

@@ -66,16 +66,23 @@
     to_table_mapping,
 )
 from sqlmesh.core.state_sync import StateReader, StateSync
+from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import get_all_model_tests, run_model_tests, run_tests
 from sqlmesh.core.user import User
 from sqlmesh.utils import UniqueKeyDict, sys_path
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import TimeLike, yesterday_ds
-from sqlmesh.utils.errors import ConfigError, MissingDependencyError, PlanError
+from sqlmesh.utils.errors import (
+    ConfigError,
+    MissingDependencyError,
+    PlanError,
+    SQLMeshError,
+)
 from sqlmesh.utils.jinja import JinjaMacroRegistry
 
 if t.TYPE_CHECKING:
     import graphviz
+    from typing_extensions import Literal
 
     from sqlmesh.core.engine_adapter._typing import DF, PySparkDataFrame, PySparkSession
 
@@ -252,6 +259,12 @@ def engine_adapter(self) -> EngineAdapter:
         """Returns an engine adapter."""
         return self._engine_adapter
 
+    def execution_context(self, is_dev: bool = False) -> ExecutionContext:
+        """Returns an execution context."""
+        return ExecutionContext(
+            engine_adapter=self._engine_adapter, snapshots=self.snapshots, is_dev=is_dev
+        )
+
     def upsert_model(self, model: t.Union[str, Model], **kwargs: t.Any) -> Model:
         """Update or insert a model.
 
@@ -371,9 +384,76 @@ def run(
         if not skip_janitor and environment.lower() == c.PROD:
             self._run_janitor()
 
-    def get_model(self, name: str) -> t.Optional[Model]:
-        """Returns a model with the given name or None if a model with such name doesn't exist."""
-        return self._models.get(name)
+    @t.overload
+    def get_model(
+        self, model_or_snapshot: ModelOrSnapshot, raise_if_missing: Literal[True] = True
+    ) -> Model:
+        ...
+
+    @t.overload
+    def get_model(
+        self, model_or_snapshot: ModelOrSnapshot, raise_if_missing: Literal[False] = False
+    ) -> t.Optional[Model]:
+        ...
+
+    def get_model(
+        self, model_or_snapshot: ModelOrSnapshot, raise_if_missing: bool = False
+    ) -> t.Optional[Model]:
+        """Returns a model with the given name or None if a model with such name doesn't exist.
+
+        Args:
+            model_or_snapshot: A model name, model, or snapshot.
+            raise_if_missing: Raises an error if a model is not found.
+
+        Returns:
+            The expected model.
+        """
+        if isinstance(model_or_snapshot, str):
+            model = self._models.get(model_or_snapshot)
+        elif isinstance(model_or_snapshot, Snapshot):
+            model = model_or_snapshot.model
+        else:
+            model = model_or_snapshot
+
+        if raise_if_missing and not model:
+            raise SQLMeshError(f"Cannot find model for '{model_or_snapshot}'")
+        return model
+
+    @t.overload
+    def get_snapshot(
+        self, model_or_snapshot: ModelOrSnapshot, raise_if_missing: Literal[True]
+    ) -> Snapshot:
+        ...
+
+    @t.overload
+    def get_snapshot(
+        self, model_or_snapshot: ModelOrSnapshot, raise_if_missing: Literal[False]
+    ) -> t.Optional[Snapshot]:
+        ...
+
+    def get_snapshot(
+        self, model_or_snapshot: ModelOrSnapshot, raise_if_missing: bool = False
+    ) -> t.Optional[Snapshot]:
+        """Returns a snapshot with the given name or None if a snapshot with such name doesn't exist.
+
+        Args:
+            model_or_snapshot: A model name, model, or snapshot.
+            raise_if_missing: Raises an error if a snapshot is not found.
+
+        Returns:
+            The expected snapshot.
+        """
+        if isinstance(model_or_snapshot, str):
+            snapshot = self.snapshots.get(model_or_snapshot)
+        elif isinstance(model_or_snapshot, Snapshot):
+            snapshot = model_or_snapshot
+        else:
+            snapshot = self.snapshots.get(model_or_snapshot.name)
+
+        if raise_if_missing and not snapshot:
+            raise SQLMeshError(f"Cannot find snapshot for '{model_or_snapshot}'")
+
+        return snapshot
 
     def config_for_path(self, path: Path) -> Config:
         for config_path, config in self.configs.items():
@@ -478,18 +558,17 @@ def render(
         """
         latest = latest or yesterday_ds()
 
-        if isinstance(model_or_snapshot, str):
-            model = self._models[model_or_snapshot]
-        elif isinstance(model_or_snapshot, Snapshot):
-            model = model_or_snapshot.model
-        else:
-            model = model_or_snapshot
+        model = self.get_model(model_or_snapshot, raise_if_missing=True)
 
         expand = self.dag.upstream(model.name) if expand is True else expand or []
 
         if model.is_seed:
-            df = next(model.render(self, start=start, end=end, latest=latest, **kwargs))
-            return next(pandas_to_sql(df, model.columns_to_types))
+            df = next(
+                model.render(
+                    context=self.execution_context(), start=start, end=end, latest=latest, **kwargs
+                )
+            )
+            return next(pandas_to_sql(t.cast(pd.DataFrame, df), model.columns_to_types))
 
         return model.render_query(
             start=start,
@@ -520,12 +599,7 @@ def evaluate(
             latest: The latest time used for non incremental datasets.
             limit: A limit applied to the model.
         """
-        if isinstance(model_or_snapshot, str):
-            snapshot = self.snapshots[model_or_snapshot]
-        elif isinstance(model_or_snapshot, Snapshot):
-            snapshot = model_or_snapshot
-        else:
-            snapshot = self.snapshots[model_or_snapshot.name]
+        snapshot = self.get_snapshot(model_or_snapshot, raise_if_missing=True)
 
         df = self.snapshot_evaluator.evaluate(
             snapshot,
@@ -671,6 +745,55 @@ def diff(self, environment: t.Optional[str] = None, detailed: bool = False) -> N
             self._context_diff(environment or c.PROD), detailed
         )
 
+    def table_diff(
+        self,
+        source: str,
+        target: str,
+        on: t.List[str] | exp.Condition,
+        model_or_snapshot: t.Optional[ModelOrSnapshot] = None,
+        where: t.Optional[str | exp.Condition] = None,
+        limit: int = 20,
+        show: bool = True,
+    ) -> TableDiff:
+        """Show a diff between two tables.
+
+        Args:
+            source: The source environment or table.
+            target: The target environment or table.
+            on: The join condition, table aliases must be "s" and "t" for source and target.
+            model_or_snapshot: The model or snapshot to use when environments are passed in.
+            where: An optional where statement to filter results.
+            limit: The limit of the sample dataframe.
+            show: Show the table diff in the console.
+
+        Returns:
+            The TableDiff object containing schema and summary differences.
+        """
+        if model_or_snapshot:
+            model = self.get_model(model_or_snapshot, raise_if_missing=True)
+            source_env = self.state_reader.get_environment(source)
+            target_env = self.state_reader.get_environment(target)
+
+            if not source_env:
+                raise SQLMeshError(f"Could not find environment '{source}'")
+            if not target_env:
+                raise SQLMeshError(f"Could not find environment '{target}')")
+
+            source = next(
+                snapshot for snapshot in source_env.snapshots if snapshot.name == model.name
+            ).table_name()
+            target = next(
+                snapshot for snapshot in target_env.snapshots if snapshot.name == model.name
+            ).table_name()
+
+        table_diff = TableDiff(
+            adapter=self._engine_adapter, source=source, target=target, on=on, where=where
+        )
+        if show:
+            self.console.show_schema_diff(table_diff.schema_diff())
+            self.console.show_row_diff(table_diff.row_diff())
+        return table_diff
+
     def get_dag(self, format: str = "svg") -> graphviz.Digraph:
         """Gets a graphviz dag.
 

sqlmesh/core/engine_adapter/base.py

@@ -781,6 +781,29 @@ def execute(
             logger.debug(f"Executing SQL:\n{sql}")
             self.cursor.execute(sql, **kwargs)
 
+    @contextlib.contextmanager
+    def temp_table(self, query_or_df: QueryOrDF, name: str = "diff") -> t.Iterator[exp.Table]:
+        """A context manager for working a temp table.
+
+        The table will be created with a random guid and cleaned up after the block.
+
+        Args:
+            query_or_df: The query or df to create a temp table for.
+            name: The base name of the temp table.
+
+        Yields:
+            The table expression
+        """
+
+        with self.transaction(TransactionType.DDL):
+            table = self._get_temp_table(name)
+            self.ctas(table, query_or_df)
+
+            try:
+                yield table
+            finally:
+                self.drop_table(table)
+
     def _create_table_properties(
         self,
         storage_format: t.Optional[str] = None,