Feat(table_diff): Add flag to warn when models lack grain and diff the remaining (#4449)

Author: themisvaltinos

docs/guides/tablediff.md

@@ -169,7 +169,10 @@ sqlmesh table_diff prod:dev -m "tag:finance" -m "metrics.*_daily"
 
 When multiple selectors are provided, they are combined with OR logic, meaning a model matching any of the selectors will be included.
 
-> Note: All models being compared must have their `grain` defined that is unique and not null, as this is used to perform the join between the tables in the two environments.
+!!! note
+    All models being compared must have their `grain` defined that is unique and not null, as this is used to perform the join between the tables in the two environments.
+
+    If the `--warn-grain-check` option is used, this requirement is not enforced. Instead of raising an error, a warning is displayed for the models without a defined grain and diffs are computed for the remaining models.
 
 ## Diffing tables or views
 

docs/reference/cli.md

@@ -557,6 +557,8 @@ Options:
                            floating point columns. Default: 3
   --skip-grain-check       Disable the check for a primary key (grain) that is
                            missing or is not unique.
+  --warn-grain-check       Warn if any selected model is missing a grain,
+                           and compute diffs for the remaining models.
   --temp-schema TEXT       Schema used for temporary tables. It can be
                            `CATALOG.SCHEMA` or `SCHEMA`. Default:
                            `sqlmesh_temp`

docs/reference/notebook.md

@@ -300,8 +300,8 @@ Create a schema file containing external model schemas.
 %table_diff [--on [ON ...]] [--skip-columns [SKIP_COLUMNS ...]]
                 [--model MODEL] [--where WHERE] [--limit LIMIT]
                 [--show-sample] [--decimals DECIMALS] [--skip-grain-check]
-                [--temp-schema SCHEMA] [--select-model [SELECT_MODEL ...]]
-                SOURCE:TARGET
+                [--warn-grain-check] [--temp-schema SCHEMA]
+                [--select-model [SELECT_MODEL ...]] SOURCE:TARGET
 
 Show the diff between two tables.
 
@@ -326,6 +326,8 @@ options:
                         floating point columns. Default: 3
   --skip-grain-check    Disable the check for a primary key (grain) that is
                         missing or is not unique.
+  --warn-grain-check    Warn if any selected model is missing a grain,
+                        and compute diffs for the remaining models.
   --temp-schema SCHEMA  The schema to use for temporary tables.
   --select-model <[SELECT_MODEL ...]>
                         Select specific models to diff using a pattern.

sqlmesh/cli/main.py

@@ -906,6 +906,11 @@ def create_external_models(obj: Context, **kwargs: t.Any) -> None:
     is_flag=True,
     help="Disable the check for a primary key (grain) that is missing or is not unique.",
 )
