From e2671f3fc63bbeb5fc11ae0ec88518371404edf5 Mon Sep 17 00:00:00 2001 From: Anant <75747269+Anant-gif@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:36:47 +0530 Subject: [PATCH] fix(table_diff): support key columns stored with non-normalized casing Signed-off-by: Anant <75747269+Anant-gif@users.noreply.github.com> --- sqlmesh/core/table_diff.py | 18 +++- tests/core/test_table_diff.py | 160 ++++++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 3 deletions(-) diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py index 97cb0c19ba..a4b3758cde 100644 --- a/sqlmesh/core/table_diff.py +++ b/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,9 +297,12 @@ 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)) @@ -305,6 +310,13 @@ def key_columns(self) -> t.Tuple[t.List[exp.Column], t.List[exp.Column], t.List[ 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 diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py index c2e293e4c2..cab237bca1 100644 --- a/tests/core/test_table_diff.py +++ b/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()