Skip to content

Commit 63774eb

Browse files
kvlongekvlonge
committed
fix: resolve Snowflake output aliases
1 parent 02a3f98 commit 63774eb

2 files changed

Lines changed: 182 additions & 0 deletions

File tree

‎src/sqlbuild/compiler/sql_analysis/_helpers/schema_validation.py‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
"E223": "B005",
2525
}
2626
_ERROR_SEVERITY: str = "error"
27+
_NATIVE_UNKNOWN_COLUMN_CODE: str = "E201"
28+
_NATIVE_AMBIGUOUS_COLUMN_CODE: str = "E221"
2729
_ORDER_BY_CONTEXT: str = "ORDER BY"
2830
_WINDOW_ORDER_BY_CONTEXT: str = "WINDOW ORDER BY"
2931
_SNOWFLAKE_DIALECT: str = "snowflake"
@@ -32,6 +34,12 @@
3234
r"^Unknown column '(column([1-9][0-9]*))'$",
3335
flags=re.IGNORECASE,
3436
)
37+
_UNKNOWN_COLUMN_PATTERN: re.Pattern[str] = re.compile(
38+
r"^Unknown column '([^']+)'(?: in table '[^']+'| \(not found in any referenced table\))$"
39+
)
40+
_AMBIGUOUS_UNQUALIFIED_COLUMN_PATTERN: re.Pattern[str] = re.compile(
41+
r"^Ambiguous unqualified column '([^']+)' found in [0-9]+ referenced tables$"
42+
)
3543

3644

