34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import typing as t
|
|
|
|
from sqlglot import exp
|
|
|
|
|
|
def expand_laterals(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
Expand lateral column alias references.
|
|
|
|
This assumes `qualify_columns` as already run.
|
|
|
|
Example:
|
|
>>> import sqlglot
|
|
>>> sql = "SELECT x.a + 1 AS b, b + 1 AS c FROM x"
|
|
>>> expression = sqlglot.parse_one(sql)
|
|
>>> expand_laterals(expression).sql()
|
|
'SELECT x.a + 1 AS b, x.a + 1 + 1 AS c FROM x'
|
|
|
|
Args:
|
|
expression: expression to optimize
|
|
Returns:
|
|
optimized expression
|
|
"""
|
|
for select in expression.find_all(exp.Select):
|
|
alias_to_expression: t.Dict[str, exp.Expression] = {}
|
|
for projection in select.expressions:
|
|
for column in projection.find_all(exp.Column):
|
|
if not column.table and column.name in alias_to_expression:
|
|
column.replace(alias_to_expression[column.name].copy())
|
|
if isinstance(projection, exp.Alias):
|
|
alias_to_expression[projection.alias] = projection.this
|
|
return expression
|