sqlglot.optimizer.qualify_columns

  1from __future__ import annotations
  2
  3import itertools
  4import typing as t
  5
  6from sqlglot import alias, exp
  7from sqlglot.dialects.dialect import Dialect, DialectType
  8from sqlglot.errors import OptimizeError
  9from sqlglot.helper import seq_get, SingleValuedMapping
 10from sqlglot.optimizer.annotate_types import TypeAnnotator
 11from sqlglot.optimizer.scope import Scope, build_scope, traverse_scope, walk_in_scope
 12from sqlglot.optimizer.simplify import simplify_parens
 13from sqlglot.schema import Schema, ensure_schema
 14
 15if t.TYPE_CHECKING:
 16    from sqlglot._typing import E
 17
 18
 19def qualify_columns(
 20    expression: exp.Expression,
 21    schema: t.Dict | Schema,
 22    expand_alias_refs: bool = True,
 23    expand_stars: bool = True,
 24    infer_schema: t.Optional[bool] = None,
 25    allow_partial_qualification: bool = False,
 26) -> exp.Expression:
 27    """
 28    Rewrite sqlglot AST to have fully qualified columns.
 29
 30    Example:
 31        >>> import sqlglot
 32        >>> schema = {"tbl": {"col": "INT"}}
 33        >>> expression = sqlglot.parse_one("SELECT col FROM tbl")
 34        >>> qualify_columns(expression, schema).sql()
 35        'SELECT tbl.col AS col FROM tbl'
 36
 37    Args:
 38        expression: Expression to qualify.
 39        schema: Database schema.
 40        expand_alias_refs: Whether to expand references to aliases.
 41        expand_stars: Whether to expand star queries. This is a necessary step
 42            for most of the optimizer's rules to work; do not set to False unless you
 43            know what you're doing!
 44        infer_schema: Whether to infer the schema if missing.
 45        allow_partial_qualification: Whether to allow partial qualification.
 46
 47    Returns:
 48        The qualified expression.
 49
 50    Notes:
 51        - Currently only handles a single PIVOT or UNPIVOT operator
 52    """
 53    schema = ensure_schema(schema)
 54    annotator = TypeAnnotator(schema)
 55    infer_schema = schema.empty if infer_schema is None else infer_schema
 56    dialect = Dialect.get_or_raise(schema.dialect)
 57    pseudocolumns = dialect.PSEUDOCOLUMNS
 58
 59    for scope in traverse_scope(expression):
 60        resolver = Resolver(scope, schema, infer_schema=infer_schema)
 61        _pop_table_column_aliases(scope.ctes)
 62        _pop_table_column_aliases(scope.derived_tables)
 63        using_column_tables = _expand_using(scope, resolver)
 64
 65        if (schema.empty or dialect.FORCE_EARLY_ALIAS_REF_EXPANSION) and expand_alias_refs:
 66            _expand_alias_refs(
 67                scope,
 68                resolver,
 69                dialect,
 70                expand_only_groupby=dialect.EXPAND_ALIAS_REFS_EARLY_ONLY_IN_GROUP_BY,
 71            )
 72
 73        _convert_columns_to_dots(scope, resolver)
 74        _qualify_columns(scope, resolver, allow_partial_qualification=allow_partial_qualification)
 75
 76        if not schema.empty and expand_alias_refs:
 77            _expand_alias_refs(scope, resolver, dialect)
 78
 79        if isinstance(scope.expression, exp.Select):
 80            if expand_stars:
 81                _expand_stars(
 82                    scope,
 83                    resolver,
 84                    using_column_tables,
 85                    pseudocolumns,
 86                    annotator,
 87                )
 88            qualify_outputs(scope)
 89
 90        _expand_group_by(scope, dialect)
 91        _expand_order_by(scope, resolver)
 92
 93        if dialect == "bigquery":
 94            annotator.annotate_scope(scope)
 95
 96    return expression
 97
 98
 99def validate_qualify_columns(expression: E) -> E:
