Edit on GitHub

# sqlglot.transforms

View Source

      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}

def unalias_group( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    14def unalias_group(expression: exp.Expression) -> exp.Expression:
    15    """
    16    Replace references to select aliases in GROUP BY clauses.
    17
    18    Example:
    19        >>> import sqlglot
    20        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
    21        'SELECT a AS b FROM x GROUP BY 1'
    22
    23    Args:
    24        expression: the expression that will be transformed.
    25
    26    Returns:
    27        The transformed expression.
    28    """
    29    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
    30        aliased_selects = {
    31            e.alias: i
    32            for i, e in enumerate(expression.parent.expressions, start=1)
    33            if isinstance(e, exp.Alias)
    34        }
    35
    36        for group_by in expression.expressions:
    37            if (
    38                isinstance(group_by, exp.Column)
    39                and not group_by.table
    40                and group_by.name in aliased_selects
    41            ):
    42                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
    43
    44    return expression

Replace references to select aliases in GROUP BY clauses.

###### Example:

>     >>> import sqlglot
>     >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
>     'SELECT a AS b FROM x GROUP BY 1'

###### Arguments:

- expression: the expression that will be transformed.

###### Returns:

> The transformed expression.

def eliminate_distinct_on( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    47def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
    48    """
    49    Convert SELECT DISTINCT ON statements to a subquery with a window function.
    50
    51    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
    52
    53    Args:
    54        expression: the expression that will be transformed.
    55
    56    Returns:
    57        The transformed expression.
    58    """
    59    if (
    60        isinstance(expression, exp.Select)
    61        and expression.args.get("distinct")
    62        and expression.args["distinct"].args.get("on")
    63        and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
    64    ):
    65        distinct_cols = expression.args["distinct"].args["on"].expressions
    66        expression.args["distinct"].pop()
    67        outer_selects = expression.selects
    68        row_number = find_new_name(expression.named_selects, "_row_number")
    69        window = exp.Window(
    70            this=exp.RowNumber(),
    71            partition_by=distinct_cols,
    72        )
    73        order = expression.args.get("order")
    74        if order:
    75            window.set("order", order.copy())
    76            order.pop()
    77        window = exp.alias_(window, row_number)
    78        expression.select(window, copy=False)
    79        return exp.select(*outer_selects).from_(expression.subquery()).where(f'"{row_number}" = 1')
    80    return expression

Convert SELECT DISTINCT ON statements to a subquery with a window function.

This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.

###### Arguments:

- expression: the expression that will be transformed.

###### Returns:

> The transformed expression.

def remove_precision_parameterized_types( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

     83def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
     84    """
     85    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in other expressions.
     86    This transforms removes the precision from parameterized types in expressions.
     87    """
     88    return expression.transform(
     89        lambda node: exp.DataType(
     90            **{
     91                **node.args,
     92                "expressions": [
     93                    node_expression
     94                    for node_expression in node.expressions
     95                    if isinstance(node_expression, exp.DataType)
     96                ],
     97            }
     98        )
     99        if isinstance(node, exp.DataType)
    100        else node,
    101    )

Some dialects only allow the precision for parameterized types to be defined in the DDL and not in other expressions. This transforms removes the precision from parameterized types in expressions.

def preprocess( transforms: List[Callable[[sqlglot.expressions.Expression], sqlglot.expressions.Expression]], to_sql: Callable[[sqlglot.generator.Generator, sqlglot.expressions.Expression], str]) -> Callable[[sqlglot.generator.Generator, sqlglot.expressions.Expression], str]: View Source

    104def preprocess(
    105    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
    106    to_sql: t.Callable[[Generator, exp.Expression], str],
    107) -> t.Callable[[Generator, exp.Expression], str]:
    108    """
    109    Creates a new transform by chaining a sequence of transformations and converts the resulting
    110    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
    111
    112    Args:
    113        transforms: sequence of transform functions. These will be called in order.
    114        to_sql: final transform that converts the resulting expression to a SQL string.
    115
    116    Returns:
    117        Function that can be used as a generator transform.
    118    """
    119
    120    def _to_sql(self, expression):
    121        expression = transforms[0](expression.copy())
    122        for t in transforms[1:]:
    123            expression = t(expression)
    124        return to_sql(self, expression)
    125
    126    return _to_sql

Creates a new transform by chaining a sequence of transformations and converts the resulting expression to SQL, using an appropriate `Generator.TRANSFORMS` function.

###### Arguments:

- transforms: sequence of transform functions. These will be called in order.
- to_sql: final transform that converts the resulting expression to a SQL string.

###### Returns:

> Function that can be used as a generator transform.

def delegate(attr: str) -> Callable: View Source

    129def delegate(attr: str) -> t.Callable:
    130    """
    131    Create a new method that delegates to `attr`. This is useful for creating `Generator.TRANSFORMS`
    132    functions that delegate to existing generator methods.
    133    """
    134
    135    def _transform(self, *args, **kwargs):
    136        return getattr(self, attr)(*args, **kwargs)
    137
    138    return _transform

Create a new method that delegates to `attr`. This is useful for creating `Generator.TRANSFORMS` functions that delegate to existing generator methods.