sqlglot 27.29.0__py3-none-any.whl → 28.4.0__py3-none-any.whl

sqlglot/optimizer/qualify_tables.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-import itertools
 import typing as t
 
-from sqlglot import alias, exp
+from sqlglot import exp
 from sqlglot.dialects.dialect import Dialect, DialectType
-from sqlglot.helper import name_sequence
+from sqlglot.helper import name_sequence, seq_get, ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.scope import Scope, traverse_scope
 
@@ -17,8 +16,9 @@ def qualify_tables(
     expression: E,
     db: t.Optional[str | exp.Identifier] = None,
     catalog: t.Optional[str | exp.Identifier] = None,
-    on_qualify: t.Optional[t.Callable[[exp.Expression], None]] = None,
+    on_qualify: t.Optional[t.Callable[[exp.Table], None]] = None,
     dialect: DialectType = None,
+    canonicalize_table_aliases: bool = False,
 ) -> E:
     """
     Rewrite sqlglot AST to have fully qualified tables. Join constructs such as
@@ -40,16 +40,14 @@ def qualify_tables(
         catalog: Catalog name
         on_qualify: Callback after a table has been qualified.
         dialect: The dialect to parse catalog and schema into.
+        canonicalize_table_aliases: Whether to use canonical aliases (_0, _1, ...) for all sources
+            instead of preserving table names. Defaults to False.
 
     Returns:
         The qualified expression.
     """
     dialect = Dialect.get_or_raise(dialect)
-
-    alias_sequence = name_sequence("_q_")
-
-    def next_alias_name() -> str:
-        return normalize_identifiers(alias_sequence(), dialect=dialect).name
+    next_alias_name = name_sequence("_")
 
     if db := db or None:
         db = exp.parse_identifier(db, dialect=dialect)
