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}
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.
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.
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.
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.
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.