-
Notifications
You must be signed in to change notification settings - Fork 146
Fix escaping during check execution #1486
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
95ffb4b
0b987ab
b5ac9d2
9376f75
d88f8de
eda7515
c96ae40
0ffe3ac
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -78,6 +78,9 @@ def _validate_spark_column(value: Any) -> Any: | |
| COLUMN_NORMALIZE_EXPRESSION = re.compile("[^a-zA-Z0-9]+") | ||
| COLUMN_PATTERN = re.compile(r"Column<'(.*?)(?: AS (\w+))?'>$", re.DOTALL) | ||
| INVALID_COLUMN_NAME_PATTERN = re.compile(r"[\s,;{}\(\)\n\t=]+") | ||
| VALID_UNQUOTED_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") | ||
| SQL_EXPRESSION_STRUCTURE_PATTERN = re.compile(r"""[()\[\]{},;'"`*]""") | ||
| SQL_EXPRESSION_OPERATOR_PATTERN = re.compile(r"[-+/%<>=!&|~^]\s|\s[-+/%<>=!&|~^]") | ||
| _UNRESOLVED_PLACEHOLDER_PATTERN = re.compile(r"\{\{[^}]*\}\}") | ||
|
|
||
| # Destructive SQL statement keywords rejected by `is_sql_query_safe`. SELECT is intentionally | ||
|
|
@@ -149,6 +152,8 @@ def get_column_name_or_alias( | |
| Args: | ||
| column: Column, ConnectColumn (if PySpark Connect available), or string representing a column. | ||
| normalize: If True, normalizes the column name (removes special characters, converts to lowercase). | ||
| For string inputs, any surrounding back-quotes the user supplied (e.g. "`Customer Name`") are | ||
| stripped first so quotes do not leak into display names or generated check names. | ||
| allow_simple_expressions_only: If True, raises an error if the column expression is not a simple expression. | ||
| Complex PySpark expressions (e.g., conditionals, arithmetic, or nested transformations), cannot be fully | ||
| reconstructed correctly when converting to string (e.g. F.col("a") + F.lit(1)). | ||
|
|
@@ -161,7 +166,9 @@ def get_column_name_or_alias( | |
| InvalidParameterError: If the column expression is invalid or unsupported. | ||
| """ | ||
| if isinstance(column, str): | ||
| col_str = column | ||
| col_str = unquote_column_name(column) | ||
| if normalize: | ||
| col_str = normalize_col_str(col_str) | ||
| else: | ||
| # Extract the last alias or column name from the PySpark Column string representation. | ||
| # Strip the representation first to guard against trailing whitespace or CRLF line endings | ||
|
|
@@ -308,6 +315,62 @@ def quote_column_name(name: str) -> str: | |
| return f"`{escaped}`" | ||
|
|
||
|
|
||
| def unquote_column_name(name: str) -> str: | ||
| """ | ||
| Removes surrounding back-quotes from a column name, reversing :func:`quote_column_name`. | ||
|
|
||
| A column reference that the user already back-quoted (e.g. "`Customer Name`") is unwrapped to its | ||
| plain form ("Customer Name") for use in display names and messages. Strings that are not a single | ||
| back-quoted identifier (plain names, dotted paths, SQL expressions) are returned unchanged. | ||
|
|
||
| Args: | ||
| name: Column reference provided as a string. | ||
|
|
||
| Returns: | ||
| The column name without surrounding back-quotes. | ||
| """ | ||
| if len(name) >= 2 and name.startswith("`") and name.endswith("`"): | ||
| return name[1:-1].replace("``", "`") | ||
| return name | ||
|
|
||
|
|
||
| def normalize_column_expr(column: str) -> str: | ||
| """ | ||
| Prepares a column reference string for use with ``F.expr``, back-quoting names that require SQL | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nit: docstring uses double-backtick object references, against the repo convention AGENTS.md (Docstrings) says: No backticks around object names — use italics instead (e.g., arg1, column). Backticks cause rendering issues in API docs. This docstring uses |
||
| identifier escaping while leaving SQL expressions untouched. | ||
|
|
||
| Check functions accept a column as either a plain name or a SQL expression string. A plain name that | ||
| is not a valid bare identifier (spaces, dashes, non-ASCII letters, etc., e.g. "Customer Name", | ||
| "gross-margin" or "Päivämäärä") does not parse when passed to *F.expr* and must be back-quoted. A | ||
| SQL expression (e.g. "a + b", "substr(x, 1, 2)", "*") must be passed through as-is. | ||
|
|
||
| Name and expression cannot be told apart with certainty from a string alone. This function uses a | ||
| conservative heuristic. A string is treated as an expression, and returned unchanged, when it | ||
| contains characters that typically appear in expressions (parentheses, brackets, quotes, comma, star) | ||
| or arithmetic/comparison operators that are whitespace-separated (e.g. "a + b"). Everything else is | ||
| treated as a dotted column path: each segment is left alone if it is a valid bare identifier. Otherwise | ||
| it is back-quoted. | ||
|
|
||
| This heuristic intentionally excludes two ambiguous cases: names that contain expression characters | ||
| (e.g. "amount (usd)"), and operator-free SQL expressions such as "col IS NOT NULL" (which is treated | ||
| as a name). Callers with such columns should back-quote the name themselves, pass a Column expression, | ||
| or use the *sql_expression* check. | ||
|
|
||
| Args: | ||
| column: Column reference provided as a string (plain name, nested path, or SQL expression). | ||
|
|
||
|
mwojtyczka marked this conversation as resolved.
|
||
| Returns: | ||
| A string safe to pass to *F.expr*. | ||
| """ | ||
| if SQL_EXPRESSION_STRUCTURE_PATTERN.search(column) or SQL_EXPRESSION_OPERATOR_PATTERN.search(column): | ||
| return column | ||
|
|
||
| return ".".join( | ||
| segment if VALID_UNQUOTED_IDENTIFIER_PATTERN.match(segment) else quote_column_name(segment) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Finding: a column whose name equals a zero-arg SQL function resolves to the function, not the column (silent wrong result) A bare identifier that matches Failure scenario: a DataFrame has a column named
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is also DBSQL's convention. Users must escape zero-arg identifiers in any DBSQL query if they use them as column names. Should we change this or align with DBSQL?
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. sounds good, I think it's better to stay aligned with DBSQL |
||
| for segment in column.split(".") | ||
| ) | ||
|
|
||
|
|
||
| def normalize_col_str(col_str: str) -> str: | ||
| """ | ||
| Normalizes string to be compatible with metastore column names by applying the following transformations: | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Finding: dotted path of separately back-quoted segments is mangled (cosmetic)
unquote_column_nameonly strips the outer back-quotes, so a path whose segments are individually quoted comes back malformed.Failure scenario:
column="Odd Name.field"→unquote_column_name(...)returns"Odd Name.field", from whichnormalize_col_strderives a garbled check name and the_errors/_warningsmessage shows a corrupted column name. Resolution itself is fine (normalize_column_exprleaves the string unchanged andF.exprresolves it), so this is display-only — but user-visible. Low severity.