916 lines
35 KiB
Python
916 lines
35 KiB
Python
from __future__ import annotations
|
|
|
|
import typing as t
|
|
|
|
from sqlglot import expressions as exp
|
|
from sqlglot.errors import UnsupportedError
|
|
from sqlglot.helper import find_new_name, name_sequence
|
|
|
|
|
|
if t.TYPE_CHECKING:
|
|
from sqlglot._typing import E
|
|
from sqlglot.generator import Generator
|
|
|
|
|
|
def preprocess(
|
|
transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
|
|
) -> t.Callable[[Generator, exp.Expression], str]:
|
|
"""
|
|
Creates a new transform by chaining a sequence of transformations and converts the resulting
|
|
expression to SQL, using either the "_sql" method corresponding to the resulting expression,
|
|
or the appropriate `Generator.TRANSFORMS` function (when applicable -- see below).
|
|
|
|
Args:
|
|
transforms: sequence of transform functions. These will be called in order.
|
|
|
|
Returns:
|
|
Function that can be used as a generator transform.
|
|
"""
|
|
|
|
def _to_sql(self, expression: exp.Expression) -> str:
|
|
expression_type = type(expression)
|
|
|
|
try:
|
|
expression = transforms[0](expression)
|
|
for transform in transforms[1:]:
|
|
expression = transform(expression)
|
|
except UnsupportedError as unsupported_error:
|
|
self.unsupported(str(unsupported_error))
|
|
|
|
_sql_handler = getattr(self, expression.key + "_sql", None)
|
|
if _sql_handler:
|
|
return _sql_handler(expression)
|
|
|
|
transforms_handler = self.TRANSFORMS.get(type(expression))
|
|
if transforms_handler:
|
|
if expression_type is type(expression):
|
|
if isinstance(expression, exp.Func):
|
|
return self.function_fallback_sql(expression)
|
|
|
|
# Ensures we don't enter an infinite loop. This can happen when the original expression
|
|
# has the same type as the final expression and there's no _sql method available for it,
|
|
# because then it'd re-enter _to_sql.
|
|
raise ValueError(
|
|
f"Expression type {expression.__class__.__name__} requires a _sql method in order to be transformed."
|
|
)
|
|
|
|
return transforms_handler(self, expression)
|
|
|
|
raise ValueError(f"Unsupported expression type {expression.__class__.__name__}.")
|
|
|
|
return _to_sql
|
|
|
|
|
|
def unnest_generate_date_array_using_recursive_cte(expression: exp.Expression) -> exp.Expression:
|
|
if isinstance(expression, exp.Select):
|
|
count = 0
|
|
recursive_ctes = []
|
|
|
|
for unnest in expression.find_all(exp.Unnest):
|
|
if (
|
|
not isinstance(unnest.parent, (exp.From, exp.Join))
|
|
or len(unnest.expressions) != 1
|
|
or not isinstance(unnest.expressions[0], exp.GenerateDateArray)
|
|
):
|
|
continue
|
|
|
|
generate_date_array = unnest.expressions[0]
|
|
start = generate_date_array.args.get("start")
|
|
end = generate_date_array.args.get("end")
|
|
step = generate_date_array.args.get("step")
|
|
|
|
if not start or not end or not isinstance(step, exp.Interval):
|
|
continue
|
|
|
|
alias = unnest.args.get("alias")
|
|
column_name = alias.columns[0] if isinstance(alias, exp.TableAlias) else "date_value"
|
|
|
|
start = exp.cast(start, "date")
|
|
date_add = exp.func(
|
|
"date_add", column_name, exp.Literal.number(step.name), step.args.get("unit")
|
|
)
|
|
cast_date_add = exp.cast(date_add, "date")
|
|
|
|
cte_name = "_generated_dates" + (f"_{count}" if count else "")
|
|
|
|
base_query = exp.select(start.as_(column_name))
|
|
recursive_query = (
|
|
exp.select(cast_date_add)
|
|
.from_(cte_name)
|
|
.where(cast_date_add <= exp.cast(end, "date"))
|
|
)
|
|
cte_query = base_query.union(recursive_query, distinct=False)
|
|
|
|
generate_dates_query = exp.select(column_name).from_(cte_name)
|
|
unnest.replace(generate_dates_query.subquery(cte_name))
|
|
|
|
recursive_ctes.append(
|
|
exp.alias_(exp.CTE(this=cte_query), cte_name, table=[column_name])
|
|
)
|
|
count += 1
|
|
|
|
if recursive_ctes:
|
|
with_expression = expression.args.get("with") or exp.With()
|
|
with_expression.set("recursive", True)
|
|
with_expression.set("expressions", [*recursive_ctes, *with_expression.expressions])
|
|
expression.set("with", with_expression)
|
|
|
|
return expression
|
|
|
|
|
|
def unnest_generate_series(expression: exp.Expression) -> exp.Expression:
|
|
"""Unnests GENERATE_SERIES or SEQUENCE table references."""
|
|
this = expression.this
|
|
if isinstance(expression, exp.Table) and isinstance(this, exp.GenerateSeries):
|
|
unnest = exp.Unnest(expressions=[this])
|
|
if expression.alias:
|
|
return exp.alias_(unnest, alias="_u", table=[expression.alias], copy=False)
|
|
|
|
return unnest
|
|
|
|
return expression
|
|
|
|
|
|
def unalias_group(expression: exp.Expression) -> exp.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'
|
|
|
|
Args:
|
|
expression: the expression that will be transformed.
|
|
|
|
Returns:
|
|
The transformed expression.
|
|
"""
|
|
if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
|
|
aliased_selects = {
|
|
e.alias: i
|
|
for i, e in enumerate(expression.parent.expressions, start=1)
|
|
if isinstance(e, exp.Alias)
|
|
}
|
|
|
|
for group_by in expression.expressions:
|
|
if (
|
|
isinstance(group_by, exp.Column)
|
|
and not group_by.table
|
|
and group_by.name in aliased_selects
|
|
):
|
|
group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
|
|
|
|
return expression
|
|
|
|
|
|
def eliminate_distinct_on(expression: exp.Expression) -> exp.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.
|
|
|
|
Args:
|
|
expression: the expression that will be transformed.
|
|
|
|
Returns:
|
|
The transformed expression.
|
|
"""
|
|
if (
|
|
isinstance(expression, exp.Select)
|
|
and expression.args.get("distinct")
|
|
and expression.args["distinct"].args.get("on")
|
|
and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
|
|
):
|
|
distinct_cols = expression.args["distinct"].pop().args["on"].expressions
|
|
outer_selects = expression.selects
|
|
row_number = find_new_name(expression.named_selects, "_row_number")
|
|
window = exp.Window(this=exp.RowNumber(), partition_by=distinct_cols)
|
|
order = expression.args.get("order")
|
|
|
|
if order:
|
|
window.set("order", order.pop())
|
|
else:
|
|
window.set("order", exp.Order(expressions=[c.copy() for c in distinct_cols]))
|
|
|
|
window = exp.alias_(window, row_number)
|
|
expression.select(window, copy=False)
|
|
|
|
return (
|
|
exp.select(*outer_selects, copy=False)
|
|
.from_(expression.subquery("_t", copy=False), copy=False)
|
|
.where(exp.column(row_number).eq(1), copy=False)
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
|
|
|
|
The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
|
|
https://docs.snowflake.com/en/sql-reference/constructs/qualify
|
|
|
|
Some dialects don't support window functions in the WHERE clause, so we need to include them as
|
|
projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
|
|
if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
|
|
otherwise we won't be able to refer to it in the outer query's WHERE clause. Finally, if a
|
|
newly aliased projection is referenced in the QUALIFY clause, it will be replaced by the
|
|
corresponding expression to avoid creating invalid column references.
|
|
"""
|
|
if isinstance(expression, exp.Select) and expression.args.get("qualify"):
|
|
taken = set(expression.named_selects)
|
|
for select in expression.selects:
|
|
if not select.alias_or_name:
|
|
alias = find_new_name(taken, "_c")
|
|
select.replace(exp.alias_(select, alias))
|
|
taken.add(alias)
|
|
|
|
def _select_alias_or_name(select: exp.Expression) -> str | exp.Column:
|
|
alias_or_name = select.alias_or_name
|
|
identifier = select.args.get("alias") or select.this
|
|
if isinstance(identifier, exp.Identifier):
|
|
return exp.column(alias_or_name, quoted=identifier.args.get("quoted"))
|
|
return alias_or_name
|
|
|
|
outer_selects = exp.select(*list(map(_select_alias_or_name, expression.selects)))
|
|
qualify_filters = expression.args["qualify"].pop().this
|
|
expression_by_alias = {
|
|
select.alias: select.this
|
|
for select in expression.selects
|
|
if isinstance(select, exp.Alias)
|
|
}
|
|
|
|
select_candidates = exp.Window if expression.is_star else (exp.Window, exp.Column)
|
|
for select_candidate in qualify_filters.find_all(select_candidates):
|
|
if isinstance(select_candidate, exp.Window):
|
|
if expression_by_alias:
|
|
for column in select_candidate.find_all(exp.Column):
|
|
expr = expression_by_alias.get(column.name)
|
|
if expr:
|
|
column.replace(expr)
|
|
|
|
alias = find_new_name(expression.named_selects, "_w")
|
|
expression.select(exp.alias_(select_candidate, alias), copy=False)
|
|
column = exp.column(alias)
|
|
|
|
if isinstance(select_candidate.parent, exp.Qualify):
|
|
qualify_filters = column
|
|
else:
|
|
select_candidate.replace(column)
|
|
elif select_candidate.name not in expression.named_selects:
|
|
expression.select(select_candidate.copy(), copy=False)
|
|
|
|
return outer_selects.from_(expression.subquery(alias="_t", copy=False), copy=False).where(
|
|
qualify_filters, copy=False
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
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.
|
|
"""
|
|
for node in expression.find_all(exp.DataType):
|
|
node.set(
|
|
"expressions", [e for e in node.expressions if not isinstance(e, exp.DataTypeParam)]
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def unqualify_unnest(expression: exp.Expression) -> exp.Expression:
|
|
"""Remove references to unnest table aliases, added by the optimizer's qualify_columns step."""
|
|
from sqlglot.optimizer.scope import find_all_in_scope
|
|
|
|
if isinstance(expression, exp.Select):
|
|
unnest_aliases = {
|
|
unnest.alias
|
|
for unnest in find_all_in_scope(expression, exp.Unnest)
|
|
if isinstance(unnest.parent, (exp.From, exp.Join))
|
|
}
|
|
if unnest_aliases:
|
|
for column in expression.find_all(exp.Column):
|
|
if column.table in unnest_aliases:
|
|
column.set("table", None)
|
|
elif column.db in unnest_aliases:
|
|
column.set("db", None)
|
|
|
|
return expression
|
|
|
|
|
|
def unnest_to_explode(
|
|
expression: exp.Expression,
|
|
unnest_using_arrays_zip: bool = True,
|
|
) -> exp.Expression:
|
|
"""Convert cross join unnest into lateral view explode."""
|
|
|
|
def _unnest_zip_exprs(
|
|
u: exp.Unnest, unnest_exprs: t.List[exp.Expression], has_multi_expr: bool
|
|
) -> t.List[exp.Expression]:
|
|
if has_multi_expr:
|
|
if not unnest_using_arrays_zip:
|
|
raise UnsupportedError("Cannot transpile UNNEST with multiple input arrays")
|
|
|
|
# Use INLINE(ARRAYS_ZIP(...)) for multiple expressions
|
|
zip_exprs: t.List[exp.Expression] = [
|
|
exp.Anonymous(this="ARRAYS_ZIP", expressions=unnest_exprs)
|
|
]
|
|
u.set("expressions", zip_exprs)
|
|
return zip_exprs
|
|
return unnest_exprs
|
|
|
|
def _udtf_type(u: exp.Unnest, has_multi_expr: bool) -> t.Type[exp.Func]:
|
|
if u.args.get("offset"):
|
|
return exp.Posexplode
|
|
return exp.Inline if has_multi_expr else exp.Explode
|
|
|
|
if isinstance(expression, exp.Select):
|
|
from_ = expression.args.get("from")
|
|
|
|
if from_ and isinstance(from_.this, exp.Unnest):
|
|
unnest = from_.this
|
|
alias = unnest.args.get("alias")
|
|
exprs = unnest.expressions
|
|
has_multi_expr = len(exprs) > 1
|
|
this, *expressions = _unnest_zip_exprs(unnest, exprs, has_multi_expr)
|
|
|
|
unnest.replace(
|
|
exp.Table(
|
|
this=_udtf_type(unnest, has_multi_expr)(
|
|
this=this,
|
|
expressions=expressions,
|
|
),
|
|
alias=exp.TableAlias(this=alias.this, columns=alias.columns) if alias else None,
|
|
)
|
|
)
|
|
|
|
for join in expression.args.get("joins") or []:
|
|
join_expr = join.this
|
|
|
|
is_lateral = isinstance(join_expr, exp.Lateral)
|
|
|
|
unnest = join_expr.this if is_lateral else join_expr
|
|
|
|
if isinstance(unnest, exp.Unnest):
|
|
if is_lateral:
|
|
alias = join_expr.args.get("alias")
|
|
else:
|
|
alias = unnest.args.get("alias")
|
|
exprs = unnest.expressions
|
|
# The number of unnest.expressions will be changed by _unnest_zip_exprs, we need to record it here
|
|
has_multi_expr = len(exprs) > 1
|
|
exprs = _unnest_zip_exprs(unnest, exprs, has_multi_expr)
|
|
|
|
expression.args["joins"].remove(join)
|
|
|
|
alias_cols = alias.columns if alias else []
|
|
for e, column in zip(exprs, alias_cols):
|
|
expression.append(
|
|
"laterals",
|
|
exp.Lateral(
|
|
this=_udtf_type(unnest, has_multi_expr)(this=e),
|
|
view=True,
|
|
alias=exp.TableAlias(
|
|
this=alias.this, # type: ignore
|
|
columns=alias_cols if unnest_using_arrays_zip else [column], # type: ignore
|
|
),
|
|
),
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def explode_to_unnest(index_offset: int = 0) -> t.Callable[[exp.Expression], exp.Expression]:
|
|
"""Convert explode/posexplode into unnest."""
|
|
|
|
def _explode_to_unnest(expression: exp.Expression) -> exp.Expression:
|
|
if isinstance(expression, exp.Select):
|
|
from sqlglot.optimizer.scope import Scope
|
|
|
|
taken_select_names = set(expression.named_selects)
|
|
taken_source_names = {name for name, _ in Scope(expression).references}
|
|
|
|
def new_name(names: t.Set[str], name: str) -> str:
|
|
name = find_new_name(names, name)
|
|
names.add(name)
|
|
return name
|
|
|
|
arrays: t.List[exp.Condition] = []
|
|
series_alias = new_name(taken_select_names, "pos")
|
|
series = exp.alias_(
|
|
exp.Unnest(
|
|
expressions=[exp.GenerateSeries(start=exp.Literal.number(index_offset))]
|
|
),
|
|
new_name(taken_source_names, "_u"),
|
|
table=[series_alias],
|
|
)
|
|
|
|
# we use list here because expression.selects is mutated inside the loop
|
|
for select in list(expression.selects):
|
|
explode = select.find(exp.Explode)
|
|
|
|
if explode:
|
|
pos_alias = ""
|
|
explode_alias = ""
|
|
|
|
if isinstance(select, exp.Alias):
|
|
explode_alias = select.args["alias"]
|
|
alias = select
|
|
elif isinstance(select, exp.Aliases):
|
|
pos_alias = select.aliases[0]
|
|
explode_alias = select.aliases[1]
|
|
alias = select.replace(exp.alias_(select.this, "", copy=False))
|
|
else:
|
|
alias = select.replace(exp.alias_(select, ""))
|
|
explode = alias.find(exp.Explode)
|
|
assert explode
|
|
|
|
is_posexplode = isinstance(explode, exp.Posexplode)
|
|
explode_arg = explode.this
|
|
|
|
if isinstance(explode, exp.ExplodeOuter):
|
|
bracket = explode_arg[0]
|
|
bracket.set("safe", True)
|
|
bracket.set("offset", True)
|
|
explode_arg = exp.func(
|
|
"IF",
|
|
exp.func(
|
|
"ARRAY_SIZE", exp.func("COALESCE", explode_arg, exp.Array())
|
|
).eq(0),
|
|
exp.array(bracket, copy=False),
|
|
explode_arg,
|
|
)
|
|
|
|
# This ensures that we won't use [POS]EXPLODE's argument as a new selection
|
|
if isinstance(explode_arg, exp.Column):
|
|
taken_select_names.add(explode_arg.output_name)
|
|
|
|
unnest_source_alias = new_name(taken_source_names, "_u")
|
|
|
|
if not explode_alias:
|
|
explode_alias = new_name(taken_select_names, "col")
|
|
|
|
if is_posexplode:
|
|
pos_alias = new_name(taken_select_names, "pos")
|
|
|
|
if not pos_alias:
|
|
pos_alias = new_name(taken_select_names, "pos")
|
|
|
|
alias.set("alias", exp.to_identifier(explode_alias))
|
|
|
|
series_table_alias = series.args["alias"].this
|
|
column = exp.If(
|
|
this=exp.column(series_alias, table=series_table_alias).eq(
|
|
exp.column(pos_alias, table=unnest_source_alias)
|
|
),
|
|
true=exp.column(explode_alias, table=unnest_source_alias),
|
|
)
|
|
|
|
explode.replace(column)
|
|
|
|
if is_posexplode:
|
|
expressions = expression.expressions
|
|
expressions.insert(
|
|
expressions.index(alias) + 1,
|
|
exp.If(
|
|
this=exp.column(series_alias, table=series_table_alias).eq(
|
|
exp.column(pos_alias, table=unnest_source_alias)
|
|
),
|
|
true=exp.column(pos_alias, table=unnest_source_alias),
|
|
).as_(pos_alias),
|
|
)
|
|
expression.set("expressions", expressions)
|
|
|
|
if not arrays:
|
|
if expression.args.get("from"):
|
|
expression.join(series, copy=False, join_type="CROSS")
|
|
else:
|
|
expression.from_(series, copy=False)
|
|
|
|
size: exp.Condition = exp.ArraySize(this=explode_arg.copy())
|
|
arrays.append(size)
|
|
|
|
# trino doesn't support left join unnest with on conditions
|
|
# if it did, this would be much simpler
|
|
expression.join(
|
|
exp.alias_(
|
|
exp.Unnest(
|
|
expressions=[explode_arg.copy()],
|
|
offset=exp.to_identifier(pos_alias),
|
|
),
|
|
unnest_source_alias,
|
|
table=[explode_alias],
|
|
),
|
|
join_type="CROSS",
|
|
copy=False,
|
|
)
|
|
|
|
if index_offset != 1:
|
|
size = size - 1
|
|
|
|
expression.where(
|
|
exp.column(series_alias, table=series_table_alias)
|
|
.eq(exp.column(pos_alias, table=unnest_source_alias))
|
|
.or_(
|
|
(exp.column(series_alias, table=series_table_alias) > size).and_(
|
|
exp.column(pos_alias, table=unnest_source_alias).eq(size)
|
|
)
|
|
),
|
|
copy=False,
|
|
)
|
|
|
|
if arrays:
|
|
end: exp.Condition = exp.Greatest(this=arrays[0], expressions=arrays[1:])
|
|
|
|
if index_offset != 1:
|
|
end = end - (1 - index_offset)
|
|
series.expressions[0].set("end", end)
|
|
|
|
return expression
|
|
|
|
return _explode_to_unnest
|
|
|
|
|
|
def add_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
|
|
"""Transforms percentiles by adding a WITHIN GROUP clause to them."""
|
|
if (
|
|
isinstance(expression, exp.PERCENTILES)
|
|
and not isinstance(expression.parent, exp.WithinGroup)
|
|
and expression.expression
|
|
):
|
|
column = expression.this.pop()
|
|
expression.set("this", expression.expression.pop())
|
|
order = exp.Order(expressions=[exp.Ordered(this=column)])
|
|
expression = exp.WithinGroup(this=expression, expression=order)
|
|
|
|
return expression
|
|
|
|
|
|
def remove_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
|
|
"""Transforms percentiles by getting rid of their corresponding WITHIN GROUP clause."""
|
|
if (
|
|
isinstance(expression, exp.WithinGroup)
|
|
and isinstance(expression.this, exp.PERCENTILES)
|
|
and isinstance(expression.expression, exp.Order)
|
|
):
|
|
quantile = expression.this.this
|
|
input_value = t.cast(exp.Ordered, expression.find(exp.Ordered)).this
|
|
return expression.replace(exp.ApproxQuantile(this=input_value, quantile=quantile))
|
|
|
|
return expression
|
|
|
|
|
|
def add_recursive_cte_column_names(expression: exp.Expression) -> exp.Expression:
|
|
"""Uses projection output names in recursive CTE definitions to define the CTEs' columns."""
|
|
if isinstance(expression, exp.With) and expression.recursive:
|
|
next_name = name_sequence("_c_")
|
|
|
|
for cte in expression.expressions:
|
|
if not cte.args["alias"].columns:
|
|
query = cte.this
|
|
if isinstance(query, exp.SetOperation):
|
|
query = query.this
|
|
|
|
cte.args["alias"].set(
|
|
"columns",
|
|
[exp.to_identifier(s.alias_or_name or next_name()) for s in query.selects],
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def epoch_cast_to_ts(expression: exp.Expression) -> exp.Expression:
|
|
"""Replace 'epoch' in casts by the equivalent date literal."""
|
|
if (
|
|
isinstance(expression, (exp.Cast, exp.TryCast))
|
|
and expression.name.lower() == "epoch"
|
|
and expression.to.this in exp.DataType.TEMPORAL_TYPES
|
|
):
|
|
expression.this.replace(exp.Literal.string("1970-01-01 00:00:00"))
|
|
|
|
return expression
|
|
|
|
|
|
def eliminate_semi_and_anti_joins(expression: exp.Expression) -> exp.Expression:
|
|
"""Convert SEMI and ANTI joins into equivalent forms that use EXIST instead."""
|
|
if isinstance(expression, exp.Select):
|
|
for join in expression.args.get("joins") or []:
|
|
on = join.args.get("on")
|
|
if on and join.kind in ("SEMI", "ANTI"):
|
|
subquery = exp.select("1").from_(join.this).where(on)
|
|
exists = exp.Exists(this=subquery)
|
|
if join.kind == "ANTI":
|
|
exists = exists.not_(copy=False)
|
|
|
|
join.pop()
|
|
expression.where(exists, copy=False)
|
|
|
|
return expression
|
|
|
|
|
|
def eliminate_full_outer_join(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
Converts a query with a FULL OUTER join to a union of identical queries that
|
|
use LEFT/RIGHT OUTER joins instead. This transformation currently only works
|
|
for queries that have a single FULL OUTER join.
|
|
"""
|
|
if isinstance(expression, exp.Select):
|
|
full_outer_joins = [
|
|
(index, join)
|
|
for index, join in enumerate(expression.args.get("joins") or [])
|
|
if join.side == "FULL"
|
|
]
|
|
|
|
if len(full_outer_joins) == 1:
|
|
expression_copy = expression.copy()
|
|
expression.set("limit", None)
|
|
index, full_outer_join = full_outer_joins[0]
|
|
|
|
tables = (expression.args["from"].alias_or_name, full_outer_join.alias_or_name)
|
|
join_conditions = full_outer_join.args.get("on") or exp.and_(
|
|
*[
|
|
exp.column(col, tables[0]).eq(exp.column(col, tables[1]))
|
|
for col in full_outer_join.args.get("using")
|
|
]
|
|
)
|
|
|
|
full_outer_join.set("side", "left")
|
|
anti_join_clause = exp.select("1").from_(expression.args["from"]).where(join_conditions)
|
|
expression_copy.args["joins"][index].set("side", "right")
|
|
expression_copy = expression_copy.where(exp.Exists(this=anti_join_clause).not_())
|
|
expression_copy.args.pop("with", None) # remove CTEs from RIGHT side
|
|
expression.args.pop("order", None) # remove order by from LEFT side
|
|
|
|
return exp.union(expression, expression_copy, copy=False, distinct=False)
|
|
|
|
return expression
|
|
|
|
|
|
def move_ctes_to_top_level(expression: E) -> E:
|
|
"""
|
|
Some dialects (e.g. Hive, T-SQL, Spark prior to version 3) only allow CTEs to be
|
|
defined at the top-level, so for example queries like:
|
|
|
|
SELECT * FROM (WITH t(c) AS (SELECT 1) SELECT * FROM t) AS subq
|
|
|
|
are invalid in those dialects. This transformation can be used to ensure all CTEs are
|
|
moved to the top level so that the final SQL code is valid from a syntax standpoint.
|
|
|
|
TODO: handle name clashes whilst moving CTEs (it can get quite tricky & costly).
|
|
"""
|
|
top_level_with = expression.args.get("with")
|
|
for inner_with in expression.find_all(exp.With):
|
|
if inner_with.parent is expression:
|
|
continue
|
|
|
|
if not top_level_with:
|
|
top_level_with = inner_with.pop()
|
|
expression.set("with", top_level_with)
|
|
else:
|
|
if inner_with.recursive:
|
|
top_level_with.set("recursive", True)
|
|
|
|
parent_cte = inner_with.find_ancestor(exp.CTE)
|
|
inner_with.pop()
|
|
|
|
if parent_cte:
|
|
i = top_level_with.expressions.index(parent_cte)
|
|
top_level_with.expressions[i:i] = inner_with.expressions
|
|
top_level_with.set("expressions", top_level_with.expressions)
|
|
else:
|
|
top_level_with.set(
|
|
"expressions", top_level_with.expressions + inner_with.expressions
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def ensure_bools(expression: exp.Expression) -> exp.Expression:
|
|
"""Converts numeric values used in conditions into explicit boolean expressions."""
|
|
from sqlglot.optimizer.canonicalize import ensure_bools
|
|
|
|
def _ensure_bool(node: exp.Expression) -> None:
|
|
if (
|
|
node.is_number
|
|
or (
|
|
not isinstance(node, exp.SubqueryPredicate)
|
|
and node.is_type(exp.DataType.Type.UNKNOWN, *exp.DataType.NUMERIC_TYPES)
|
|
)
|
|
or (isinstance(node, exp.Column) and not node.type)
|
|
):
|
|
node.replace(node.neq(0))
|
|
|
|
for node in expression.walk():
|
|
ensure_bools(node, _ensure_bool)
|
|
|
|
return expression
|
|
|
|
|
|
def unqualify_columns(expression: exp.Expression) -> exp.Expression:
|
|
for column in expression.find_all(exp.Column):
|
|
# We only wanna pop off the table, db, catalog args
|
|
for part in column.parts[:-1]:
|
|
part.pop()
|
|
|
|
return expression
|
|
|
|
|
|
def remove_unique_constraints(expression: exp.Expression) -> exp.Expression:
|
|
assert isinstance(expression, exp.Create)
|
|
for constraint in expression.find_all(exp.UniqueColumnConstraint):
|
|
if constraint.parent:
|
|
constraint.parent.pop()
|
|
|
|
return expression
|
|
|
|
|
|
def ctas_with_tmp_tables_to_create_tmp_view(
|
|
expression: exp.Expression,
|
|
tmp_storage_provider: t.Callable[[exp.Expression], exp.Expression] = lambda e: e,
|
|
) -> exp.Expression:
|
|
assert isinstance(expression, exp.Create)
|
|
properties = expression.args.get("properties")
|
|
temporary = any(
|
|
isinstance(prop, exp.TemporaryProperty)
|
|
for prop in (properties.expressions if properties else [])
|
|
)
|
|
|
|
# CTAS with temp tables map to CREATE TEMPORARY VIEW
|
|
if expression.kind == "TABLE" and temporary:
|
|
if expression.expression:
|
|
return exp.Create(
|
|
kind="TEMPORARY VIEW",
|
|
this=expression.this,
|
|
expression=expression.expression,
|
|
)
|
|
return tmp_storage_provider(expression)
|
|
|
|
return expression
|
|
|
|
|
|
def move_schema_columns_to_partitioned_by(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
In Hive, the PARTITIONED BY property acts as an extension of a table's schema. When the
|
|
PARTITIONED BY value is an array of column names, they are transformed into a schema.
|
|
The corresponding columns are removed from the create statement.
|
|
"""
|
|
assert isinstance(expression, exp.Create)
|
|
has_schema = isinstance(expression.this, exp.Schema)
|
|
is_partitionable = expression.kind in {"TABLE", "VIEW"}
|
|
|
|
if has_schema and is_partitionable:
|
|
prop = expression.find(exp.PartitionedByProperty)
|
|
if prop and prop.this and not isinstance(prop.this, exp.Schema):
|
|
schema = expression.this
|
|
columns = {v.name.upper() for v in prop.this.expressions}
|
|
partitions = [col for col in schema.expressions if col.name.upper() in columns]
|
|
schema.set("expressions", [e for e in schema.expressions if e not in partitions])
|
|
prop.replace(exp.PartitionedByProperty(this=exp.Schema(expressions=partitions)))
|
|
expression.set("this", schema)
|
|
|
|
return expression
|
|
|
|
|
|
def move_partitioned_by_to_schema_columns(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
Spark 3 supports both "HIVEFORMAT" and "DATASOURCE" formats for CREATE TABLE.
|
|
|
|
Currently, SQLGlot uses the DATASOURCE format for Spark 3.
|
|
"""
|
|
assert isinstance(expression, exp.Create)
|
|
prop = expression.find(exp.PartitionedByProperty)
|
|
if (
|
|
prop
|
|
and prop.this
|
|
and isinstance(prop.this, exp.Schema)
|
|
and all(isinstance(e, exp.ColumnDef) and e.kind for e in prop.this.expressions)
|
|
):
|
|
prop_this = exp.Tuple(
|
|
expressions=[exp.to_identifier(e.this) for e in prop.this.expressions]
|
|
)
|
|
schema = expression.this
|
|
for e in prop.this.expressions:
|
|
schema.append("expressions", e)
|
|
prop.set("this", prop_this)
|
|
|
|
return expression
|
|
|
|
|
|
def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
|
|
"""Converts struct arguments to aliases, e.g. STRUCT(1 AS y)."""
|
|
if isinstance(expression, exp.Struct):
|
|
expression.set(
|
|
"expressions",
|
|
[
|
|
exp.alias_(e.expression, e.this) if isinstance(e, exp.PropertyEQ) else e
|
|
for e in expression.expressions
|
|
],
|
|
)
|
|
|
|
return expression
|
|
|
|
|
|
def eliminate_join_marks(expression: exp.Expression) -> exp.Expression:
|
|
"""
|
|
Remove join marks from an AST. This rule assumes that all marked columns are qualified.
|
|
If this does not hold for a query, consider running `sqlglot.optimizer.qualify` first.
|
|
|
|
For example,
|
|
SELECT * FROM a, b WHERE a.id = b.id(+) -- ... is converted to
|
|
SELECT * FROM a LEFT JOIN b ON a.id = b.id -- this
|
|
|
|
Args:
|
|
expression: The AST to remove join marks from.
|
|
|
|
Returns:
|
|
The AST with join marks removed.
|
|
"""
|
|
from sqlglot.optimizer.scope import traverse_scope
|
|
|
|
for scope in traverse_scope(expression):
|
|
query = scope.expression
|
|
|
|
where = query.args.get("where")
|
|
joins = query.args.get("joins")
|
|
|
|
if not where or not joins:
|
|
continue
|
|
|
|
query_from = query.args["from"]
|
|
|
|
# These keep track of the joins to be replaced
|
|
new_joins: t.Dict[str, exp.Join] = {}
|
|
old_joins = {join.alias_or_name: join for join in joins}
|
|
|
|
for column in scope.columns:
|
|
if not column.args.get("join_mark"):
|
|
continue
|
|
|
|
predicate = column.find_ancestor(exp.Predicate, exp.Select)
|
|
assert isinstance(
|
|
predicate, exp.Binary
|
|
), "Columns can only be marked with (+) when involved in a binary operation"
|
|
|
|
predicate_parent = predicate.parent
|
|
join_predicate = predicate.pop()
|
|
|
|
left_columns = [
|
|
c for c in join_predicate.left.find_all(exp.Column) if c.args.get("join_mark")
|
|
]
|
|
right_columns = [
|
|
c for c in join_predicate.right.find_all(exp.Column) if c.args.get("join_mark")
|
|
]
|
|
|
|
assert not (
|
|
left_columns and right_columns
|
|
), "The (+) marker cannot appear in both sides of a binary predicate"
|
|
|
|
marked_column_tables = set()
|
|
for col in left_columns or right_columns:
|
|
table = col.table
|
|
assert table, f"Column {col} needs to be qualified with a table"
|
|
|
|
col.set("join_mark", False)
|
|
marked_column_tables.add(table)
|
|
|
|
assert (
|
|
len(marked_column_tables) == 1
|
|
), "Columns of only a single table can be marked with (+) in a given binary predicate"
|
|
|
|
join_this = old_joins.get(col.table, query_from).this
|
|
new_join = exp.Join(this=join_this, on=join_predicate, kind="LEFT")
|
|
|
|
# Upsert new_join into new_joins dictionary
|
|
new_join_alias_or_name = new_join.alias_or_name
|
|
existing_join = new_joins.get(new_join_alias_or_name)
|
|
if existing_join:
|
|
existing_join.set("on", exp.and_(existing_join.args.get("on"), new_join.args["on"]))
|
|
else:
|
|
new_joins[new_join_alias_or_name] = new_join
|
|
|
|
# If the parent of the target predicate is a binary node, then it now has only one child
|
|
if isinstance(predicate_parent, exp.Binary):
|
|
if predicate_parent.left is None:
|
|
predicate_parent.replace(predicate_parent.right)
|
|
else:
|
|
predicate_parent.replace(predicate_parent.left)
|
|
|
|
if query_from.alias_or_name in new_joins:
|
|
only_old_joins = old_joins.keys() - new_joins.keys()
|
|
assert (
|
|
len(only_old_joins) >= 1
|
|
), "Cannot determine which table to use in the new FROM clause"
|
|
|
|
new_from_name = list(only_old_joins)[0]
|
|
query.set("from", exp.From(this=old_joins[new_from_name].this))
|
|
|
|
query.set("joins", list(new_joins.values()))
|
|
|
|
if not where.this:
|
|
where.pop()
|
|
|
|
return expression
|