diff options
Diffstat (limited to 'sqlglot/transforms.py')
-rw-r--r-- | sqlglot/transforms.py | 46 |
1 files changed, 40 insertions, 6 deletions
diff --git a/sqlglot/transforms.py b/sqlglot/transforms.py index aa7d240..2eafb0b 100644 --- a/sqlglot/transforms.py +++ b/sqlglot/transforms.py @@ -2,13 +2,12 @@ from __future__ import annotations import typing as t +from sqlglot import expressions as exp from sqlglot.helper import find_new_name if t.TYPE_CHECKING: from sqlglot.generator import Generator -from sqlglot import expressions as exp - def unalias_group(expression: exp.Expression) -> exp.Expression: """ @@ -61,8 +60,7 @@ def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression: and expression.args["distinct"].args.get("on") and isinstance(expression.args["distinct"].args["on"], exp.Tuple) ): - distinct_cols = expression.args["distinct"].args["on"].expressions - expression.args["distinct"].pop() + distinct_cols = expression.args["distinct"].pop().args["on"].expressions outer_selects = expression.selects row_number = find_new_name(expression.named_selects, "_row_number") window = exp.Window( @@ -71,14 +69,49 @@ def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression: ) order = expression.args.get("order") if order: - window.set("order", order.copy()) - order.pop() + window.set("order", order.pop().copy()) window = exp.alias_(window, row_number) expression.select(window, copy=False) return exp.select(*outer_selects).from_(expression.subquery()).where(f'"{row_number}" = 1') return expression +def eliminate_qualify(expression: exp.Expression) -> exp.Expression: + """ + Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently. + + The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY: + https://docs.snowflake.com/en/sql-reference/constructs/qualify + + Some dialects don't support window functions in the WHERE clause, so we need to include them as + projections in the subquery, in order to refer to them in the outer filter using aliases. Also, + if a column is referenced in the QUALIFY clause but is not selected, we need to include it too, + otherwise we won't be able to refer to it in the outer query's WHERE clause. + """ + if isinstance(expression, exp.Select) and expression.args.get("qualify"): + taken = set(expression.named_selects) + for select in expression.selects: + if not select.alias_or_name: + alias = find_new_name(taken, "_c") + select.replace(exp.alias_(select.copy(), alias)) + taken.add(alias) + + outer_selects = exp.select(*[select.alias_or_name for select in expression.selects]) + qualify_filters = expression.args["qualify"].pop().this + + for expr in qualify_filters.find_all((exp.Window, exp.Column)): + if isinstance(expr, exp.Window): + alias = find_new_name(expression.named_selects, "_w") + expression.select(exp.alias_(expr.copy(), alias), copy=False) + expr.replace(exp.column(alias)) + elif expr.name not in expression.named_selects: + expression.select(expr.copy(), copy=False) + + return outer_selects.from_(expression.subquery(alias="_t")).where(qualify_filters) + + return expression + + def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression: """ Some dialects only allow the precision for parameterized types to be defined in the DDL and not in other expressions. @@ -139,6 +172,7 @@ def delegate(attr: str) -> t.Callable: UNALIAS_GROUP = {exp.Group: preprocess([unalias_group], delegate("group_sql"))} ELIMINATE_DISTINCT_ON = {exp.Select: preprocess([eliminate_distinct_on], delegate("select_sql"))} +ELIMINATE_QUALIFY = {exp.Select: preprocess([eliminate_qualify], delegate("select_sql"))} REMOVE_PRECISION_PARAMETERIZED_TYPES = { exp.Cast: preprocess([remove_precision_parameterized_types], delegate("cast_sql")) } |