@@ -68,95 +66,132 @@ def qualify_tables(
                 table.set("catalog", catalog.copy())
 
     if (db or catalog) and not isinstance(expression, exp.Query):
-        with_ = expression.args.get("with") or exp.With()
+        with_ = expression.args.get("with_") or exp.With()
         cte_names = {cte.alias_or_name for cte in with_.expressions}
 
         for node in expression.walk(prune=lambda n: isinstance(n, exp.Query)):
             if isinstance(node, exp.Table) and node.name not in cte_names:
                 _qualify(node)
 
+    def _set_alias(
+        expression: exp.Expression,
+        canonical_aliases: t.Dict[str, str],
+        target_alias: t.Optional[str] = None,
+        scope: t.Optional[Scope] = None,
+        normalize: bool = False,
+        columns: t.Optional[t.List[t.Union[str, exp.Identifier]]] = None,
+    ) -> None:
+        alias = expression.args.get("alias") or exp.TableAlias()
+
+        if canonicalize_table_aliases:
+            new_alias_name = next_alias_name()
+            canonical_aliases[alias.name or target_alias or ""] = new_alias_name
+        elif not alias.name:
+            new_alias_name = target_alias or next_alias_name()
+            if normalize and target_alias:
+                new_alias_name = normalize_identifiers(new_alias_name, dialect=dialect).name
+        else:
+            return
+
+        alias.set("this", exp.to_identifier(new_alias_name))
+
+        if columns:
+            alias.set("columns", [exp.to_identifier(c) for c in columns])
+
+        expression.set("alias", alias)
+
+        if scope:
+            scope.rename_source(None, new_alias_name)
+
     for scope in traverse_scope(expression):
-        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
-            if isinstance(derived_table, exp.Subquery):
-                unnested = derived_table.unnest()
-                if isinstance(unnested, exp.Table):
-                    joins = unnested.args.get("joins")
-                    unnested.set("joins", None)
-                    derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
-                    derived_table.this.set("joins", joins)
-
-            if not derived_table.args.get("alias"):
-                alias_ = next_alias_name()
-                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
-                scope.rename_source(None, alias_)
-
-            pivots = derived_table.args.get("pivots")
-            if pivots and not pivots[0].alias:
-                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
+        local_columns = scope.local_columns
+        canonical_aliases: t.Dict[str, str] = {}
+
+        for query in scope.subqueries:
+            subquery = query.parent
+            if isinstance(subquery, exp.Subquery):
+                subquery.unwrap().replace(subquery)
+
+        for derived_table in scope.derived_tables:
+            unnested = derived_table.unnest()
+            if isinstance(unnested, exp.Table):
+                joins = unnested.args.get("joins")
+                unnested.set("joins", None)
+                derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
+                derived_table.this.set("joins", joins)
+
+            _set_alias(derived_table, canonical_aliases, scope=scope)
+            if pivot := seq_get(derived_table.args.get("pivots") or [], 0):
+                _set_alias(pivot, canonical_aliases)
 
         table_aliases = {}
 
         for name, source in scope.sources.items():
             if isinstance(source, exp.Table):
-                pivots = source.args.get("pivots")
-                if not source.alias:
-                    # Don't add the pivot's alias to the pivoted table, use the table's name instead
-                    if pivots and pivots[0].alias == name:
-                        name = source.name
-
-                    # Mutates the source by attaching an alias to it
-                    normalized_alias = normalize_identifiers(
-                        name or source.name or alias_sequence(), dialect=dialect
-                    )
-                    alias(source, normalized_alias, copy=False, table=True)
-
-                table_aliases[".".join(p.name for p in source.parts)] = exp.to_identifier(
-                    source.alias
+                # When the name is empty, it means that we have a non-table source, e.g. a pivoted cte
+                is_real_table_source = bool(name)
+
+                if pivot := seq_get(source.args.get("pivots") or [], 0):
+                    name = source.name
+
+                table_this = source.this
+                table_alias = source.args.get("alias")
+                function_columns: t.List[t.Union[str, exp.Identifier]] = []
+                if isinstance(table_this, exp.Func):
+                    if not table_alias:
+                        function_columns = ensure_list(
+                            dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES.get(type(table_this))
+                        )
+                    elif columns := table_alias.columns:
+                        function_columns = columns
+                    elif type(table_this) in dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES:
+                        function_columns = ensure_list(source.alias_or_name)
+                        source.set("alias", None)
+                        name = None
+
+                _set_alias(
+                    source,
+                    canonical_aliases,
+                    target_alias=name or source.name or None,
+                    normalize=True,
+                    columns=function_columns,
                 )
 
-                if pivots:
-                    pivot = pivots[0]
-                    if not pivot.alias:
-                        pivot_alias = normalize_identifiers(
-                            source.alias if pivot.unpivot else alias_sequence(),
-                            dialect=dialect,
-                        )
-                        pivot.set("alias", exp.TableAlias(this=exp.to_identifier(pivot_alias)))
+                source_fqn = ".".join(p.name for p in source.parts)
+                table_aliases[source_fqn] = source.args["alias"].this.copy()
+
+                if pivot:
+                    target_alias = source.alias if pivot.unpivot else None
+                    _set_alias(pivot, canonical_aliases, target_alias=target_alias, normalize=True)
 
                     # This case corresponds to a pivoted CTE, we don't want to qualify that
                     if isinstance(scope.sources.get(source.alias_or_name), Scope):
                         continue
 
-                _qualify(source)
+                if is_real_table_source:
+                    _qualify(source)
 
-                if on_qualify:
-                    on_qualify(source)
+                    if on_qualify:
+                        on_qualify(source)
             elif isinstance(source, Scope) and source.is_udtf:
-                udtf = source.expression
-                table_alias = udtf.args.get("alias") or exp.TableAlias(
-                    this=exp.to_identifier(next_alias_name())
-                )
-                udtf.set("alias", table_alias)
+                _set_alias(udtf := source.expression, canonical_aliases)
+
+                table_alias = udtf.args["alias"]
 
-                if not table_alias.name:
-                    table_alias.set("this", exp.to_identifier(next_alias_name()))
                 if isinstance(udtf, exp.Values) and not table_alias.columns:
                     column_aliases = [
                         normalize_identifiers(i, dialect=dialect)
                         for i in dialect.generate_values_aliases(udtf)
                     ]
                     table_alias.set("columns", column_aliases)
-            else:
-                for node in scope.walk():
-                    if (
-                        isinstance(node, exp.Table)
-                        and not node.alias
-                        and isinstance(node.parent, (exp.From, exp.Join))
-                    ):
-                        # Mutates the table by attaching an alias to it
-                        alias(node, node.name, copy=False, table=True)
-
-        for column in scope.columns:
+
+        for table in scope.tables:
+            if not table.alias and isinstance(table.parent, (exp.From, exp.Join)):
+                _set_alias(table, canonical_aliases, target_alias=table.name)
+
+        for column in local_columns:
+            table = column.table
+
             if column.db:
                 table_alias = table_aliases.get(".".join(p.name for p in column.parts[0:-1]))
 
@@ -165,5 +200,12 @@ def qualify_tables(
                         column.set(p, None)
 
                     column.set("table", table_alias.copy())
+            elif (
+                canonical_aliases
+                and table
+                and (canonical_table := canonical_aliases.get(table, "")) != column.table
+            ):
+                # Amend existing aliases, e.g. t.c -> _0.c if t is aliased to _0
+                column.set("table", exp.to_identifier(canonical_table))
 
     return expression

sqlglot/optimizer/resolver.py

@@ -0,0 +1,374 @@
+from __future__ import annotations
+
+import itertools
+import typing as t
+
+from sqlglot import exp
+from sqlglot.dialects.dialect import Dialect
+from sqlglot.errors import OptimizeError
+from sqlglot.helper import seq_get, SingleValuedMapping
+from sqlglot.optimizer.scope import Scope
+
+if t.TYPE_CHECKING:
+    from sqlglot.schema import Schema
+
+
+class Resolver:
+    """
+    Helper for resolving columns.
+
+    This is a class so we can lazily load some things and easily share them across functions.
+    """
+
+    def __init__(self, scope: Scope, schema: Schema, infer_schema: bool = True):
+        self.scope = scope
+        self.schema = schema
+        self.dialect = schema.dialect or Dialect()
+        self._source_columns: t.Optional[t.Dict[str, t.Sequence[str]]] = None
+        self._unambiguous_columns: t.Optional[t.Mapping[str, str]] = None
+        self._all_columns: t.Optional[t.Set[str]] = None
+        self._infer_schema = infer_schema
+        self._get_source_columns_cache: t.Dict[t.Tuple[str, bool], t.Sequence[str]] = {}
+
+    def get_table(self, column: str | exp.Column) -> t.Optional[exp.Identifier]:
+        """
+        Get the table for a column name.
+
+        Args:
+            column: The column expression (or column name) to find the table for.
+        Returns:
+            The table name if it can be found/inferred.
+        """
+        column_name = column if isinstance(column, str) else column.name
+
+        table_name = self._get_table_name_from_sources(column_name)
+
+        if not table_name and isinstance(column, exp.Column):
+            # Fall-back case: If we couldn't find the `table_name` from ALL of the sources,
+            # attempt to disambiguate the column based on other characteristics e.g if this column is in a join condition,
+            # we may be able to disambiguate based on the source order.
+            if join_context := self._get_column_join_context(column):
+                # In this case, the return value will be the join that _may_ be able to disambiguate the column
+                # and we can use the source columns available at that join to get the table name
+                # catch OptimizeError if column is still ambiguous and try to resolve with schema inference below
+                try:
+                    table_name = self._get_table_name_from_sources(
+                        column_name, self._get_available_source_columns(join_context)
+                    )
+                except OptimizeError:
+                    pass
+
+        if not table_name and self._infer_schema:
+            sources_without_schema = tuple(
+                source
+                for source, columns in self._get_all_source_columns().items()
+                if not columns or "*" in columns
+            )
+            if len(sources_without_schema) == 1:
+                table_name = sources_without_schema[0]
+
+        if table_name not in self.scope.selected_sources:
+            return exp.to_identifier(table_name)
+
+        node, _ = self.scope.selected_sources.get(table_name)
+
+        if isinstance(node, exp.Query):
+            while node and node.alias != table_name:
+                node = node.parent
+
+        node_alias = node.args.get("alias")
+        if node_alias:
+            return exp.to_identifier(node_alias.this)
+
+        return exp.to_identifier(table_name)
+
+    @property
+    def all_columns(self) -> t.Set[str]:
+        """All available columns of all sources in this scope"""
+        if self._all_columns is None:
+            self._all_columns = {
+                column for columns in self._get_all_source_columns().values() for column in columns
+            }
+        return self._all_columns
+
+    def get_source_columns_from_set_op(self, expression: exp.Expression) -> t.List[str]:
+        if isinstance(expression, exp.Select):
+            return expression.named_selects
+        if isinstance(expression, exp.Subquery) and isinstance(expression.this, exp.SetOperation):
+            # Different types of SET modifiers can be chained together if they're explicitly grouped by nesting
+            return self.get_source_columns_from_set_op(expression.this)
+        if not isinstance(expression, exp.SetOperation):
+            raise OptimizeError(f"Unknown set operation: {expression}")
+
+        set_op = expression
+
+        # BigQuery specific set operations modifiers, e.g INNER UNION ALL BY NAME
+        on_column_list = set_op.args.get("on")
+
+        if on_column_list:
+            # The resulting columns are the columns in the ON clause:
+            # {INNER | LEFT | FULL} UNION ALL BY NAME ON (col1, col2, ...)
+            columns = [col.name for col in on_column_list]
+        elif set_op.side or set_op.kind:
+            side = set_op.side
+            kind = set_op.kind
+
+            # Visit the children UNIONs (if any) in a post-order traversal
+            left = self.get_source_columns_from_set_op(set_op.left)
+            right = self.get_source_columns_from_set_op(set_op.right)
+
+            # We use dict.fromkeys to deduplicate keys and maintain insertion order
+            if side == "LEFT":
+                columns = left
+            elif side == "FULL":
+                columns = list(dict.fromkeys(left + right))
+            elif kind == "INNER":
+                columns = list(dict.fromkeys(left).keys() & dict.fromkeys(right).keys())
+        else:
+            columns = set_op.named_selects
+
+        return columns
+
+    def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequence[str]:
+        """Resolve the source columns for a given source `name`."""
+        cache_key = (name, only_visible)
+        if cache_key not in self._get_source_columns_cache:
+            if name not in self.scope.sources:
+                raise OptimizeError(f"Unknown table: {name}")
+
+            source = self.scope.sources[name]
+
+            if isinstance(source, exp.Table):
+                columns = self.schema.column_names(source, only_visible)
+            elif isinstance(source, Scope) and isinstance(
+                source.expression, (exp.Values, exp.Unnest)
+            ):
+                columns = source.expression.named_selects
+
+                # in bigquery, unnest structs are automatically scoped as tables, so you can
+                # directly select a struct field in a query.
+                # this handles the case where the unnest is statically defined.
+                if self.dialect.UNNEST_COLUMN_ONLY and isinstance(source.expression, exp.Unnest):
+                    unnest = source.expression
+
+                    # if type is not annotated yet, try to get it from the schema
+                    if not unnest.type or unnest.type.is_type(exp.DataType.Type.UNKNOWN):
+                        unnest_expr = seq_get(unnest.expressions, 0)
+                        if isinstance(unnest_expr, exp.Column) and self.scope.parent:
+                            col_type = self._get_unnest_column_type(unnest_expr)
+                            # extract element type if it's an ARRAY
+                            if col_type and col_type.is_type(exp.DataType.Type.ARRAY):
+                                element_types = col_type.expressions
+                                if element_types:
+                                    unnest.type = element_types[0].copy()
+                            else:
+                                if col_type:
+                                    unnest.type = col_type.copy()
+                    # check if the result type is a STRUCT - extract struct field names
+                    if unnest.is_type(exp.DataType.Type.STRUCT):
+                        for k in unnest.type.expressions:  # type: ignore
+                            columns.append(k.name)
+            elif isinstance(source, Scope) and isinstance(source.expression, exp.SetOperation):
+                columns = self.get_source_columns_from_set_op(source.expression)
+
+            else:
+                select = seq_get(source.expression.selects, 0)
+
+                if isinstance(select, exp.QueryTransform):
+                    # https://spark.apache.org/docs/3.5.1/sql-ref-syntax-qry-select-transform.html
+                    schema = select.args.get("schema")
+                    columns = [c.name for c in schema.expressions] if schema else ["key", "value"]
+                else:
+                    columns = source.expression.named_selects
+
+            node, _ = self.scope.selected_sources.get(name) or (None, None)
+            if isinstance(node, Scope):
+                column_aliases = node.expression.alias_column_names
+            elif isinstance(node, exp.Expression):
+                column_aliases = node.alias_column_names
+            else:
+                column_aliases = []
+
+            if column_aliases:
+                # If the source's columns are aliased, their aliases shadow the corresponding column names.
+                # This can be expensive if there are lots of columns, so only do this if column_aliases exist.
+                columns = [
+                    alias or name
+                    for (name, alias) in itertools.zip_longest(columns, column_aliases)
+                ]
+
+            self._get_source_columns_cache[cache_key] = columns
+
+        return self._get_source_columns_cache[cache_key]
+
+    def _get_all_source_columns(self) -> t.Dict[str, t.Sequence[str]]:
+        if self._source_columns is None:
+            self._source_columns = {
+                source_name: self.get_source_columns(source_name)
+                for source_name, source in itertools.chain(
+                    self.scope.selected_sources.items(), self.scope.lateral_sources.items()
+                )
+            }
+        return self._source_columns
+
+    def _get_table_name_from_sources(
+        self, column_name: str, source_columns: t.Optional[t.Dict[str, t.Sequence[str]]] = None
+    ) -> t.Optional[str]:
+        if not source_columns:
+            # If not supplied, get all sources to calculate unambiguous columns
+            if self._unambiguous_columns is None:
+                self._unambiguous_columns = self._get_unambiguous_columns(
+                    self._get_all_source_columns()
+                )
+
+            unambiguous_columns = self._unambiguous_columns
+        else:
+            unambiguous_columns = self._get_unambiguous_columns(source_columns)
+
+        return unambiguous_columns.get(column_name)
+
+    def _get_column_join_context(self, column: exp.Column) -> t.Optional[exp.Join]:
+        """
+        Check if a column participating in a join can be qualified based on the source order.
+        """
+        args = self.scope.expression.args
+        joins = args.get("joins")
+
+        if not joins or args.get("laterals") or args.get("pivots"):
+            # Feature gap: We currently don't try to disambiguate columns if other sources
+            # (e.g laterals, pivots) exist alongside joins
+            return None
+
+        join_ancestor = column.find_ancestor(exp.Join, exp.Select)
+
+        if (
+            isinstance(join_ancestor, exp.Join)
+            and join_ancestor.alias_or_name in self.scope.selected_sources
+        ):
+            # Ensure that the found ancestor is a join that contains an actual source,
+            # e.g in Clickhouse `b` is an array expression in `a ARRAY JOIN b`
+            return join_ancestor
+
+        return None
+
+    def _get_available_source_columns(
+        self, join_ancestor: exp.Join
+    ) -> t.Dict[str, t.Sequence[str]]:
+        """
+        Get the source columns that are available at the point where a column is referenced.
+
+        For columns in JOIN conditions, this only includes tables that have been joined
+        up to that point. Example:
+
+        ```
+        SELECT * FROM t_1 INNER JOIN ... INNER JOIN t_n ON t_1.a = c INNER JOIN t_n+1 ON ...
+        ```                                                        ^
+                                                                   |
+                                +----------------------------------+
+                                |
+                                ⌄
+        The unqualified column `c` is not ambiguous if no other sources up until that
+        join i.e t_1, ..., t_n, contain a column named `c`.
+
+        """
+        args = self.scope.expression.args
+
+        # Collect tables in order: FROM clause tables + joined tables up to current join
+        from_name = args["from_"].alias_or_name
+        available_sources = {from_name: self.get_source_columns(from_name)}
+
+        for join in args["joins"][: t.cast(int, join_ancestor.index) + 1]:
+            available_sources[join.alias_or_name] = self.get_source_columns(join.alias_or_name)
+
+        return available_sources
+
+    def _get_unambiguous_columns(
+        self, source_columns: t.Dict[str, t.Sequence[str]]
+    ) -> t.Mapping[str, str]:
+        """
+        Find all the unambiguous columns in sources.
+
+        Args:
+            source_columns: Mapping of names to source columns.
+
+        Returns:
+            Mapping of column name to source name.
+        """
+        if not source_columns:
+            return {}
+
+        source_columns_pairs = list(source_columns.items())
+
+        first_table, first_columns = source_columns_pairs[0]
+
+        if len(source_columns_pairs) == 1:
+            # Performance optimization - avoid copying first_columns if there is only one table.
+            return SingleValuedMapping(first_columns, first_table)
+
+        unambiguous_columns = {col: first_table for col in first_columns}
+        all_columns = set(unambiguous_columns)
+
+        for table, columns in source_columns_pairs[1:]:
+            unique = set(columns)
+            ambiguous = all_columns.intersection(unique)
+            all_columns.update(columns)
+
+            for column in ambiguous:
+                unambiguous_columns.pop(column, None)
+            for column in unique.difference(ambiguous):
+                unambiguous_columns[column] = table
+
+        return unambiguous_columns
+
+    def _get_unnest_column_type(self, column: exp.Column) -> t.Optional[exp.DataType]:
+        """
+        Get the type of a column being unnested, tracing through CTEs/subqueries to find the base table.
+
+        Args:
+            column: The column expression being unnested.
+
+        Returns:
+            The DataType of the column, or None if not found.
+        """
+        scope = self.scope.parent
+
+        # if column is qualified, use that table, otherwise disambiguate using the resolver
+        if column.table:
+            table_name = column.table
+        else:
+            # use the parent scope's resolver to disambiguate the column
+            parent_resolver = Resolver(scope, self.schema, self._infer_schema)
+            table_identifier = parent_resolver.get_table(column)
+            if not table_identifier:
+                return None
+            table_name = table_identifier.name
+
+        source = scope.sources.get(table_name)
+        return self._get_column_type_from_scope(source, column) if source else None
+
+    def _get_column_type_from_scope(
+        self, source: t.Union[Scope, exp.Table], column: exp.Column
+    ) -> t.Optional[exp.DataType]:
+        """
+        Get a column's type by tracing through scopes/tables to find the base table.
+
+        Args:
+            source: The source to search - can be a Scope (to iterate its sources) or a Table.
+            column: The column to find the type for.
+
+        Returns:
+            The DataType of the column, or None if not found.
+        """
+        if isinstance(source, exp.Table):
+            # base table - get the column type from schema
+            col_type: t.Optional[exp.DataType] = self.schema.get_column_type(source, column)
+            if col_type and not col_type.is_type(exp.DataType.Type.UNKNOWN):
+                return col_type
+        elif isinstance(source, Scope):
+            # iterate over all sources in the scope
+            for source_name, nested_source in source.sources.items():
+                col_type = self._get_column_type_from_scope(nested_source, column)
+                if col_type and not col_type.is_type(exp.DataType.Type.UNKNOWN):
+                    return col_type
+
+        return None