Repository: TobikoData/sqlmesh
Pull request #2872 (2024-07-05 14:27:09)

Fix: normalization issue in table diff
This addresses an issue encountered when running table diff, fixes: #2869 
The proposed fix is to use the project dialect for normalization instead of the gateway dialect and also account for quoted names. `test_table_diff` has been adapted to test for both quoted and unquoted identifiers in snowflake.


Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 10de8180..3546d743 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1299,6 +1299,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             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 @@ class GenericContext(BaseContext, t.Generic[C]):
             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 @@ class GenericContext(BaseContext, t.Generic[C]):
             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,
         )
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 44d9afd7..9fb98266 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -144,6 +144,7 @@ class TableDiff:
         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,6 +154,7 @@ class TableDiff:
         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
@@ -160,6 +162,7 @@ class TableDiff:
         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 @@ class TableDiff:
                         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 @@ class TableDiff:
                 .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:


Test patch:
diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py
index 99b933ed..42fc716f 100644
--- a/tests/core/test_table_diff.py
+++ b/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,
     )