fix(table_diff): support key columns stored with non-normalized casing

Signed-off-by: Anant <75747269+Anant-gif@users.noreply.github.com>

Author: Anant-gif

sqlmesh/core/table_diff.py

@@ -282,9 +282,11 @@ def key_columns(self) -> t.Tuple[t.List[exp.Column], t.List[exp.Column], t.List[
         # If the columns to join on are explicitly specified, then just return them
         if isinstance(self._on, (list, tuple)):
             identifiers = [normalize_identifiers(c, dialect=dialect) for c in self._on]
-            s_index = [exp.column(c, "s") for c in identifiers]
-            t_index = [exp.column(c, "t") for c in identifiers]
-            return s_index, t_index, [i.name for i in identifiers]
+            s_names = [self._resolve_column_name(c.name, self.source_schema) for c in identifiers]
+            t_names = [self._resolve_column_name(c.name, self.target_schema) for c in identifiers]
+            s_index = [exp.column(c, "s") for c in s_names]
+            t_index = [exp.column(c, "t") for c in t_names]
+            return s_index, t_index, s_names
 
         # Otherwise, we need to parse them out of the supplied "on" condition
         index_cols = []
@@ -295,16 +297,26 @@ def key_columns(self) -> t.Tuple[t.List[exp.Column], t.List[exp.Column], t.List[
         for col in self._on.find_all(exp.Column):
             index_cols.append(col.name)
             if col.table.lower() == "s":
+                col = exp.column(self._resolve_column_name(col.name, self.source_schema), col.table)
                 s_index.append(col)
             elif col.table.lower() == "t":
+                col = exp.column(self._resolve_column_name(col.name, self.target_schema), col.table)
                 t_index.append(col)
+            index_cols.append(col.name)
 
         index_cols = list(dict.fromkeys(index_cols))
         s_index = list(dict.fromkeys(s_index))
         t_index = list(dict.fromkeys(t_index))
 
         return s_index, t_index, index_cols
 
+    def _resolve_column_name(self, name: str, schema: t.Dict[str, exp.DataType]) -> str:
+        if name in schema:
+            return name
+
+        lowercase_name = name.lower()
+        return next((c for c in schema if c.lower() == lowercase_name), name)
+
     @property
     def source_key_expression(self) -> exp.Expr:
         s_index, _, _ = self.key_columns

tests/core/test_table_diff.py

@@ -1197,6 +1197,166 @@ def test_data_diff_sample_limit():
     assert len(diff.joined_sample) == 3
 
 
+def test_data_diff_non_lowercase_key_columns():
+    engine_adapter = DuckDBConnectionConfig().create_engine_adapter()
+
+    columns_to_types = {
+        "KEY1": exp.DataType.build("int"),
+        "Key2": exp.DataType.build("varchar"),
+        "VALUE": exp.DataType.build("varchar"),
+    }
+
+    engine_adapter.create_table("src", columns_to_types)
+    engine_adapter.create_table("target", columns_to_types)
+
+    src_records = [
+        (1, "a", "value"),
+        (2, "b", "source"),
+        (3, "c", "source only"),
+    ]
+
+    target_records = [
+        (1, "a", "value"),
+        (2, "b", "target"),
+        (4, "d", "target only"),
+    ]
+
+    src_df = pd.DataFrame(data=src_records, columns=columns_to_types.keys())
+    target_df = pd.DataFrame(data=target_records, columns=columns_to_types.keys())
+
+    engine_adapter.insert_append("src", src_df)
+    engine_adapter.insert_append("target", target_df)
+
+    # casing of the supplied key should not matter
+    for on in (["KEY1", "Key2"], ["key1", "KEY2"]):
+        table_diff = TableDiff(adapter=engine_adapter, source="src", target="target", on=on)
+
+        _, _, col_names = table_diff.key_columns
+        assert col_names == ["KEY1", "Key2"]
+
+        diff = table_diff.row_diff()
+
+        assert diff.join_count == 2
+        assert diff.full_match_count == 1
+        assert diff.partial_match_count == 1
+        assert diff.s_only_count == 1
+        assert diff.t_only_count == 1
+
+        assert diff.s_sample["VALUE"].tolist() == ["source only"]
+        assert diff.t_sample["VALUE"].tolist() == ["target only"]
+        assert diff.joined_sample[["s_VALUE", "t_VALUE"]].values.flatten().tolist() == [
+            "source",
+            "target",
+        ]
+
+    table_diff = TableDiff(adapter=engine_adapter, source="src", target="target", on=["KEY1"])
+
+    _, _, col_names = table_diff.key_columns
+    assert col_names == ["KEY1"]
+
+    diff = table_diff.row_diff()
+
+    assert diff.join_count == 2
+    assert diff.full_match_count == 1
+    assert diff.partial_match_count == 1
+    assert diff.s_only_count == 1
+    assert diff.t_only_count == 1
+
+
+def test_data_diff_key_columns_with_differing_case_between_source_and_target():
+    engine_adapter = DuckDBConnectionConfig().create_engine_adapter()
+
+    source_columns_to_types = {
+        "KEY1": exp.DataType.build("int"),
+        "Key2": exp.DataType.build("varchar"),
+        "value": exp.DataType.build("varchar"),
+    }
+    target_columns_to_types = {
+        "key1": exp.DataType.build("int"),
+        "KEY2": exp.DataType.build("varchar"),
+        "value": exp.DataType.build("varchar"),
+    }
+
+    engine_adapter.create_table("src", source_columns_to_types)
+    engine_adapter.create_table("target", target_columns_to_types)
+
+    engine_adapter.insert_append(
+        "src",
+        pd.DataFrame(
+            data=[(1, "a", "value"), (2, "b", "source")],
+            columns=source_columns_to_types.keys(),
+        ),
+    )
+    engine_adapter.insert_append(
+        "target",
+        pd.DataFrame(
+            data=[(1, "a", "value"), (2, "b", "target")],
+            columns=target_columns_to_types.keys(),
+        ),
+    )
+
+    table_diff = TableDiff(
+        adapter=engine_adapter, source="src", target="target", on=["key1", "KEY2"]
+    )
+
+    s_index, t_index, col_names = table_diff.key_columns
+    assert [c.sql() for c in s_index] == ['"s.KEY1"', '"s.Key2"']
+    assert [c.sql() for c in t_index] == ['"t.key1"', '"t.KEY2"']
+    assert col_names == ["KEY1", "Key2"]
+
+    diff = table_diff.row_diff()
+
+    assert diff.join_count == 2
+    assert diff.full_match_count == 1
+    assert diff.partial_match_count == 1
+    assert diff.s_only_count == 0
+    assert diff.t_only_count == 0
+
+    # the key columns are excluded from the per column match stats
+    assert diff.column_stats.index.tolist() == ["value"]
+
+
+def test_data_diff_non_lowercase_key_columns_in_on_condition():
+    engine_adapter = DuckDBConnectionConfig().create_engine_adapter()
+
+    columns_to_types = {
+        "KEY1": exp.DataType.build("int"),
+        "Key2": exp.DataType.build("varchar"),
+        "VALUE": exp.DataType.build("varchar"),
+    }
+
+    engine_adapter.create_table("src", columns_to_types)
+    engine_adapter.create_table("target", columns_to_types)
+
+    src_df = pd.DataFrame(
+        data=[(1, "a", "value"), (2, "b", "source")], columns=columns_to_types.keys()
+    )
+    target_df = pd.DataFrame(
+        data=[(1, "a", "value"), (2, "b", "target")], columns=columns_to_types.keys()
+    )
+
+    engine_adapter.insert_append("src", src_df)
+    engine_adapter.insert_append("target", target_df)
+
+    table_diff = TableDiff(
+        adapter=engine_adapter,
+        source="src",
+        target="target",
+        on=exp.condition('s."KEY1" = t."KEY1" AND s."Key2" = t."Key2"'),
+    )
+
+    _, col_names = table_diff.key_columns
+    assert col_names == ["KEY1", "Key2"]
+
+    diff = table_diff.row_diff()
+
+    assert diff.join_count == 2
+    assert diff.full_match_count == 1
+    assert diff.partial_match_count == 1
+    assert diff.s_only_count == 0
+    assert diff.t_only_count == 0
+
+
 def test_data_diff_nulls_in_some_grain_columns():
     engine_adapter = DuckDBConnectionConfig().create_engine_adapter()