python/sqlglot
1
0
Fork 0
Code Issues Activity

Merging upstream version 25.6.1.

Browse source 
Signed-off-by: Daniel Baumann <daniel@debian.org>
This commit is contained in:
Daniel Baumann 2025-02-13 21:43:00 +01:00
parent 78f79d1d22
commit 4a7feb3eaa
Signed by: daniel
GPG key ID: FBB4F0E80A80222F
69 changed files with 46817 additions and 45778 deletions
Download patch file Download diff file Expand all files Collapse all files

13
sqlglot/optimizer/annotate_types.py
View file

@ -184,7 +184,6 @@ class TypeAnnotator(metaclass=_TypeAnnotator):
exp.Ceil,
exp.DatetimeDiff,
exp.DateDiff,
exp.Extract,
exp.TimestampDiff,
exp.TimeDiff,
exp.DateToDi,
@ -268,6 +267,7 @@ class TypeAnnotator(metaclass=_TypeAnnotator):
exp.Div: lambda self, e: self._annotate_div(e),
exp.Dot: lambda self, e: self._annotate_dot(e),
exp.Explode: lambda self, e: self._annotate_explode(e),
exp.Extract: lambda self, e: self._annotate_extract(e),
exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
exp.GenerateDateArray: lambda self, e: self._annotate_with_type(
e, exp.DataType.build("ARRAY<DATE>")
@ -680,3 +680,14 @@ class TypeAnnotator(metaclass=_TypeAnnotator):
self._set_type(expression, map_type)
return expression
def _annotate_extract(self, expression: exp.Extract) -> exp.Extract:
self._annotate_args(expression)
part = expression.name
if part == "TIME":
self._set_type(expression, exp.DataType.Type.TIME)
elif part == "DATE":
self._set_type(expression, exp.DataType.Type.DATE)
else:
self._set_type(expression, exp.DataType.Type.INT)
return expression

34
sqlglot/optimizer/qualify_columns.py
View file

@ -513,7 +513,9 @@ def _expand_stars(
new_selections = []
except_columns: t.Dict[int, t.Set[str]] = {}
replace_columns: t.Dict[int, t.Dict[str, str]] = {}
replace_columns: t.Dict[int, t.Dict[str, exp.Alias]] = {}
rename_columns: t.Dict[int, t.Dict[str, str]] = {}
coalesced_columns = set()
dialect = resolver.schema.dialect
@ -548,11 +550,13 @@ def _expand_stars(
tables.extend(scope.selected_sources)
_add_except_columns(expression, tables, except_columns)
_add_replace_columns(expression, tables, replace_columns)
_add_rename_columns(expression, tables, rename_columns)
elif expression.is_star:
if not isinstance(expression, exp.Dot):
tables.append(expression.table)
_add_except_columns(expression.this, tables, except_columns)
_add_replace_columns(expression.this, tables, replace_columns)
_add_rename_columns(expression.this, tables, rename_columns)
elif is_bigquery:
struct_fields = _expand_struct_stars(expression)
if struct_fields:
@ -578,6 +582,8 @@ def _expand_stars(
table_id = id(table)
columns_to_exclude = except_columns.get(table_id) or set()
renamed_columns = rename_columns.get(table_id, {})
replaced_columns = replace_columns.get(table_id, {})
if pivot:
if pivot_output_columns and pivot_exclude_columns:
@ -606,10 +612,12 @@ def _expand_stars(
alias(exp.func("coalesce", *coalesce_args), alias=name, copy=False)
)
else:
alias_ = replace_columns.get(table_id, {}).get(name, name)
column = exp.column(name, table=table)
alias_ = renamed_columns.get(name, name)
selection_expr = replaced_columns.get(name) or exp.column(name, table=table)
new_selections.append(
alias(column, alias_, copy=False) if alias_ != name else column
alias(selection_expr, alias_, copy=False)
if alias_ != name
else selection_expr
)
# Ensures we don't overwrite the initial selections with an empty list
@ -631,15 +639,29 @@ def _add_except_columns(
except_columns[id(table)] = columns
def _add_rename_columns(
expression: exp.Expression, tables, rename_columns: t.Dict[int, t.Dict[str, str]]
) -> None:
rename = expression.args.get("rename")
if not rename:
return
columns = {e.this.name: e.alias for e in rename}
for table in tables:
rename_columns[id(table)] = columns
def _add_replace_columns(
expression: exp.Expression, tables, replace_columns: t.Dict[int, t.Dict[str, str]]
expression: exp.Expression, tables, replace_columns: t.Dict[int, t.Dict[str, exp.Alias]]
) -> None:
replace = expression.args.get("replace")
if not replace:
return
columns = {e.this.name: e.alias for e in replace}
columns = {e.alias: e for e in replace}
for table in tables:
replace_columns[id(table)] = columns

1
sqlglot/optimizer/scope.py
View file

@ -284,6 +284,7 @@ class Scope:
or column.name not in named_selects
)
)
or (isinstance(ancestor, exp.Star) and not column.arg_key == "except")
):
self._columns.append(column)