Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 38 additions & 11 deletions sqlmesh/core/table_diff.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -275,16 +268,50 @@ 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

# 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 = []
Expand Down
78 changes: 78 additions & 0 deletions tests/core/test_table_diff.py
Original file line number Diff line number Diff line change
Expand Up @@ -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