Repository: TobikoData/sqlmesh
Pull request #2895 (2024-07-11 19:28:23)

Feat: support skipping columns from being matched in table diff
Sometimes we can't compare certain columns using the equality (`=`) operator. For example, in BigQuery we can't compare arrays, i.e. `SELECT [1, 2, 3] = [1, 2, 3]` is invalid. Such SQL is generated [here](https://github.com/TobikoData/sqlmesh/blob/main/sqlmesh/core/table_diff.py#L243), resulting in a crash when the query is executed.

This PR adds a flag that can be used to skip certain columns from being matched, so users can selectively discard array columns from the table diffing process.


Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 5efcc6a6..e1957f0e 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -26,7 +26,7 @@ Commands:
   diff                    Show the diff between the local state and the...
   evaluate                Evaluate a model and return a dataframe with a...
   fetchdf                 Run a SQL query and display the results.
-  format                  Format all SQL models.
+  format                  Format all SQL models and audits.
   info                    Print information about a SQLMesh project.
   init                    Create a new SQLMesh repository.
   invalidate              Invalidate the target environment, forcing its...
@@ -166,7 +166,7 @@ Options:
 ```
 Usage: sqlmesh format [OPTIONS]
 
-  Format all SQL models.
+  Format all SQL models and audits.
 
 Options:
   -t, --transpile TEXT        Transpile project models to the specified
@@ -410,18 +410,20 @@ Usage: sqlmesh table_diff [OPTIONS] SOURCE:TARGET [MODEL]
   Show the diff between two tables.
 
 Options:
-  -o, --on TEXT           The column to join on. Can be specified multiple
-                          times. The model grain will be used if not
-                          specified.
-  --where TEXT            An optional where statement to filter results.
-  --limit INTEGER         The limit of the sample dataframe.
-  --show-sample           Show a sample of the rows that differ. With many
-                          columns, the output can be very wide.
-  --skip-grain-check      Disable the check for a primary key (grain) that 
-                          is missing or is not unique.
-  -d, --decimals INTEGER  The number of decimal places to keep when comparing
-                          floating point columns. Default: 3
-  --help                  Show this message and exit.
+  -o, --on TEXT            The column to join on. Can be specified multiple
+                           times. The model grain will be used if not
+                           specified.
+  -s, --skip-columns TEXT  The column(s) to skip when comparing the source and
+                           target table.
+  --where TEXT             An optional where statement to filter results.
+  --limit INTEGER          The limit of the sample dataframe.
+  --show-sample            Show a sample of the rows that differ. With many
+                           columns, the output can be very wide.
+  -d, --decimals INTEGER   The number of decimal places to keep when comparing
+                           floating point columns. Default: 3
+  --skip-grain-check       Disable the check for a primary key (grain) that is
+                           missing or is not unique.
+  --help                   Show this message and exit.
 ```
 
 ## table_name
diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md
index 0932ddc6..5b34da79 100644
--- a/docs/reference/notebook.md
+++ b/docs/reference/notebook.md
@@ -266,27 +266,34 @@ Create a schema file containing external model schemas.
 
 #### table_diff
 ```
-%table_diff [--on [ON ...]] [--model MODEL] [--where WHERE]
-                  [--limit LIMIT] [--show-sample] [--skip-grain-check]
-                  SOURCE:TARGET
+%table_diff [--on [ON ...]] [--skip-columns [SKIP_COLUMNS ...]]
+                [--model MODEL] [--where WHERE] [--limit LIMIT]
+                [--show-sample] [--decimals DECIMALS] [--skip-grain-check]
+                SOURCE:TARGET
 
 Show the diff between two tables.
 
 Can either be two tables or two environments and a model.
 
 positional arguments:
-  SOURCE:TARGET    Source and target in `SOURCE:TARGET` format
+  SOURCE:TARGET         Source and target in `SOURCE:TARGET` format
 
 options:
-  --on <[ON ...]>     The column to join on. Can be specified multiple times. The
-                      model grain will be used if not specified.
-  --model MODEL       The model to diff against when source and target are
-                      environments and not tables.
-  --where WHERE       An optional where statement to filter results.
-  --limit LIMIT       The limit of the sample dataframe.
-  --show-sample       Show a sample of the rows that differ. With many columns,
-                      the output can be very wide.  
-  --skip-grain-check  Disable the check for a primary key (grain) that is missing or is not unique.
+  --on <[ON ...]>       The column to join on. Can be specified multiple
+                        times. The model grain will be used if not specified.
+  --skip-columns <[SKIP_COLUMNS ...]>
+                        The column(s) to skip when comparing the source and
+                        target table.
+  --model MODEL         The model to diff against when source and target are
+                        environments and not tables.
+  --where WHERE         An optional where statement to filter results.
+  --limit LIMIT         The limit of the sample dataframe.
+  --show-sample         Show a sample of the rows that differ. With many
+                        columns, the output can be very wide.
+  --decimals DECIMALS   The number of decimal places to keep when comparing
+                        floating point columns. Default: 3
+  --skip-grain-check    Disable the check for a primary key (grain) that is
+                        missing or is not unique.
 ```
 
 #### model
diff --git a/docs/reference/overview.md b/docs/reference/overview.md
index c48f886d..a10d10ca 100644
--- a/docs/reference/overview.md
+++ b/docs/reference/overview.md
@@ -27,7 +27,7 @@ Runs all audits.
 Read more about [auditing](../concepts/audits.md).
 
 ## format
-Formats all SQL model files in place.
+Formats all SQL model and audit files in place.
 
 ## diff
 Shows the diff between the local model and a model in an environment.
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 64016034..4bc28920 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -259,7 +259,7 @@ def evaluate(
 @error_handler
 @cli_analytics
 def format(ctx: click.Context, **kwargs: t.Any) -> None:
-    """Format all SQL models."""
+    """Format all SQL models and audits."""
     ctx.obj.format(**{k: v for k, v in kwargs.items() if v is not None})
 
 
@@ -720,6 +720,13 @@ def create_external_models(obj: Context) -> None:
     multiple=True,
     help="The column to join on. Can be specified multiple times. The model grain will be used if not specified.",
 )
+@click.option(
+    "-s",
+    "--skip-columns",
+    type=str,
+    multiple=True,
+    help="The column(s) to skip when comparing the source and target table.",
+)
 @click.option(
     "--where",
     type=str,
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 3546d743..14783d00 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -858,13 +858,14 @@ class GenericContext(BaseContext, t.Generic[C]):
         append_newline: t.Optional[bool] = None,
         **kwargs: t.Any,
     ) -> None:
-        """Format all SQL models."""
-        for model in self._models.values():
-            if not model._path.suffix == ".sql":
+        """Format all SQL models and audits."""
+        format_targets = {**self._models, **self._audits}
+        for target in format_targets.values():
+            if not target._path.suffix == ".sql":
                 continue
-            with open(model._path, "r+", encoding="utf-8") as file:
+            with open(target._path, "r+", encoding="utf-8") as file:
                 expressions = parse(
-                    file.read(), default_dialect=self.config_for_node(model).dialect
+                    file.read(), default_dialect=self.config_for_node(target).dialect
                 )
                 if transpile:
                     for prop in expressions[0].expressions:
@@ -872,14 +873,14 @@ class GenericContext(BaseContext, t.Generic[C]):
                             prop.replace(
                                 exp.Property(
                                     this="dialect",
-                                    value=exp.Literal.string(transpile or model.dialect),
+                                    value=exp.Literal.string(transpile or target.dialect),
                                 )
                             )
-                format = self.config_for_node(model).format
+                format = self.config_for_node(target).format
                 opts = {**format.generator_options, **kwargs}
                 file.seek(0)
                 file.write(
-                    format_model_expressions(expressions, transpile or model.dialect, **opts)
+                    format_model_expressions(expressions, transpile or target.dialect, **opts)
                 )
                 if append_newline is None:
                     append_newline = format.append_newline
@@ -1272,6 +1273,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         source: str,
         target: str,
         on: t.List[str] | exp.Condition | None = None,
+        skip_columns: t.List[str] | None = None,
         model_or_snapshot: t.Optional[ModelOrSnapshot] = None,
         where: t.Optional[str | exp.Condition] = None,
         limit: int = 20,
@@ -1287,6 +1289,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             target: The target environment or table.
             on: The join condition, table aliases must be "s" and "t" for source and target.
                 If omitted, the table's grain will be used.
+            skip_columns: The columns to skip when computing the table diff.
             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.
@@ -1340,6 +1343,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             source=source,
             target=target,
             on=on,
+            skip_columns=skip_columns,
             where=where,
             source_alias=source_alias,
             target_alias=target_alias,
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index c9e6abd6..33bd240b 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -804,7 +804,7 @@ def extend_sqlglot() -> None:
         if MacroFunc not in generator.TRANSFORMS:
             generator.TRANSFORMS.update(
                 {
-                    Audit: lambda self, e: _sqlmesh_ddl_sql(self, e, "Audit"),
+                    Audit: lambda self, e: _sqlmesh_ddl_sql(self, e, "AUDIT"),
                     DColonCast: lambda self, e: f"{self.sql(e, 'this')}::{self.sql(e, 'to')}",
                     Jinja: lambda self, e: e.name,
                     JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 9fb98266..2a686fbb 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -5,6 +5,7 @@ import typing as t
 
 import pandas as pd
 from sqlglot import exp, parse_one
+from sqlglot.helper import ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
@@ -139,6 +140,7 @@ class TableDiff:
         source: TableName,
         target: TableName,
         on: t.List[str] | exp.Condition,
+        skip_columns: t.List[str] | None = None,
         where: t.Optional[str | exp.Condition] = None,
         limit: int = 20,
         source_alias: t.Optional[str] = None,
@@ -179,6 +181,14 @@ class TableDiff:
         else:
             self.on = on
 
+        self.skip_columns = {
+            normalize_identifiers(
+                exp.parse_identifier(t.cast(str, col)),
+                dialect=self.model_dialect or self.dialect,
+            ).name
+            for col in ensure_list(skip_columns)
+        }
+
         normalize_identifiers(self.on, dialect=self.model_dialect or self.dialect)
 
         self._source_schema: t.Optional[t.Dict[str, exp.DataType]] = None
@@ -210,8 +220,15 @@ class TableDiff:
 
     def row_diff(self, skip_grain_check: bool = False) -> RowDiff:
         if self._row_diff is None:
-            s_selects = {c: exp.column(c, "s").as_(f"s__{c}") for c in self.source_schema}
-            t_selects = {c: exp.column(c, "t").as_(f"t__{c}") for c in self.target_schema}
+            source_schema = {
+                c: t for c, t in self.source_schema.items() if c not in self.skip_columns
+            }
+            target_schema = {
+                c: t for c, t in self.target_schema.items() if c not in self.skip_columns
+            }
+
+            s_selects = {c: exp.column(c, "s").as_(f"s__{c}") for c in source_schema}
+            t_selects = {c: exp.column(c, "t").as_(f"t__{c}") for c in target_schema}
 
             index_cols = []
             s_index = []
@@ -227,9 +244,7 @@ class TableDiff:
             s_index = list(dict.fromkeys(s_index))
             t_index = list(dict.fromkeys(t_index))
 
-            matched_columns = {
-                c: t for c, t in self.source_schema.items() if t == self.target_schema.get(c)
-            }
+            matched_columns = {c: t for c, t in source_schema.items() if t == target_schema.get(c)}
 
             def _column_expr(name: str, table: str) -> exp.Expression:
                 if matched_columns[name].this in exp.DataType.FLOAT_TYPES:
@@ -435,7 +450,7 @@ class TableDiff:
                 s_sample = sample[(sample["s_exists"] == 1) & (sample["row_joined"] == 0)][
                     [
                         *[f"s__{c}" for c in index_cols],
-                        *[f"s__{c}" for c in self.source_schema if c not in index_cols],
+                        *[f"s__{c}" for c in source_schema if c not in index_cols],
                     ]
                 ]
                 s_sample.rename(
@@ -445,7 +460,7 @@ class TableDiff:
                 t_sample = sample[(sample["t_exists"] == 1) & (sample["row_joined"] == 0)][
                     [
                         *[f"t__{c}" for c in index_cols],
-                        *[f"t__{c}" for c in self.target_schema if c not in index_cols],
+                        *[f"t__{c}" for c in target_schema if c not in index_cols],
                     ]
                 ]
                 t_sample.rename(
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 8d5eb692..e1c1c801 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -577,6 +577,12 @@ class SQLMeshMagics(Magics):
         nargs="*",
         help="The column to join on. Can be specified multiple times. The model grain will be used if not specified.",
     )
+    @argument(
+        "--skip-columns",
+        type=str,
+        nargs="*",
+        help="The column(s) to skip when comparing the source and target table.",
+    )
     @argument(
         "--model",
         type=str,
@@ -622,6 +628,7 @@ class SQLMeshMagics(Magics):
             source=source,
             target=target,
             on=args.on,
+            skip_columns=args.skip_columns,
             model_or_snapshot=args.model,
             where=args.where,
             limit=args.limit,
@@ -724,7 +731,7 @@ class SQLMeshMagics(Magics):
     @line_magic
     @pass_sqlmesh_context
     def format(self, context: Context, line: str) -> None:
-        """Format all SQL models."""
+        """Format all SQL models and audits."""
         args = parse_argstring(self.format, line)
         context.format(**{k: v for k, v in vars(args).items() if v is not None})
 


Test patch:
diff --git a/tests/core/test_audit.py b/tests/core/test_audit.py
index 03877c03..9ab97a43 100644
--- a/tests/core/test_audit.py
+++ b/tests/core/test_audit.py
@@ -670,7 +670,7 @@ def test_text_diff():
 
 @@ -1,5 +1,5 @@
 
- Audit (
+ AUDIT (
 -  name my_audit,
 +  name my_audit_2,
    dialect spark,
diff --git a/tests/core/test_format.py b/tests/core/test_format.py
index 60b0cb86..907e0b4f 100644
--- a/tests/core/test_format.py
+++ b/tests/core/test_format.py
@@ -3,12 +3,14 @@ import pathlib
 from sqlmesh.core.config import Config
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import parse
+from sqlmesh.core.audit import ModelAudit
 from sqlmesh.core.model import SqlModel, load_sql_based_model
 from tests.utils.test_filesystem import create_temp_file
 
 
 def test_format_files(tmp_path: pathlib.Path):
     models_dir = pathlib.Path("models")
+    audits_dir = pathlib.Path("audits")
 
     f1 = create_temp_file(
         tmp_path,
@@ -20,14 +22,36 @@ def test_format_files(tmp_path: pathlib.Path):
         pathlib.Path(models_dir, "model_2.sql"),
         "MODEL(name other.model); SELECT 2 AS another_column",
     )
+    f3 = create_temp_file(
+        tmp_path,
+        pathlib.Path(audits_dir, "audit_1.sql"),
+        "AUDIT(name assert_positive_id, dialect 'duckdb'); SELECT  * FROM @this_model WHERE  \"CaseSensitive\"_item_id < 0;",
+    )
+    f4 = create_temp_file(
+        tmp_path,
+        pathlib.Path(models_dir, "model_3.sql"),
+        "MODEL(name audit.model, audits (inline_audit)); SELECT 3 AS item_id; AUDIT(name inline_audit); SELECT  * FROM @this_model WHERE  item_id < 0;",
+    )
 
     config = Config()
     context = Context(paths=tmp_path, config=config)
     context.load()
+
     assert isinstance(context.get_model("this.model"), SqlModel)
     assert isinstance(context.get_model("other.model"), SqlModel)
+    assert isinstance(context.get_model("audit.model"), SqlModel)
+    assert isinstance(context._audits["assert_positive_id"], ModelAudit)
     assert context.get_model("this.model").query.sql() == 'SELECT 1 AS "CaseSensitive"'  # type: ignore
     assert context.get_model("other.model").query.sql() == "SELECT 2 AS another_column"  # type: ignore
+    assert context.get_model("audit.model").query.sql() == "SELECT 3 AS item_id"  # type: ignore
+    assert (
+        context.get_model("audit.model").inline_audits["inline_audit"].query.sql()
+        == "SELECT * FROM @this_model WHERE item_id < 0"
+    )
+    assert (
+        context._audits["assert_positive_id"].query.sql()
+        == 'SELECT * FROM @this_model WHERE "CaseSensitive_item_id" < 0'
+    )
 
     # Transpile project to BigQuery
     context.format(transpile="bigquery")
@@ -44,3 +68,17 @@ def test_format_files(tmp_path: pathlib.Path):
     # Ensure no dialect is added if it's not needed
     upd2 = f2.read_text(encoding="utf-8")
     assert upd2 == "MODEL (\n  name other.model\n);\n\nSELECT\n  2 AS another_column"
+
+    # Ensure audit specific dialect is updated and formatting
+    upd3 = f3.read_text(encoding="utf-8")
+    assert (
+        upd3
+        == "AUDIT (\n  name assert_positive_id,\n  dialect 'bigquery'\n);\n\nSELECT\n  *\nFROM @this_model\nWHERE\n  `CaseSensitive_item_id` < 0"
+    )
+
+    # Ensure inline audit is formatted within model definition
+    upd4 = f4.read_text(encoding="utf-8")
+    assert (
+        upd4
+        == "MODEL (\n  name audit.model,\n  audits (\n    inline_audit\n  )\n);\n\nSELECT\n  3 AS item_id;\n\nAUDIT (\n  name inline_audit\n);\n\nSELECT\n  *\nFROM @this_model\nWHERE\n  item_id < 0"
+    )
diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py
index 42fc716f..fa4175be 100644
--- a/tests/core/test_table_diff.py
+++ b/tests/core/test_table_diff.py
@@ -230,6 +230,7 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
             {
                 "key": [1, 2, 3],
                 "value": [1.0, 4.2, 4.1],
+                "ignored": [1, 2, 3],
             }
         ),
     )
@@ -240,6 +241,7 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
             {
                 "key": [1, 2, 6],
                 "value": [1.0, 3.0, -2.2],
+                "ignored": [1, 2, 3],
             }
         ),
     )
@@ -255,6 +257,7 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
         source="table_diff_source",
         target="table_diff_target",
         on=["key"],
+        skip_columns=["ignored"],
     )
 
     spy_execute.assert_any_call(query_sql)