3745
def get_schema_validations(
@@ -112,6 +120,11 @@ def _binding_result(*, sql: str, dialect: str | None, response: object) -> SqlBi
112120
message=raw_message,
113121
line=line,
114122
column=column,
123+
) or _is_proven_snowflake_output_alias(
124+
sql=sql,
125+
dialect=dialect,
126+
native_code=native_code,
127+
message=raw_message,
115128
):
116129
continue
117130
diagnostics.append(
@@ -259,6 +272,127 @@ def _projection_is_implicit_values_column(
259272
return (candidate_line, candidate_column) == (line, column)
260273

261274

275+
def _is_proven_snowflake_output_alias(
276+
*, sql: str, dialect: str | None, native_code: str, message: str
277+
) -> bool:
278+
if (dialect or "").lower() != _SNOWFLAKE_DIALECT:
279+
return False
280+
column_name: str | None = _output_alias_diagnostic_column(
281+
native_code=native_code,
282+
message=message,
283+
)
284+
if column_name is None:
285+
return False
286+
polyglot_module: Any = import_polyglot_sql()
287+
try:
288+
parsed: Any = polyglot_module.parse_one(sql, dialect=_SNOWFLAKE_DIALECT)
289+
except polyglot_module.PolyglotError:
290+
return False
291+
selects: list[Any] = []
292+
if str(getattr(parsed, "kind", "")) == _SELECT_KIND:
293+
selects.append(parsed)
294+
selects.extend(parsed.find_all(_SELECT_KIND))
295+
return any(
296+
_select_proves_output_alias(
297+
select=select,
298+
column_name=column_name,
299+
native_code=native_code,
300+
)
301+
for select in selects
302+
)
303+
304+
305+
def _output_alias_diagnostic_column(*, native_code: str, message: str) -> str | None:
306+
pattern: re.Pattern[str]
307+
if native_code == _NATIVE_UNKNOWN_COLUMN_CODE:
308+
pattern = _UNKNOWN_COLUMN_PATTERN
309+
elif native_code == _NATIVE_AMBIGUOUS_COLUMN_CODE:
310+
pattern = _AMBIGUOUS_UNQUALIFIED_COLUMN_PATTERN
311+
else:
312+
return None
313+
match: re.Match[str] | None = pattern.fullmatch(message)
314+
return match.group(1) if match is not None else None
315+
316+
317+
def _select_proves_output_alias(*, select: Any, column_name: str, native_code: str) -> bool:
318+
payload: object = select.to_dict().get(_SELECT_KIND)
319+
if not isinstance(payload, dict):
320+
return False
321+
select_payload: dict[str, object] = cast(dict[str, object], payload)
322+
projections: object = select_payload.get("expressions")
323+
if not isinstance(projections, list):
324+
return False
325+
aliases: set[str] = set()
326+
for projection in projections:
327+
expression, alias = _projection_expression_and_alias(projection)
328+
if native_code == _NATIVE_UNKNOWN_COLUMN_CODE and column_name.lower() in aliases:
329+
if _payload_contains_unqualified_column(
330+
payload=expression,
331+
column_name=column_name,
332+
):
333+
return True
334+
if alias is not None:
335+
aliases.add(alias.lower())
336+
if column_name.lower() not in aliases:
337+
return False
338+
if native_code == _NATIVE_UNKNOWN_COLUMN_CODE:
339+
return _payload_contains_unqualified_column(
340+
payload=select_payload.get("where_clause"),
341+
column_name=column_name,
342+
)
343+
return _payload_contains_unqualified_column(
344+
payload=select_payload.get("order_by"),
345+
column_name=column_name,
346+
)
347+
348+
349+
def _projection_expression_and_alias(projection: object) -> tuple[object, str | None]:
350+
if not isinstance(projection, dict):
351+
return projection, None
352+
alias_payload: object = cast(dict[str, object], projection).get("alias")
353+
if not isinstance(alias_payload, dict):
354+
return projection, None
355+
alias_dict: dict[str, object] = cast(dict[str, object], alias_payload)
356+
name_payload: object = alias_dict.get("alias")
357+
if not isinstance(name_payload, dict):
358+
return alias_dict.get("this"), None
359+
name_dict: dict[str, object] = cast(dict[str, object], name_payload)
360+
alias: str | None = (
361+
str(name_dict.get("name"))
362+
if name_dict.get("quoted") is not True and name_dict.get("name") is not None
363+
else None
364+
)
365+
return alias_dict.get("this"), alias
366+
367+
368+
def _payload_contains_unqualified_column(*, payload: object, column_name: str) -> bool:
369+
if isinstance(payload, list | tuple):
370+
return any(
371+
_payload_contains_unqualified_column(payload=value, column_name=column_name)
372+
for value in payload
373+
)
374+
if not isinstance(payload, dict):
375+
return False
376+
payload_dict: dict[str, object] = cast(dict[str, object], payload)
377+
column_payload: object = payload_dict.get("column")
378+
if isinstance(column_payload, dict):
379+
column_dict: dict[str, object] = cast(dict[str, object], column_payload)
380+
name_payload: object = column_dict.get("name")
381+
if isinstance(name_payload, dict):
382+
name_dict: dict[str, object] = cast(dict[str, object], name_payload)
383+
if (
384+
column_dict.get("table") is None
385+
and name_dict.get("quoted") is not True
386+
and str(name_dict.get("name") or "").lower() == column_name.lower()
387+
):
388+
return True
389+
return any(
390+
key != _SELECT_KIND
391+
and _payload_contains_unqualified_column(payload=value, column_name=column_name)
392+
for key, value in payload_dict.items()
393+
)
394+
395+
262396
def _optional_int(value: object) -> int | None:
263397
return value if isinstance(value, int) else None
264398

‎tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_schema_validation_scopes.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,54 @@
133133
dialects=("snowflake",),
134134
expected_diagnostic_count=1,
135135
),
136+
SchemaValidationScopeTestCase(
137+
description="snowflake order by resolves an output alias before ambiguous inputs",
138+
query_sql=(
139+
"SELECT COALESCE(placed.customer_id, returned.customer_id) AS customer_id "
140+
"FROM placed_orders placed FULL JOIN returned_orders returned "
141+
"ON placed.customer_id = returned.customer_id ORDER BY customer_id"
142+
),
143+
schema={
144+
"placed_orders": {"customer_id": "integer"},
145+
"returned_orders": {"customer_id": "integer"},
146+
},
147+
dialects=("snowflake",),
148+
expected_diagnostic_count=0,
149+
),
150+
SchemaValidationScopeTestCase(
151+
description="snowflake later projection resolves an earlier output alias",
152+
query_sql=(
153+
"SELECT amount + 1 AS adjusted_amount, adjusted_amount * 2 AS doubled_amount "
154+
"FROM orders"
155+
),
156+
schema={"orders": {"amount": "integer"}},
157+
dialects=("snowflake",),
158+
expected_diagnostic_count=0,
159+
),
160+
SchemaValidationScopeTestCase(
161+
description="snowflake where resolves an output alias",
162+
query_sql="SELECT amount > 0 AS is_positive FROM orders WHERE is_positive = TRUE",
163+
schema={"orders": {"amount": "integer"}},
164+
dialects=("snowflake",),
165+
expected_diagnostic_count=0,
166+
),
167+
SchemaValidationScopeTestCase(
168+
description="snowflake forward output alias reference remains unknown",
169+
query_sql=(
170+
"SELECT adjusted_amount * 2 AS doubled_amount, "
171+
"amount + 1 AS adjusted_amount FROM orders"
172+
),
173+
schema={"orders": {"amount": "integer"}},
174+
dialects=("snowflake",),
175+
expected_diagnostic_count=1,
176+
),
177+
SchemaValidationScopeTestCase(
178+
description="snowflake unrelated unknown column remains unknown",
179+
query_sql="SELECT amount AS order_amount FROM orders WHERE missing_amount > 0",
180+
schema={"orders": {"amount": "integer"}},
181+
dialects=("snowflake",),
182+
expected_diagnostic_count=1,
183+
),
136184
],
137185
ids=lambda case: case.description,
138186
)

0 commit comments

Comments
 (0)