+@click.option(
+    "--warn-grain-check",
+    is_flag=True,
+    help="Warn if any selected model is missing a grain, and compute diffs for the remaining models.",
+)
 @click.option(
     "--temp-schema",
     type=str,

sqlmesh/core/context.py

@@ -1629,6 +1629,7 @@ def table_diff(
         show_sample: bool = True,
         decimals: int = 3,
         skip_grain_check: bool = False,
+        warn_grain_check: bool = False,
         temp_schema: t.Optional[str] = None,
     ) -> t.List[TableDiff]:
         """Show a diff between two tables.
@@ -1703,22 +1704,29 @@ def table_diff(
                             target_env.naming_info, adapter.dialect
                         )
                         model_on = on or model.on
-                        models_to_diff.append((model, adapter, source, target, model_on))
                         if not model_on:
                             models_without_grain.append(model)
+                        else:
+                            models_to_diff.append((model, adapter, source, target, model_on))
+
+            if models_without_grain:
+                model_names = "\n".join(
+                    f"─ {model.name} \n  at '{model._path}'" for model in models_without_grain
+                )
+                message = (
+                    "SQLMesh doesn't know how to join the tables for the following models:\n"
+                    f"{model_names}\n\n"
+                    "Please specify a `grain` in each model definition. It must be unique and not null."
+                )
+                if warn_grain_check:
+                    self.console.log_warning(message)
+                else:
+                    raise SQLMeshError(message)
 
             if models_to_diff:
                 self.console.show_table_diff_details(
                     [model[0].name for model in 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 `grain` in each model definition. Must be unique and not null."
-                    )
 
                 self.console.start_table_diff_progress(len(models_to_diff))
                 try:

sqlmesh/magics.py

@@ -702,6 +702,11 @@ def create_external_models(self, context: Context, line: str) -> None:
         action="store_true",
         help="Disable the check for a primary key (grain) that is missing or is not unique.",
     )
+    @argument(
+        "--warn-grain-check",
+        action="store_true",
+        help="Warn if any selected model is missing a grain, and compute diffs for the remaining models.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def table_diff(self, context: Context, line: str) -> None:
@@ -723,6 +728,7 @@ def table_diff(self, context: Context, line: str) -> None:
             show_sample=args.show_sample,
             decimals=args.decimals,
             skip_grain_check=args.skip_grain_check,
+            warn_grain_check=args.warn_grain_check,
         )
 
     @magic_arguments()

tests/core/test_table_diff.py

@@ -14,6 +14,8 @@
 from sqlmesh.core.table_diff import TableDiff
 import numpy as np
 
+from sqlmesh.utils.errors import SQLMeshError
+
 
 def create_test_console() -> t.Tuple[StringIO, TerminalConsole]:
     """Creates a console and buffer for validating console output."""
@@ -829,3 +831,123 @@ def test_data_diff_empty_tables():
     assert (
         strip_ansi_codes(output) == "Neither the source nor the target table contained any records"
     )
+
+
+@pytest.mark.slow
+def test_data_diff_multiple_models_lacking_grain(sushi_context_fixed_date, capsys, caplog):
+    # Create first model with grain
+    expressions = d.parse(
+        """
+        MODEL (name memory.sushi.grain_model, kind full, grain(key),);
+        SELECT
+            key,
+            value,
+        FROM
+            (VALUES
+                (1, 3),
+                (2, 4),
+            ) AS t (key, value)
+    """
+    )
+    model_s = load_sql_based_model(expressions, dialect="snowflake")
+    sushi_context_fixed_date.upsert_model(model_s)
+
+    # Create second model without grain
+    expressions_2 = d.parse(
+        """
+        MODEL (name memory.sushi.no_grain_model, kind full,);
+        SELECT
+            key,
+            value as amount,
+        FROM
+            memory.sushi.grain_model
+    """
+    )
+    model_s2 = load_sql_based_model(expressions_2, dialect="snowflake")
+    sushi_context_fixed_date.upsert_model(model_s2)
+
+    sushi_context_fixed_date.plan(
+        "source_dev",
+        no_prompts=True,
+        auto_apply=True,
+        skip_tests=True,
+        start="2023-01-31",
+        end="2023-01-31",
+    )
+
+    # Modify first model
+    model = sushi_context_fixed_date.models['"MEMORY"."SUSHI"."GRAIN_MODEL"']
+    modified_model = model.dict()
+    modified_model["query"] = (
+        exp.select("*")
+        .from_(model.query.subquery())
+        .union("SELECT key, value FROM (VALUES (1, 6),(2,3),) AS t (key, value)")
+    )
+    modified_sqlmodel = SqlModel(**modified_model)
+    sushi_context_fixed_date.upsert_model(modified_sqlmodel)
+
+    # Modify second model
+    model2 = sushi_context_fixed_date.models['"MEMORY"."SUSHI"."NO_GRAIN_MODEL"']
+    modified_model2 = model2.dict()
+    modified_model2["query"] = (
+        exp.select("*")
+        .from_(model2.query.subquery())
+        .union("SELECT key, amount FROM (VALUES (5, 150.2),(6,250.2),) AS t (key, amount)")
+    )
+    modified_sqlmodel2 = SqlModel(**modified_model2)
+    sushi_context_fixed_date.upsert_model(modified_sqlmodel2)
+
+    sushi_context_fixed_date.auto_categorize_changes = CategorizerConfig(
+        sql=AutoCategorizationMode.FULL
+    )
+    sushi_context_fixed_date.plan(
+        "target_dev",
+        create_from="source_dev",
+        no_prompts=True,
+        auto_apply=True,
+        skip_tests=True,
+        start="2023-01-31",
+        end="2023-01-31",
+    )
+
+    # By default erroring out when even one model lacks a grain
+    with pytest.raises(
+        SQLMeshError,
+        match=r"SQLMesh doesn't know how to join the tables for the following models:*",
+    ):
+        sushi_context_fixed_date.table_diff(
+            source="source_dev",
+            target="target_dev",
+            select_models={"*"},
+            skip_grain_check=False,
+        )
+
+    # With warn_grain_check flag the diff will go ahead by warning
+    diffs = sushi_context_fixed_date.table_diff(
+        source="source_dev",
+        target="target_dev",
+        select_models={"*"},
+        skip_grain_check=False,
+        warn_grain_check=True,
+    )
+
+    # Check that the diff was performed only for the model with a grain
+    assert len(diffs) == 1
+    diff1 = diffs[0]
+
+    # Check the table diff corresponds to the grain model
+    row_diff1 = diff1.row_diff()
+    assert row_diff1.full_match_count == 2.0
+    assert row_diff1.full_match_pct == 50.0
+    assert row_diff1.s_only_count == 0.0
+    assert row_diff1.t_only_count == 0.0
+    assert row_diff1.stats["join_count"] == 4.0
+    assert row_diff1.stats["null_grain_count"] == 0.0
+    assert row_diff1.stats["s_count"] == 4.0
+    assert row_diff1.stats["distinct_count_s"] == 2.0
+    assert row_diff1.stats["t_count"] == 4.0
+    assert row_diff1.stats["distinct_count_t"] == 2.0
+    assert row_diff1.s_sample.shape == (0, 2)
+    assert row_diff1.t_sample.shape == (0, 2)
+    assert row_diff1.joined_sample.shape == (2, 3)
+    assert row_diff1.sample.shape == (2, 4)