From 63774eb14529610cd5ccda2b155efc480b6e069a Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:55:53 +0100 Subject: [PATCH] fix: resolve Snowflake output aliases --- .../_helpers/schema_validation.py | 134 ++++++++++++++++++ .../main/test_schema_validation_scopes.py | 48 +++++++ 2 files changed, 182 insertions(+) diff --git a/src/sqlbuild/compiler/sql_analysis/_helpers/schema_validation.py b/src/sqlbuild/compiler/sql_analysis/_helpers/schema_validation.py index 3c5eb8dd86..e37d8ac421 100644 --- a/src/sqlbuild/compiler/sql_analysis/_helpers/schema_validation.py +++ b/src/sqlbuild/compiler/sql_analysis/_helpers/schema_validation.py @@ -24,6 +24,8 @@ "E223": "B005", } _ERROR_SEVERITY: str = "error" +_NATIVE_UNKNOWN_COLUMN_CODE: str = "E201" +_NATIVE_AMBIGUOUS_COLUMN_CODE: str = "E221" _ORDER_BY_CONTEXT: str = "ORDER BY" _WINDOW_ORDER_BY_CONTEXT: str = "WINDOW ORDER BY" _SNOWFLAKE_DIALECT: str = "snowflake" @@ -32,6 +34,12 @@ r"^Unknown column '(column([1-9][0-9]*))'$", flags=re.IGNORECASE, ) +_UNKNOWN_COLUMN_PATTERN: re.Pattern[str] = re.compile( + r"^Unknown column '([^']+)'(?: in table '[^']+'| \(not found in any referenced table\))$" +) +_AMBIGUOUS_UNQUALIFIED_COLUMN_PATTERN: re.Pattern[str] = re.compile( + r"^Ambiguous unqualified column '([^']+)' found in [0-9]+ referenced tables$" +) def get_schema_validations( @@ -112,6 +120,11 @@ def _binding_result(*, sql: str, dialect: str | None, response: object) -> SqlBi message=raw_message, line=line, column=column, + ) or _is_proven_snowflake_output_alias( + sql=sql, + dialect=dialect, + native_code=native_code, + message=raw_message, ): continue diagnostics.append( @@ -259,6 +272,127 @@ def _projection_is_implicit_values_column( return (candidate_line, candidate_column) == (line, column) +def _is_proven_snowflake_output_alias( + *, sql: str, dialect: str | None, native_code: str, message: str +) -> bool: + if (dialect or "").lower() != _SNOWFLAKE_DIALECT: + return False + column_name: str | None = _output_alias_diagnostic_column( + native_code=native_code, + message=message, + ) + if column_name is None: + return False + polyglot_module: Any = import_polyglot_sql() + try: + parsed: Any = polyglot_module.parse_one(sql, dialect=_SNOWFLAKE_DIALECT) + except polyglot_module.PolyglotError: + return False + selects: list[Any] = [] + if str(getattr(parsed, "kind", "")) == _SELECT_KIND: + selects.append(parsed) + selects.extend(parsed.find_all(_SELECT_KIND)) + return any( + _select_proves_output_alias( + select=select, + column_name=column_name, + native_code=native_code, + ) + for select in selects + ) + + +def _output_alias_diagnostic_column(*, native_code: str, message: str) -> str | None: + pattern: re.Pattern[str] + if native_code == _NATIVE_UNKNOWN_COLUMN_CODE: + pattern = _UNKNOWN_COLUMN_PATTERN + elif native_code == _NATIVE_AMBIGUOUS_COLUMN_CODE: + pattern = _AMBIGUOUS_UNQUALIFIED_COLUMN_PATTERN + else: + return None + match: re.Match[str] | None = pattern.fullmatch(message) + return match.group(1) if match is not None else None + + +def _select_proves_output_alias(*, select: Any, column_name: str, native_code: str) -> bool: + payload: object = select.to_dict().get(_SELECT_KIND) + if not isinstance(payload, dict): + return False + select_payload: dict[str, object] = cast(dict[str, object], payload) + projections: object = select_payload.get("expressions") + if not isinstance(projections, list): + return False + aliases: set[str] = set() + for projection in projections: + expression, alias = _projection_expression_and_alias(projection) + if native_code == _NATIVE_UNKNOWN_COLUMN_CODE and column_name.lower() in aliases: + if _payload_contains_unqualified_column( + payload=expression, + column_name=column_name, + ): + return True + if alias is not None: + aliases.add(alias.lower()) + if column_name.lower() not in aliases: + return False + if native_code == _NATIVE_UNKNOWN_COLUMN_CODE: + return _payload_contains_unqualified_column( + payload=select_payload.get("where_clause"), + column_name=column_name, + ) + return _payload_contains_unqualified_column( + payload=select_payload.get("order_by"), + column_name=column_name, + ) + + +def _projection_expression_and_alias(projection: object) -> tuple[object, str | None]: + if not isinstance(projection, dict): + return projection, None + alias_payload: object = cast(dict[str, object], projection).get("alias") + if not isinstance(alias_payload, dict): + return projection, None + alias_dict: dict[str, object] = cast(dict[str, object], alias_payload) + name_payload: object = alias_dict.get("alias") + if not isinstance(name_payload, dict): + return alias_dict.get("this"), None + name_dict: dict[str, object] = cast(dict[str, object], name_payload) + alias: str | None = ( + str(name_dict.get("name")) + if name_dict.get("quoted") is not True and name_dict.get("name") is not None + else None + ) + return alias_dict.get("this"), alias + + +def _payload_contains_unqualified_column(*, payload: object, column_name: str) -> bool: + if isinstance(payload, list | tuple): + return any( + _payload_contains_unqualified_column(payload=value, column_name=column_name) + for value in payload + ) + if not isinstance(payload, dict): + return False + payload_dict: dict[str, object] = cast(dict[str, object], payload) + column_payload: object = payload_dict.get("column") + if isinstance(column_payload, dict): + column_dict: dict[str, object] = cast(dict[str, object], column_payload) + name_payload: object = column_dict.get("name") + if isinstance(name_payload, dict): + name_dict: dict[str, object] = cast(dict[str, object], name_payload) + if ( + column_dict.get("table") is None + and name_dict.get("quoted") is not True + and str(name_dict.get("name") or "").lower() == column_name.lower() + ): + return True + return any( + key != _SELECT_KIND + and _payload_contains_unqualified_column(payload=value, column_name=column_name) + for key, value in payload_dict.items() + ) + + def _optional_int(value: object) -> int | None: return value if isinstance(value, int) else None diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_schema_validation_scopes.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_schema_validation_scopes.py index 53704705d0..d24d314a51 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_schema_validation_scopes.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_schema_validation_scopes.py @@ -133,6 +133,54 @@ dialects=("snowflake",), expected_diagnostic_count=1, ), + SchemaValidationScopeTestCase( + description="snowflake order by resolves an output alias before ambiguous inputs", + query_sql=( + "SELECT COALESCE(placed.customer_id, returned.customer_id) AS customer_id " + "FROM placed_orders placed FULL JOIN returned_orders returned " + "ON placed.customer_id = returned.customer_id ORDER BY customer_id" + ), + schema={ + "placed_orders": {"customer_id": "integer"}, + "returned_orders": {"customer_id": "integer"}, + }, + dialects=("snowflake",), + expected_diagnostic_count=0, + ), + SchemaValidationScopeTestCase( + description="snowflake later projection resolves an earlier output alias", + query_sql=( + "SELECT amount + 1 AS adjusted_amount, adjusted_amount * 2 AS doubled_amount " + "FROM orders" + ), + schema={"orders": {"amount": "integer"}}, + dialects=("snowflake",), + expected_diagnostic_count=0, + ), + SchemaValidationScopeTestCase( + description="snowflake where resolves an output alias", + query_sql="SELECT amount > 0 AS is_positive FROM orders WHERE is_positive = TRUE", + schema={"orders": {"amount": "integer"}}, + dialects=("snowflake",), + expected_diagnostic_count=0, + ), + SchemaValidationScopeTestCase( + description="snowflake forward output alias reference remains unknown", + query_sql=( + "SELECT adjusted_amount * 2 AS doubled_amount, " + "amount + 1 AS adjusted_amount FROM orders" + ), + schema={"orders": {"amount": "integer"}}, + dialects=("snowflake",), + expected_diagnostic_count=1, + ), + SchemaValidationScopeTestCase( + description="snowflake unrelated unknown column remains unknown", + query_sql="SELECT amount AS order_amount FROM orders WHERE missing_amount > 0", + schema={"orders": {"amount": "integer"}}, + dialects=("snowflake",), + expected_diagnostic_count=1, + ), ], ids=lambda case: case.description, )