Fix: normalization issue in table diff (#2872)

Author: themisvaltinos

sqlmesh/core/context.py

@@ -1299,6 +1299,7 @@ def table_diff(
             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)
@@ -1321,7 +1322,13 @@ def table_diff(
             if not on:
                 for ref in model.all_references:
                     if ref.unique:
-                        on = ref.columns
+                        expr = ref.expression
+
+                        if isinstance(expr, exp.Tuple):
+                            on = [key.this.sql() for key in expr.expressions]
+                        else:
+                            # Handle a single Column or Paren expression
+                            on = [expr.this.sql()]
 
         if not on:
             raise SQLMeshError(
@@ -1337,6 +1344,7 @@ def table_diff(
             source_alias=source_alias,
             target_alias=target_alias,
             model_name=model.name if model_or_snapshot else None,
+            model_dialect=model.dialect if model_or_snapshot else None,
             limit=limit,
             decimals=decimals,
         )

sqlmesh/core/table_diff.py

@@ -144,6 +144,7 @@ def __init__(
         source_alias: t.Optional[str] = None,
         target_alias: t.Optional[str] = None,
         model_name: t.Optional[str] = None,
+        model_dialect: t.Optional[str] = None,
         decimals: int = 3,
     ):
         self.adapter = adapter
@@ -153,13 +154,15 @@ def __init__(
         self.where = exp.condition(where, dialect=self.dialect) if where else None
         self.limit = limit
         self.model_name = model_name
+        self.model_dialect = model_dialect
         self.decimals = decimals
 
         # Support environment aliases for diff output improvement in certain cases
         self.source_alias = source_alias
         self.target_alias = target_alias
 
         if isinstance(on, (list, tuple)):
+            join_condition = [exp.parse_identifier(key) for key in on]
             s_table = exp.to_identifier("s", quoted=True)
             t_table = exp.to_identifier("t", quoted=True)
 
@@ -170,13 +173,13 @@ def __init__(
                         exp.column(c, s_table).is_(exp.null())
                         & exp.column(c, t_table).is_(exp.null())
                     )
-                    for c in on
+                    for c in join_condition
                 )
             )
         else:
             self.on = on
 
-        normalize_identifiers(self.on, dialect=self.dialect)
+        normalize_identifiers(self.on, dialect=self.model_dialect or self.dialect)
 
         self._source_schema: t.Optional[t.Dict[str, exp.DataType]] = None
         self._target_schema: t.Optional[t.Dict[str, exp.DataType]] = None
@@ -321,7 +324,7 @@ def name(e: exp.Expression) -> str:
                 .as_("row_full_match"),
             ).from_(query.subquery("stats"))
 
-            query = quote_identifiers(query, dialect=self.dialect)
+            query = quote_identifiers(query, dialect=self.model_dialect or self.dialect)
             temp_table = exp.table_("diff", db="sqlmesh_temp", quoted=True)
 
             with self.adapter.temp_table(query, name=temp_table) as table:

tests/core/test_table_diff.py

@@ -142,10 +142,10 @@ def test_data_diff_decimals(sushi_context_fixed_date):
 def test_grain_check(sushi_context_fixed_date):
     expressions = d.parse(
         """
-        MODEL (name memory.sushi.grain_items, kind full, grain(key_1, key_2));
+        MODEL (name memory.sushi.grain_items, kind full, grain("key_1", KEY_2));
         SELECT
-            key_1,
-            key_2,
+            key_1 as "key_1",
+            KEY_2,
             value,
         FROM
             (VALUES
@@ -156,10 +156,10 @@ def test_grain_check(sushi_context_fixed_date):
                 (1, 2, 2),
                 (4, NULL, 3),
                 (2, 3, 2),
-            ) AS t (key_1,key_2, value)
+            ) AS t (key_1,KEY_2, value)
     """
     )
-    model_s = load_sql_based_model(expressions)
+    model_s = load_sql_based_model(expressions, dialect="snowflake")
     sushi_context_fixed_date.upsert_model(model_s)
     sushi_context_fixed_date.plan(
         "source_dev",
@@ -170,14 +170,14 @@ def test_grain_check(sushi_context_fixed_date):
         end="2023-01-31",
     )
 
-    model = sushi_context_fixed_date.models['"memory"."sushi"."grain_items"']
+    model = sushi_context_fixed_date.models['"MEMORY"."SUSHI"."GRAIN_ITEMS"']
 
     modified_model = model.dict()
     modified_model["query"] = (
         exp.select("*")
         .from_(model.query.subquery())
         .union(
-            "SELECT key_1, key_2, value FROM (VALUES (1, 6, 1),(1, 5, 3),(NULL, 2, 3),) AS t (key_1, key_2, value)"
+            'SELECT key_1 as "key_1", KEY_2, value FROM (VALUES (1, 6, 1),(1, 5, 3),(NULL, 2, 3),) AS t (key_1, KEY_2, value)'
         )
     )
 
@@ -200,8 +200,8 @@ def test_grain_check(sushi_context_fixed_date):
     diff = sushi_context_fixed_date.table_diff(
         source="source_dev",
         target="target_dev",
-        on=["key_1", "key_2"],
-        model_or_snapshot="sushi.grain_items",
+        on=["'key_1'", "key_2"],
+        model_or_snapshot="SUSHI.GRAIN_ITEMS",
         skip_grain_check=False,
     )