From b54a78995e975c4c6251928e756a3cb2cc43fac4 Mon Sep 17 00:00:00 2001 From: Vaibhav Naik Date: Sun, 20 Sep 2026 11:26:45 -0700 Subject: [PATCH] fix(table_diff): resolve on/skip_columns case-insensitively against schema On engines whose sqlglot dialect lowercases unquoted identifiers (e.g. BigQuery, DuckDB), TableDiff.key_columns normalized `on` column names via normalize_identifiers() and then looked them up verbatim in the schema returned by adapter.columns(), which preserves the table's actual column casing. When the two disagreed, this raised a bare KeyError before any query ran (table_diff.py:326), and in the single-key path caused a similar KeyError deeper in _fetch_sample. `skip_columns` had the same normalize-then-exact-match mismatch, so columns intended to be excluded silently were not. Add TableDiff._resolve_column_name() to fall back to a case-insensitive match against the schema when there's no exact match, and use it when building key_columns (for both the single- and multi-column paths) and when resolving skip_columns. Add regression tests reproducing the reported KeyError for both single- and multi-column `on` lists, plus a case-insensitive skip_columns test. Fixes #6067 Signed-off-by: Vaibhav Naik --- sqlmesh/core/table_diff.py | 49 +++++++++++++++++----- tests/core/test_table_diff.py | 78 +++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 11 deletions(-) diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py index 97cb0c19ba..17dbb00c14 100644 --- a/sqlmesh/core/table_diff.py +++ b/sqlmesh/core/table_diff.py @@ -255,14 +255,7 @@ def __init__( self.source_alias = source_alias self.target_alias = target_alias - cols: t.List[str] = ensure_list(skip_columns) - self.skip_columns = { - normalize_identifiers( - exp.parse_identifier(col), - dialect=self.model_dialect or self.dialect, - ).name - for col in cols - } + self._skip_columns_raw: t.List[str] = ensure_list(skip_columns) self._on = on self._row_diff: t.Optional[RowDiff] = None @@ -275,6 +268,38 @@ def source_schema(self) -> t.Dict[str, exp.DataType]: def target_schema(self) -> t.Dict[str, exp.DataType]: return self.adapter.columns(self.target_table) + @cached_property + def skip_columns(self) -> t.Set[str]: + dialect = self.model_dialect or self.dialect + names = set() + for col in self._skip_columns_raw: + normalized_name = normalize_identifiers(exp.parse_identifier(col), dialect=dialect).name + # Resolve against both schemas (case-insensitively, if needed) so a column is + # skipped even if the two tables disagree on casing, or the engine reports a + # different case than the normalized `skip_columns` name. + names.add(self._resolve_column_name(normalized_name, self.source_schema)) + names.add(self._resolve_column_name(normalized_name, self.target_schema)) + return names + + @staticmethod + def _resolve_column_name(name: str, schema: t.Dict[str, exp.DataType]) -> str: + """Resolves `name` to the corresponding key in `schema`. + + Some dialects (e.g. BigQuery, DuckDB) normalize unquoted identifiers to a + different case than what the engine's `adapter.columns()` reports for the + underlying table (which reflects however the table was actually created). + If there isn't an exact match, fall back to a case-insensitive lookup so + the normalized `on`/`skip_columns` names still resolve to the real column. + """ + if name in schema: + return name + + for actual_name in schema: + if actual_name.lower() == name.lower(): + return actual_name + + return name + @cached_property def key_columns(self) -> t.Tuple[t.List[exp.Column], t.List[exp.Column], t.List[str]]: dialect = self.model_dialect or self.dialect @@ -282,9 +307,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(i.name, self.source_schema) for i in identifiers] + t_names = [self._resolve_column_name(i.name, self.target_schema) for i in identifiers] + s_index = [exp.column(name, "s") for name in s_names] + t_index = [exp.column(name, "t") for name in t_names] + return s_index, t_index, s_names # Otherwise, we need to parse them out of the supplied "on" condition index_cols = [] diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py index c2e293e4c2..5b727c94a0 100644 --- a/tests/core/test_table_diff.py +++ b/tests/core/test_table_diff.py @@ -1246,3 +1246,81 @@ def test_data_diff_nulls_in_some_grain_columns(): "null value", "null value modified", ] + + +def test_data_diff_on_columns_with_non_lowercase_names(): + # On engines whose sqlglot dialect lowercases unquoted identifiers (e.g. BigQuery, DuckDB), + # `on` columns are normalized to lowercase before being looked up in the schema returned by + # `adapter.columns()`, which preserves the original (non-lowercase) casing. This used to raise + # a KeyError instead of resolving case-insensitively (issue #6067). + engine_adapter = DuckDBConnectionConfig().create_engine_adapter() + + columns_to_types = { + "KEY1": exp.DataType.build("int"), + "KEY2": exp.DataType.build("int"), + "VALUE": exp.DataType.build("varchar"), + } + + engine_adapter.create_table("src", columns_to_types) + engine_adapter.create_table("target", columns_to_types) + + src_records = [(1, 1, "a"), (2, 2, "source only")] + target_records = [(1, 1, "a"), (3, 3, "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) + + # multiple key columns, referenced with a case that doesn't match the schema + multi_key_diff = TableDiff( + adapter=engine_adapter, source="src", target="target", on=["key1", "key2"] + ).row_diff() + + assert multi_key_diff.full_match_count == 1 + assert multi_key_diff.s_only_count == 1 + assert multi_key_diff.t_only_count == 1 + + # single key column, referenced with a case that doesn't match the schema + single_key_diff = TableDiff( + adapter=engine_adapter, source="src", target="target", on=["KEY1"] + ).row_diff() + + assert single_key_diff.full_match_count == 1 + assert single_key_diff.s_only_count == 1 + assert single_key_diff.t_only_count == 1 + + +def test_data_diff_skip_columns_with_non_lowercase_names(): + # `skip_columns` goes through the same normalize-then-exact-match lookup as `on`, so it is + # subject to the same casing mismatch on engines that lowercase unquoted identifiers. + engine_adapter = DuckDBConnectionConfig().create_engine_adapter() + + columns_to_types = { + "KEY1": exp.DataType.build("int"), + "IGNORE_ME": 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) + + # IGNORE_ME differs between source and target, but should be excluded from comparison + src_df = pd.DataFrame(data=[(1, "src-only-value", "a")], columns=columns_to_types.keys()) + target_df = pd.DataFrame(data=[(1, "target-only-value", "a")], columns=columns_to_types.keys()) + + engine_adapter.insert_append("src", src_df) + engine_adapter.insert_append("target", target_df) + + diff = TableDiff( + adapter=engine_adapter, + source="src", + target="target", + on=["KEY1"], + skip_columns=["ignore_me"], + ).row_diff() + + assert diff.full_match_count == 1 + assert diff.s_only_count == 0 + assert diff.t_only_count == 0