sqlglot.transforms

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import expressions as exp
  6from sqlglot.helper import find_new_name
  7
  8if t.TYPE_CHECKING:
  9    from sqlglot.generator import Generator
 10
 11
 12def unalias_group(expression: exp.Expression) -> exp.Expression:
 13    """
 14    Replace references to select aliases in GROUP BY clauses.
 15
 16    Example:
 17        >>> import sqlglot
 18        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
 19        'SELECT a AS b FROM x GROUP BY 1'
 20
 21    Args:
 22        expression: the expression that will be transformed.
 23
 24    Returns:
 25        The transformed expression.
 26    """
 27    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
 28        aliased_selects = {
 29            e.alias: i
 30            for i, e in enumerate(expression.parent.expressions, start=1)
 31            if isinstance(e, exp.Alias)
 32        }
 33
 34        for group_by in expression.expressions:
 35            if (
 36                isinstance(group_by, exp.Column)
 37                and not group_by.table
 38                and group_by.name in aliased_selects
 39            ):
 40                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
 41
 42    return expression
 43
 44
 45def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
 46    """
 47    Convert SELECT DISTINCT ON statements to a subquery with a window function.
 48
 49    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
 50
 51    Args:
 52        expression: the expression that will be transformed.
 53
 54    Returns:
 55        The transformed expression.
 56    """
 57    if (
 58        isinstance(expression, exp.Select)
 59        and expression.args.get("distinct")
 60        and expression.args["distinct"].args.get("on")
 61        and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
 62    ):
 63        distinct_cols = expression.args["distinct"].pop().args["on"].expressions
 64        outer_selects = expression.selects
 65        row_number = find_new_name(expression.named_selects, "_row_number")
 66        window = exp.Window(
 67            this=exp.RowNumber(),
 68            partition_by=distinct_cols,
 69        )
 70        order = expression.args.get("order")
 71        if order:
 72            window.set("order", order.pop().copy())
 73        window = exp.alias_(window, row_number)
 74        expression.select(window, copy=False)
 75        return exp.select(*outer_selects).from_(expression.subquery()).where(f'"{row_number}" = 1')
 76    return expression
 77
 78
 79def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
 80    """
 81    Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
 82
 83    The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
 84    https://docs.snowflake.com/en/sql-reference/constructs/qualify
 85
 86    Some dialects don't support window functions in the WHERE clause, so we need to include them as
 87    projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
 88    if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
 89    otherwise we won't be able to refer to it in the outer query's WHERE clause.
 90    """
 91    if isinstance(expression, exp.Select) and expression.args.get("qualify"):
 92        taken = set(expression.named_selects)
 93        for select in expression.selects:
 94            if not select.alias_or_name:
 95                alias = find_new_name(taken, "_c")
 96                select.replace(exp.alias_(select.copy(), alias))
 97                taken.add(alias)
 98
 99        outer_selects = exp.select(*[select.alias_or_name for select in expression.selects])
100        qualify_filters = expression.args["qualify"].pop().this
101
102        for expr in qualify_filters.find_all((exp.Window, exp.Column)):
103            if isinstance(expr, exp.Window):
104                alias = find_new_name(expression.named_selects, "_w")
105                expression.select(exp.alias_(expr.copy(), alias), copy=False)
106                expr.replace(exp.column(alias))
107            elif expr.name not in expression.named_selects:
108                expression.select(expr.copy(), copy=False)
109
110        return outer_selects.from_(expression.subquery(alias="_t")).where(qualify_filters)
111
112    return expression
113
114
115def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
116    """
117    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in other expressions.
118    This transforms removes the precision from parameterized types in expressions.
119    """
120    return expression.transform(
121        lambda node: exp.DataType(
122            **{
123                **node.args,
124                "expressions": [
125                    node_expression
126                    for node_expression in node.expressions
127                    if isinstance(node_expression, exp.DataType)
128                ],
129            }
130        )
131        if isinstance(node, exp.DataType)
132        else node,
133    )
134
135
136def preprocess(
137    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
138    to_sql: t.Callable[[Generator, exp.Expression], str],
139) -> t.Callable[[Generator, exp.Expression], str]:
140    """
141    Creates a new transform by chaining a sequence of transformations and converts the resulting
142    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
143
144    Args:
145        transforms: sequence of transform functions. These will be called in order.
146        to_sql: final transform that converts the resulting expression to a SQL string.
147
148    Returns:
149        Function that can be used as a generator transform.
150    """
151
152    def _to_sql(self, expression):
153        expression = transforms[0](expression.copy())
154        for t in transforms[1:]:
155            expression = t(expression)
156        return to_sql(self, expression)
157
158    return _to_sql
159
160
161def delegate(attr: str) -> t.Callable:
162    """
163    Create a new method that delegates to `attr`. This is useful for creating `Generator.TRANSFORMS`
164    functions that delegate to existing generator methods.
165    """
166
167    def _transform(self, *args, **kwargs):
168        return getattr(self, attr)(*args, **kwargs)
169
170    return _transform
171
172
173UNALIAS_GROUP = {exp.Group: preprocess([unalias_group], delegate("group_sql"))}
174ELIMINATE_DISTINCT_ON = {exp.Select: preprocess([eliminate_distinct_on], delegate("select_sql"))}
175ELIMINATE_QUALIFY = {exp.Select: preprocess([eliminate_qualify], delegate("select_sql"))}
176REMOVE_PRECISION_PARAMETERIZED_TYPES = {
177    exp.Cast: preprocess([remove_precision_parameterized_types], delegate("cast_sql"))
178}