100    """Raise an `OptimizeError` if any columns aren't qualified"""
101    all_unqualified_columns = []
102    for scope in traverse_scope(expression):
103        if isinstance(scope.expression, exp.Select):
104            unqualified_columns = scope.unqualified_columns
105
106            if scope.external_columns and not scope.is_correlated_subquery and not scope.pivots:
107                column = scope.external_columns[0]
108                for_table = f" for table: '{column.table}'" if column.table else ""
109                raise OptimizeError(f"Column '{column}' could not be resolved{for_table}")
110
111            if unqualified_columns and scope.pivots and scope.pivots[0].unpivot:
112                # New columns produced by the UNPIVOT can't be qualified, but there may be columns
113                # under the UNPIVOT's IN clause that can and should be qualified. We recompute
114                # this list here to ensure those in the former category will be excluded.
115                unpivot_columns = set(_unpivot_columns(scope.pivots[0]))
116                unqualified_columns = [c for c in unqualified_columns if c not in unpivot_columns]
117
118            all_unqualified_columns.extend(unqualified_columns)
119
120    if all_unqualified_columns:
121        raise OptimizeError(f"Ambiguous columns: {all_unqualified_columns}")
122
123    return expression
124
125
126def _unpivot_columns(unpivot: exp.Pivot) -> t.Iterator[exp.Column]:
127    name_column = []
128    field = unpivot.args.get("field")
129    if isinstance(field, exp.In) and isinstance(field.this, exp.Column):
130        name_column.append(field.this)
131
132    value_columns = (c for e in unpivot.expressions for c in e.find_all(exp.Column))
133    return itertools.chain(name_column, value_columns)
134
135
136def _pop_table_column_aliases(derived_tables: t.List[exp.CTE | exp.Subquery]) -> None:
137    """
138    Remove table column aliases.
139
140    For example, `col1` and `col2` will be dropped in SELECT ... FROM (SELECT ...) AS foo(col1, col2)
141    """
142    for derived_table in derived_tables:
143        if isinstance(derived_table.parent, exp.With) and derived_table.parent.recursive:
144            continue
145        table_alias = derived_table.args.get("alias")
146        if table_alias:
147            table_alias.args.pop("columns", None)
148
149
150def _expand_using(scope: Scope, resolver: Resolver) -> t.Dict[str, t.Any]:
151    columns = {}
152
153    def _update_source_columns(source_name: str) -> None:
154        for column_name in resolver.get_source_columns(source_name):
155            if column_name not in columns:
156                columns[column_name] = source_name
157
158    joins = list(scope.find_all(exp.Join))
159    names = {join.alias_or_name for join in joins}
160    ordered = [key for key in scope.selected_sources if key not in names]
161
162    # Mapping of automatically joined column names to an ordered set of source names (dict).
163    column_tables: t.Dict[str, t.Dict[str, t.Any]] = {}
164
165    for source_name in ordered:
166        _update_source_columns(source_name)
167
168    for i, join in enumerate(joins):
169        source_table = ordered[-1]
170        if source_table:
171            _update_source_columns(source_table)
172
173        join_table = join.alias_or_name
174        ordered.append(join_table)
175
176        using = join.args.get("using")
177        if not using:
178            continue
179
180        join_columns = resolver.get_source_columns(join_table)
181        conditions = []
182        using_identifier_count = len(using)
183
184        for identifier in using:
185            identifier = identifier.name
186            table = columns.get(identifier)
187
188            if not table or identifier not in join_columns:
189                if (columns and "*" not in columns) and join_columns:
190                    raise OptimizeError(f"Cannot automatically join: {identifier}")
191
192            table = table or source_table
193
194            if i == 0 or using_identifier_count == 1:
195                lhs: exp.Expression = exp.column(identifier, table=table)
196            else:
197                coalesce_columns = [
198                    exp.column(identifier, table=t)
199                    for t in ordered[:-1]
200                    if identifier in resolver.get_source_columns(t)
201                ]
202                if len(coalesce_columns) > 1:
203                    lhs = exp.func("coalesce", *coalesce_columns)
204                else:
205                    lhs = exp.column(identifier, table=table)
206
207            conditions.append(lhs.eq(exp.column(identifier, table=join_table)))
208
209            # Set all values in the dict to None, because we only care about the key ordering
210            tables = column_tables.setdefault(identifier, {})
211            if table not in tables:
212                tables[table] = None
213            if join_table not in tables:
214                tables[join_table] = None
215
216        join.args.pop("using")
217        join.set("on", exp.and_(*conditions, copy=False))
218
219    if column_tables:
220        for column in scope.columns:
221            if not column.table and column.name in column_tables:
222                tables = column_tables[column.name]
223                coalesce_args = [exp.column(column.name, table=table) for table in tables]
224                replacement: exp.Expression = exp.func("coalesce", *coalesce_args)
225
226                if isinstance(column.parent, exp.Select):
227                    # Ensure the USING column keeps its name if it's projected
228                    replacement = alias(replacement, alias=column.name, copy=False)
229                elif isinstance(column.parent, exp.Struct):
230                    # Ensure the USING column keeps its name if it's an anonymous STRUCT field
231                    replacement = exp.PropertyEQ(
232                        this=exp.to_identifier(column.name), expression=replacement
233                    )
234
235                scope.replace(column, replacement)
236
237    return column_tables
238
239
240def _expand_alias_refs(
241    scope: Scope, resolver: Resolver, dialect: Dialect, expand_only_groupby: bool = False
242) -> None:
243    """
244    Expand references to aliases.
245    Example:
246        SELECT y.foo AS bar, bar * 2 AS baz FROM y
247     => SELECT y.foo AS bar, y.foo * 2 AS baz FROM y
248    """
249    expression = scope.expression
250
251    if not isinstance(expression, exp.Select):
252        return
253
254    alias_to_expression: t.Dict[str, t.Tuple[exp.Expression, int]] = {}
255
256    def replace_columns(
257        node: t.Optional[exp.Expression], resolve_table: bool = False, literal_index: bool = False
258    ) -> None:
259        is_group_by = isinstance(node, exp.Group)
260        if not node or (expand_only_groupby and not is_group_by):
261            return
262
263        for column in walk_in_scope(node, prune=lambda node: node.is_star):
264            if not isinstance(column, exp.Column):
265                continue
266
267            # BigQuery's GROUP BY allows alias expansion only for standalone names, e.g:
268            #   SELECT FUNC(col) AS col FROM t GROUP BY col --> Can be expanded
269            #   SELECT FUNC(col) AS col FROM t GROUP BY FUNC(col)  --> Shouldn't be expanded, will result to FUNC(FUNC(col))
270            # This not required for the HAVING clause as it can evaluate expressions using both the alias & the table columns
271            if expand_only_groupby and is_group_by and column.parent is not node:
272                continue
273
274            table = resolver.get_table(column.name) if resolve_table and not column.table else None
275            alias_expr, i = alias_to_expression.get(column.name, (None, 1))
276            double_agg = (
277                (
278                    alias_expr.find(exp.AggFunc)
279                    and (
280                        column.find_ancestor(exp.AggFunc)
281                        and not isinstance(column.find_ancestor(exp.Window, exp.Select), exp.Window)
282                    )
283                )
284                if alias_expr
285                else False
286            )
287
288            if table and (not alias_expr or double_agg):
289                column.set("table", table)
290            elif not column.table and alias_expr and not double_agg:
291                if isinstance(alias_expr, exp.Literal) and (literal_index or resolve_table):
292                    if literal_index:
293                        column.replace(exp.Literal.number(i))
294                else:
295                    column = column.replace(exp.paren(alias_expr))
296                    simplified = simplify_parens(column)
297                    if simplified is not column:
298                        column.replace(simplified)
299
300    for i, projection in enumerate(expression.selects):
301        replace_columns(projection)
302        if isinstance(projection, exp.Alias):
303            alias_to_expression[projection.alias] = (projection.this, i + 1)
304
305    parent_scope = scope
306    while parent_scope.is_union:
307        parent_scope = parent_scope.parent
308
309    # We shouldn't expand aliases if they match the recursive CTE's columns
310    if parent_scope.is_cte:
311        cte = parent_scope.expression.parent
312        if cte.find_ancestor(exp.With).recursive:
313            for recursive_cte_column in cte.args["alias"].columns or cte.this.selects:
314                alias_to_expression.pop(recursive_cte_column.output_name, None)
315
316    replace_columns(expression.args.get("where"))
317    replace_columns(expression.args.get("group"), literal_index=True)
318    replace_columns(expression.args.get("having"), resolve_table=True)
319    replace_columns(expression.args.get("qualify"), resolve_table=True)
320
321    # Snowflake allows alias expansion in the JOIN ... ON clause (and almost everywhere else)
322    # https://docs.snowflake.com/en/sql-reference/sql/select#usage-notes
323    if dialect == "snowflake":
324        for join in expression.args.get("joins") or []:
325            replace_columns(join)
326
327    scope.clear_cache()
328
329
330def _expand_group_by(scope: Scope, dialect: DialectType) -> None:
331    expression = scope.expression
332    group = expression.args.get("group")
333    if not group:
334        return
335
336    group.set("expressions", _expand_positional_references(scope, group.expressions, dialect))
337    expression.set("group", group)
338
339
340def _expand_order_by(scope: Scope, resolver: Resolver) -> None:
341    order = scope.expression.args.get("order")
342    if not order:
343        return
344
345    ordereds = order.expressions
346    for ordered, new_expression in zip(
347        ordereds,
348        _expand_positional_references(
349            scope, (o.this for o in ordereds), resolver.schema.dialect, alias=True
350        ),
351    ):
352        for agg in ordered.find_all(exp.AggFunc):
353            for col in agg.find_all(exp.Column):
354                if not col.table:
355                    col.set("table", resolver.get_table(col.name))
356
357        ordered.set("this", new_expression)
358
359    if scope.expression.args.get("group"):
360        selects = {s.this: exp.column(s.alias_or_name) for s in scope.expression.selects}
361
362        for ordered in ordereds:
363            ordered = ordered.this
364
365            ordered.replace(
366                exp.to_identifier(_select_by_pos(scope, ordered).alias)
367                if ordered.is_int
368                else selects.get(ordered, ordered)
369            )
370
371
372def _expand_positional_references(
373    scope: Scope, expressions: t.Iterable[exp.Expression], dialect: DialectType, alias: bool = False
374) -> t.List[exp.Expression]:
375    new_nodes: t.List[exp.Expression] = []
376    ambiguous_projections = None
377
378    for node in expressions:
379        if node.is_int:
380            select = _select_by_pos(scope, t.cast(exp.Literal, node))
381
382            if alias:
383                new_nodes.append(exp.column(select.args["alias"].copy()))
384            else:
385                select = select.this
386
387                if dialect == "bigquery":
388                    if ambiguous_projections is None:
389                        # When a projection name is also a source name and it is referenced in the
390                        # GROUP BY clause, BQ can't understand what the identifier corresponds to
391                        ambiguous_projections = {
392                            s.alias_or_name
393                            for s in scope.expression.selects
394                            if s.alias_or_name in scope.selected_sources
395                        }
396
397                    ambiguous = any(
398                        column.parts[0].name in ambiguous_projections
399                        for column in select.find_all(exp.Column)
400                    )
401                else:
402                    ambiguous = False
403
404                if (
405                    isinstance(select, exp.CONSTANTS)
406                    or select.find(exp.Explode, exp.Unnest)
407                    or ambiguous
408                ):
409                    new_nodes.append(node)
410                else:
411                    new_nodes.append(select.copy())
412        else:
413            new_nodes.append(node)
414
415    return new_nodes
416
417
418def _select_by_pos(scope: Scope, node: exp.Literal) -> exp.Alias:
419    try:
420        return scope.expression.selects[int(node.this) - 1].assert_is(exp.Alias)
421    except IndexError:
422        raise OptimizeError(f"Unknown output column: {node.name}")
423
424
425def _convert_columns_to_dots(scope: Scope, resolver: Resolver) -> None:
426    """
427    Converts `Column` instances that represent struct field lookup into chained `Dots`.
428
429    Struct field lookups look like columns (e.g. "struct"."field"), but they need to be
430    qualified separately and represented as Dot(Dot(...(<table>.<column>, field1), field2, ...)).
431    """
432    converted = False
433    for column in itertools.chain(scope.columns, scope.stars):
434        if isinstance(column, exp.Dot):
435            continue
436
437        column_table: t.Optional[str | exp.Identifier] = column.table
438        if (
439            column_table
440            and column_table not in scope.sources
441            and (
442                not scope.parent
443                or column_table not in scope.parent.sources
444                or not scope.is_correlated_subquery
445            )
446        ):
447            root, *parts = column.parts
448
449            if root.name in scope.sources:
450                # The struct is already qualified, but we still need to change the AST
451                column_table = root
452                root, *parts = parts
453            else:
454                column_table = resolver.get_table(root.name)
455
456            if column_table:
457                converted = True
458                column.replace(exp.Dot.build([exp.column(root, table=column_table), *parts]))
459
460    if converted:
461        # We want to re-aggregate the converted columns, otherwise they'd be skipped in
462        # a `for column in scope.columns` iteration, even though they shouldn't be
463        scope.clear_cache()
464
465
466def _qualify_columns(scope: Scope, resolver: Resolver, allow_partial_qualification: bool) -> None:
467    """Disambiguate columns, ensuring each column specifies a source"""
468    for column in scope.columns:
469        column_table = column.table
470        column_name = column.name
471
472        if column_table and column_table in scope.sources:
473            source_columns = resolver.get_source_columns(column_table)
474            if (
475                not allow_partial_qualification
476                and source_columns
477                and column_name not in source_columns
478                and "*" not in source_columns
479            ):
480                raise OptimizeError(f"Unknown column: {column_name}")
481
482        if not column_table:
483            if scope.pivots and not column.find_ancestor(exp.Pivot):
484                # If the column is under the Pivot expression, we need to qualify it
485                # using the name of the pivoted source instead of the pivot's alias
486                column.set("table", exp.to_identifier(scope.pivots[0].alias))
487                continue
488
489            # column_table can be a '' because bigquery unnest has no table alias
490            column_table = resolver.get_table(column_name)
491            if column_table:
492                column.set("table", column_table)
493
494    for pivot in scope.pivots:
495        for column in pivot.find_all(exp.Column):
496            if not column.table and column.name in resolver.all_columns:
497                column_table = resolver.get_table(column.name)
498                if column_table:
499                    column.set("table", column_table)
500
501
502def _expand_struct_stars(
503    expression: exp.Dot,
504) -> t.List[exp.Alias]:
505    """[BigQuery] Expand/Flatten foo.bar.* where bar is a struct column"""
506
507    dot_column = t.cast(exp.Column, expression.find(exp.Column))
508    if not dot_column.is_type(exp.DataType.Type.STRUCT):
509        return []
510
511    # All nested struct values are ColumnDefs, so normalize the first exp.Column in one
512    dot_column = dot_column.copy()
513    starting_struct = exp.ColumnDef(this=dot_column.this, kind=dot_column.type)
514
515    # First part is the table name and last part is the star so they can be dropped
516    dot_parts = expression.parts[1:-1]
517
518    # If we're expanding a nested struct eg. t.c.f1.f2.* find the last struct (f2 in this case)
519    for part in dot_parts[1:]:
520        for field in t.cast(exp.DataType, starting_struct.kind).expressions:
521            # Unable to expand star unless all fields are named
522            if not isinstance(field.this, exp.Identifier):
523                return []
524
525            if field.name == part.name and field.kind.is_type(exp.DataType.Type.STRUCT):
526                starting_struct = field
527                break
528        else:
529            # There is no matching field in the struct
530            return []
531
532    taken_names = set()
533    new_selections = []
534
535    for field in t.cast(exp.DataType, starting_struct.kind).expressions:
536        name = field.name
537
538        # Ambiguous or anonymous fields can't be expanded
539        if name in taken_names or not isinstance(field.this, exp.Identifier):
540            return []
541
542        taken_names.add(name)
543
544        this = field.this.copy()
545        root, *parts = [part.copy() for part in itertools.chain(dot_parts, [this])]
546        new_column = exp.column(
547            t.cast(exp.Identifier, root),
548            table=dot_column.args.get("table"),
549            fields=t.cast(t.List[exp.Identifier], parts),
550        )
551        new_selections.append(alias(new_column, this, copy=False))
552
553    return new_selections
554
555
556def _expand_stars(
557    scope: Scope,
558    resolver: Resolver,
559    using_column_tables: t.Dict[str, t.Any],
560    pseudocolumns: t.Set[str],
561    annotator: TypeAnnotator,
562) -> None:
563    """Expand stars to lists of column selections"""
564
565    new_selections: t.List[exp.Expression] = []
566    except_columns: t.Dict[int, t.Set[str]] = {}
567    replace_columns: t.Dict[int, t.Dict[str, exp.Alias]] = {}
568    rename_columns: t.Dict[int, t.Dict[str, str]] = {}
569
570    coalesced_columns = set()
571    dialect = resolver.schema.dialect
572
573    pivot_output_columns = None
574    pivot_exclude_columns = None
575
576    pivot = t.cast(t.Optional[exp.Pivot], seq_get(scope.pivots, 0))
577    if isinstance(pivot, exp.Pivot) and not pivot.alias_column_names:
578        if pivot.unpivot:
579            pivot_output_columns = [c.output_name for c in _unpivot_columns(pivot)]
580
581            field = pivot.args.get("field")
582            if isinstance(field, exp.In):
583                pivot_exclude_columns = {
584                    c.output_name for e in field.expressions for c in e.find_all(exp.Column)
585                }
586        else:
587            pivot_exclude_columns = set(c.output_name for c in pivot.find_all(exp.Column))
588
589            pivot_output_columns = [c.output_name for c in pivot.args.get("columns", [])]
590            if not pivot_output_columns:
591                pivot_output_columns = [c.alias_or_name for c in pivot.expressions]
592
593    is_bigquery = dialect == "bigquery"
594    if is_bigquery and any(isinstance(col, exp.Dot) for col in scope.stars):
595        # Found struct expansion, annotate scope ahead of time
596        annotator.annotate_scope(scope)
597
598    for expression in scope.expression.selects:
599        tables = []
600        if isinstance(expression, exp.Star):
601            tables.extend(scope.selected_sources)
602            _add_except_columns(expression, tables, except_columns)
603            _add_replace_columns(expression, tables, replace_columns)
604            _add_rename_columns(expression, tables, rename_columns)
605        elif expression.is_star:
606            if not isinstance(expression, exp.Dot):
607                tables.append(expression.table)
608                _add_except_columns(expression.this, tables, except_columns)
609                _add_replace_columns(expression.this, tables, replace_columns)
610                _add_rename_columns(expression.this, tables, rename_columns)
611            elif is_bigquery:
612                struct_fields = _expand_struct_stars(expression)
613                if struct_fields:
614                    new_selections.extend(struct_fields)
615                    continue
616
617        if not tables:
618            new_selections.append(expression)
619            continue
620
621        for table in tables:
622            if table not in scope.sources:
623                raise OptimizeError(f"Unknown table: {table}")
624
625            columns = resolver.get_source_columns(table, only_visible=True)
626            columns = columns or scope.outer_columns
627
628            if pseudocolumns:
629                columns = [name for name in columns if name.upper() not in pseudocolumns]
630
631            if not columns or "*" in columns:
632                return
633
634            table_id = id(table)
635            columns_to_exclude = except_columns.get(table_id) or set()
636            renamed_columns = rename_columns.get(table_id, {})
637            replaced_columns = replace_columns.get(table_id, {})
638
639            if pivot:
640                if pivot_output_columns and pivot_exclude_columns:
641                    pivot_columns = [c for c in columns if c not in pivot_exclude_columns]
642                    pivot_columns.extend(pivot_output_columns)
643                else:
644                    pivot_columns = pivot.alias_column_names
645
646                if pivot_columns:
647                    new_selections.extend(
648                        alias(exp.column(name, table=pivot.alias), name, copy=False)
649                        for name in pivot_columns
650                        if name not in columns_to_exclude
651                    )
652                    continue
653
654            for name in columns:
655                if name in columns_to_exclude or name in coalesced_columns:
656                    continue
657                if name in using_column_tables and table in using_column_tables[name]:
658                    coalesced_columns.add(name)
659                    tables = using_column_tables[name]
660                    coalesce_args = [exp.column(name, table=table) for table in tables]
661
662                    new_selections.append(
663                        alias(exp.func("coalesce", *coalesce_args), alias=name, copy=False)
664                    )
665                else:
666                    alias_ = renamed_columns.get(name, name)
667                    selection_expr = replaced_columns.get(name) or exp.column(name, table=table)
668                    new_selections.append(
669                        alias(selection_expr, alias_, copy=False)
670                        if alias_ != name
671                        else selection_expr
672                    )
673
674    # Ensures we don't overwrite the initial selections with an empty list
675    if new_selections and isinstance(scope.expression, exp.Select):
676        scope.expression.set("expressions", new_selections)
677
678
679def _add_except_columns(
680    expression: exp.Expression, tables, except_columns: t.Dict[int, t.Set[str]]
681) -> None:
682    except_ = expression.args.get("except")
683
684    if not except_:
685        return
686
687    columns = {e.name for e in except_}
688
689    for table in tables:
690        except_columns[id(table)] = columns
691
692
693def _add_rename_columns(
694    expression: exp.Expression, tables, rename_columns: t.Dict[int, t.Dict[str, str]]
695) -> None:
696    rename = expression.args.get("rename")
697
698    if not rename:
699        return
700
701    columns = {e.this.name: e.alias for e in rename}
702
703    for table in tables:
704        rename_columns[id(table)] = columns
705
706
707def _add_replace_columns(
708    expression: exp.Expression, tables, replace_columns: t.Dict[int, t.Dict[str, exp.Alias]]
709) -> None:
710    replace = expression.args.get("replace")
711
712    if not replace:
713        return
714
715    columns = {e.alias: e for e in replace}
716
717    for table in tables:
718        replace_columns[id(table)] = columns
719
720
721def qualify_outputs(scope_or_expression: Scope | exp.Expression) -> None:
722    """Ensure all output columns are aliased"""
723    if isinstance(scope_or_expression, exp.Expression):
724        scope = build_scope(scope_or_expression)
725        if not isinstance(scope, Scope):
726            return
727    else:
728        scope = scope_or_expression
729
730    new_selections = []
731    for i, (selection, aliased_column) in enumerate(
732        itertools.zip_longest(scope.expression.selects, scope.outer_columns)
733    ):
734        if selection is None:
735            break
736
737        if isinstance(selection, exp.Subquery):
738            if not selection.output_name:
739                selection.set("alias", exp.TableAlias(this=exp.to_identifier(f"_col_{i}")))
740        elif not isinstance(selection, exp.Alias) and not selection.is_star:
741            selection = alias(
742                selection,
743                alias=selection.output_name or f"_col_{i}",
744                copy=False,
745            )
746        if aliased_column:
747            selection.set("alias", exp.to_identifier(aliased_column))
748
749        new_selections.append(selection)
750
751    if isinstance(scope.expression, exp.Select):
752        scope.expression.set("expressions", new_selections)
753
754
755def quote_identifiers(expression: E, dialect: DialectType = None, identify: bool = True) -> E:
756    """Makes sure all identifiers that need to be quoted are quoted."""
757    return expression.transform(
758        Dialect.get_or_raise(dialect).quote_identifier, identify=identify, copy=False
759    )  # type: ignore
760
761
762def pushdown_cte_alias_columns(expression: exp.Expression) -> exp.Expression:
763    """
764    Pushes down the CTE alias columns into the projection,
765
766    This step is useful in Snowflake where the CTE alias columns can be referenced in the HAVING.
767
768    Example:
769        >>> import sqlglot
770        >>> expression = sqlglot.parse_one("WITH y (c) AS (SELECT SUM(a) FROM ( SELECT 1 a ) AS x HAVING c > 0) SELECT c FROM y")
771        >>> pushdown_cte_alias_columns(expression).sql()
772        'WITH y(c) AS (SELECT SUM(a) AS c FROM (SELECT 1 AS a) AS x HAVING c > 0) SELECT c FROM y'
773
774    Args:
775        expression: Expression to pushdown.
776
777    Returns:
778        The expression with the CTE aliases pushed down into the projection.
779    """
780    for cte in expression.find_all(exp.CTE):
781        if cte.alias_column_names:
782            new_expressions = []
783            for _alias, projection in zip(cte.alias_column_names, cte.this.expressions):
784                if isinstance(projection, exp.Alias):
785                    projection.set("alias", _alias)
786                else:
787                    projection = alias(projection, alias=_alias)
788                new_expressions.append(projection)
789            cte.this.set("expressions", new_expressions)
790
791    return expression
792
793
794class Resolver:
795    """
796    Helper for resolving columns.
797
798    This is a class so we can lazily load some things and easily share them across functions.
799    """
800
801    def __init__(self, scope: Scope, schema: Schema, infer_schema: bool = True):
802        self.scope = scope
803        self.schema = schema
804        self._source_columns: t.Optional[t.Dict[str, t.Sequence[str]]] = None
805        self._unambiguous_columns: t.Optional[t.Mapping[str, str]] = None
806        self._all_columns: t.Optional[t.Set[str]] = None
807        self._infer_schema = infer_schema
808        self._get_source_columns_cache: t.Dict[t.Tuple[str, bool], t.Sequence[str]] = {}
809
810    def get_table(self, column_name: str) -> t.Optional[exp.Identifier]:
811        """
812        Get the table for a column name.
813
814        Args:
815            column_name: The column name to find the table for.
816        Returns:
817            The table name if it can be found/inferred.
818        """
819        if self._unambiguous_columns is None:
820            self._unambiguous_columns = self._get_unambiguous_columns(
821                self._get_all_source_columns()
822            )
823
824        table_name = self._unambiguous_columns.get(column_name)
825
826        if not table_name and self._infer_schema:
827            sources_without_schema = tuple(
828                source
829                for source, columns in self._get_all_source_columns().items()
830                if not columns or "*" in columns
831            )
832            if len(sources_without_schema) == 1:
833                table_name = sources_without_schema[0]
834
835        if table_name not in self.scope.selected_sources:
836            return exp.to_identifier(table_name)
837
838        node, _ = self.scope.selected_sources.get(table_name)
839
840        if isinstance(node, exp.Query):
841            while node and node.alias != table_name:
842                node = node.parent
843
844        node_alias = node.args.get("alias")
845        if node_alias:
846            return exp.to_identifier(node_alias.this)
847
848        return exp.to_identifier(table_name)
849
850    @property
851    def all_columns(self) -> t.Set[str]:
852        """All available columns of all sources in this scope"""
853        if self._all_columns is None:
854            self._all_columns = {
855                column for columns in self._get_all_source_columns().values() for column in columns
856            }
857        return self._all_columns
858
859    def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequence[str]:
860        """Resolve the source columns for a given source `name`."""
861        cache_key = (name, only_visible)
862        if cache_key not in self._get_source_columns_cache:
863            if name not in self.scope.sources:
864                raise OptimizeError(f"Unknown table: {name}")
865
866            source = self.scope.sources[name]
867
868            if isinstance(source, exp.Table):
869                columns = self.schema.column_names(source, only_visible)
870            elif isinstance(source, Scope) and isinstance(
871                source.expression, (exp.Values, exp.Unnest)
872            ):
873                columns = source.expression.named_selects
874
875                # in bigquery, unnest structs are automatically scoped as tables, so you can
876                # directly select a struct field in a query.
877                # this handles the case where the unnest is statically defined.
878                if self.schema.dialect == "bigquery":
879                    if source.expression.is_type(exp.DataType.Type.STRUCT):
880                        for k in source.expression.type.expressions:  # type: ignore
881                            columns.append(k.name)
882            else:
883                columns = source.expression.named_selects
884
885            node, _ = self.scope.selected_sources.get(name) or (None, None)
886            if isinstance(node, Scope):
887                column_aliases = node.expression.alias_column_names
888            elif isinstance(node, exp.Expression):
889                column_aliases = node.alias_column_names
890            else:
891                column_aliases = []
892
893            if column_aliases:
894                # If the source's columns are aliased, their aliases shadow the corresponding column names.
895                # This can be expensive if there are lots of columns, so only do this if column_aliases exist.
896                columns = [
897                    alias or name
898                    for (name, alias) in itertools.zip_longest(columns, column_aliases)
899                ]
900
901            pseudocolumns = self._get_source_pseudocolumns(name)
902            if pseudocolumns:
903                columns = list(columns)
904                columns.extend(c for c in pseudocolumns if c not in columns)
905
906            self._get_source_columns_cache[cache_key] = columns
907
908        return self._get_source_columns_cache[cache_key]
909
910    def _get_source_pseudocolumns(self, name: str) -> t.Sequence[str]:
911        if self.schema.dialect == "snowflake" and self.scope.expression.args.get("connect"):
912            # When there is a CONNECT BY clause, there is only one table being scanned
913            # See: https://docs.snowflake.com/en/sql-reference/constructs/connect-by
914            return ["LEVEL"]
915        return []
916
917    def _get_all_source_columns(self) -> t.Dict[str, t.Sequence[str]]:
918        if self._source_columns is None:
919            self._source_columns = {
920                source_name: self.get_source_columns(source_name)
921                for source_name, source in itertools.chain(
922                    self.scope.selected_sources.items(), self.scope.lateral_sources.items()
923                )
924            }
925        return self._source_columns
926
927    def _get_unambiguous_columns(
928        self, source_columns: t.Dict[str, t.Sequence[str]]
929    ) -> t.Mapping[str, str]:
930        """
931        Find all the unambiguous columns in sources.
932
933        Args:
934            source_columns: Mapping of names to source columns.
935
936        Returns:
937            Mapping of column name to source name.
938        """
939        if not source_columns:
940            return {}
941
942        source_columns_pairs = list(source_columns.items())
943
944        first_table, first_columns = source_columns_pairs[0]
945
946        if len(source_columns_pairs) == 1:
947            # Performance optimization - avoid copying first_columns if there is only one table.
948            return SingleValuedMapping(first_columns, first_table)
949
950        unambiguous_columns = {col: first_table for col in first_columns}
951        all_columns = set(unambiguous_columns)
952
953        for table, columns in source_columns_pairs[1:]:
954            unique = set(columns)
955            ambiguous = all_columns.intersection(unique)
956            all_columns.update(columns)
957
958            for column in ambiguous:
959                unambiguous_columns.pop(column, None)
960            for column in unique.difference(ambiguous):
961                unambiguous_columns[column] = table
962
963        return unambiguous_columns