sqlglot.transforms

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot.helper import find_new_name
  6
  7if t.TYPE_CHECKING:
  8    from sqlglot.generator import Generator
  9
 10from sqlglot import expressions as exp
 11
 12
 13def unalias_group(expression: exp.Expression) -> exp.Expression:
 14    """
 15    Replace references to select aliases in GROUP BY clauses.
 16
 17    Example:
 18        >>> import sqlglot
 19        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
 20        'SELECT a AS b FROM x GROUP BY 1'
 21
 22    Args:
 23        expression: the expression that will be transformed.
 24
 25    Returns:
 26        The transformed expression.
 27    """
 28    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
 29        aliased_selects = {
 30            e.alias: i
 31            for i, e in enumerate(expression.parent.expressions, start=1)
 32            if isinstance(e, exp.Alias)
 33        }
 34
 35        for group_by in expression.expressions:
 36            if (
 37                isinstance(group_by, exp.Column)
 38                and not group_by.table
 39                and group_by.name in aliased_selects
 40            ):
 41                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
 42
 43    return expression
 44
 45
 46def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
 47    """
 48    Convert SELECT DISTINCT ON statements to a subquery with a window function.
 49
 50    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
 51
 52    Args:
 53        expression: the expression that will be transformed.
 54
 55    Returns:
 56        The transformed expression.
 57    """
 58    if (
 59        isinstance(expression, exp.Select)
 60        and expression.args.get("distinct")
 61        and expression.args["distinct"].args.get("on")
 62        and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
 63    ):
 64        distinct_cols = expression.args["distinct"].args["on"].expressions
 65        expression.args["distinct"].pop()
 66        outer_selects = expression.selects
 67        row_number = find_new_name(expression.named_selects, "_row_number")
 68        window = exp.Window(
 69            this=exp.RowNumber(),
 70            partition_by=distinct_cols,
 71        )
 72        order = expression.args.get("order")
 73        if order:
 74            window.set("order", order.copy())
 75            order.pop()
 76        window = exp.alias_(window, row_number)
 77        expression.select(window, copy=False)
 78        return exp.select(*outer_selects).from_(expression.subquery()).where(f'"{row_number}" = 1')
 79    return expression
 80
 81
 82def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
 83    """
 84    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in other expressions.
 85    This transforms removes the precision from parameterized types in expressions.
 86    """
 87    return expression.transform(
 88        lambda node: exp.DataType(
 89            **{
 90                **node.args,
 91                "expressions": [
 92                    node_expression
 93                    for node_expression in node.expressions
 94                    if isinstance(node_expression, exp.DataType)
 95                ],
 96            }
 97        )
 98        if isinstance(node, exp.DataType)
 99        else node,
100    )
101
102
103def preprocess(
104    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
105    to_sql: t.Callable[[Generator, exp.Expression], str],
106) -> t.Callable[[Generator, exp.Expression], str]:
107    """
108    Creates a new transform by chaining a sequence of transformations and converts the resulting
109    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
110
111    Args:
112        transforms: sequence of transform functions. These will be called in order.
113        to_sql: final transform that converts the resulting expression to a SQL string.
114
115    Returns:
116        Function that can be used as a generator transform.
117    """
118
119    def _to_sql(self, expression):
120        expression = transforms[0](expression.copy())
121        for t in transforms[1:]:
122            expression = t(expression)
123        return to_sql(self, expression)
124
125    return _to_sql
126
127
128def delegate(attr: str) -> t.Callable:
129    """
130    Create a new method that delegates to `attr`. This is useful for creating `Generator.TRANSFORMS`
131    functions that delegate to existing generator methods.
132    """
133
134    def _transform(self, *args, **kwargs):
135        return getattr(self, attr)(*args, **kwargs)
136
137    return _transform
138
139
140UNALIAS_GROUP = {exp.Group: preprocess([unalias_group], delegate("group_sql"))}
141ELIMINATE_DISTINCT_ON = {exp.Select: preprocess([eliminate_distinct_on], delegate("select_sql"))}
142REMOVE_PRECISION_PARAMETERIZED_TYPES = {
143    exp.Cast: preprocess([remove_precision_parameterized_types], delegate("cast_sql"))
144}