sqlglot.optimizer.qualify_tables

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.dialects.dialect import Dialect, DialectType
  7from sqlglot.helper import name_sequence, seq_get, ensure_list
  8from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
  9from sqlglot.optimizer.scope import Scope, traverse_scope
 10
 11if t.TYPE_CHECKING:
 12    from sqlglot._typing import E
 13    from collections.abc import Sequence
 14
 15
 16def qualify_tables(
 17    expression: E,
 18    db: str | exp.Identifier | None = None,
 19    catalog: str | exp.Identifier | None = None,
 20    on_qualify: t.Callable[[exp.Table], None] | None = None,
 21    dialect: DialectType = None,
 22    canonicalize_table_aliases: bool = False,
 23) -> E:
 24    """
 25    Rewrite sqlglot AST to have fully qualified tables. Join constructs such as
 26    (t1 JOIN t2) AS t will be expanded into (SELECT * FROM t1 AS t1, t2 AS t2) AS t.
 27
 28    Examples:
 29        >>> import sqlglot
 30        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
 31        >>> qualify_tables(expression, db="db").sql()
 32        'SELECT 1 FROM db.tbl AS tbl'
 33        >>>
 34        >>> expression = sqlglot.parse_one("SELECT 1 FROM (t1 JOIN t2) AS t")
 35        >>> qualify_tables(expression).sql()
 36        'SELECT 1 FROM (SELECT * FROM t1 AS t1, t2 AS t2) AS t'
 37
 38    Args:
 39        expression: Expr to qualify
 40        db: Database name
 41        catalog: Catalog name
 42        on_qualify: Callback after a table has been qualified.
 43        dialect: The dialect to parse catalog and schema into.
 44        canonicalize_table_aliases: Whether to use canonical aliases (_0, _1, ...) for all sources
 45            instead of preserving table names. Defaults to False.
 46
 47    Returns:
 48        The qualified expression.
 49    """
 50    dialect = Dialect.get_or_raise(dialect)
 51    next_alias_name = name_sequence("_")
 52
 53    if db := db or None:
 54        db = exp.parse_identifier(db, dialect=dialect)
 55        db.meta["is_table"] = True
 56        db = normalize_identifiers(db, dialect=dialect)
 57    if catalog := catalog or None:
 58        catalog = exp.parse_identifier(catalog, dialect=dialect)
 59        catalog.meta["is_table"] = True
 60        catalog = normalize_identifiers(catalog, dialect=dialect)
 61
 62    def _qualify(table: exp.Table) -> None:
 63        if isinstance(table.this, exp.Identifier):
 64            if db and not table.args.get("db"):
 65                table.set("db", db.copy())
 66            if catalog and not table.args.get("catalog") and table.args.get("db"):
 67                table.set("catalog", catalog.copy())
 68
 69    if (db or catalog) and not isinstance(expression, exp.Query):
 70        with_ = expression.args.get("with_") or exp.With()
 71        cte_names = {cte.alias_or_name for cte in with_.expressions}
 72
 73        for node in expression.walk(prune=lambda n: isinstance(n, exp.Query)):
 74            if isinstance(node, exp.Table) and node.name not in cte_names:
 75                _qualify(node)
 76
 77    def _set_alias(
 78        expression: exp.Expr,
 79        canonical_aliases: dict[str, str],
 80        target_alias: str | None = None,
 81        scope: Scope | None = None,
 82        normalize: bool = False,
 83        columns: Sequence[str | exp.Identifier] | None = None,
 84    ) -> None:
 85        alias = expression.args.get("alias") or exp.TableAlias()
 86
 87        if canonicalize_table_aliases:
 88            new_alias_name = next_alias_name()
 89            canonical_aliases[alias.name or target_alias or ""] = new_alias_name
 90        elif not alias.name:
 91            new_alias_name = target_alias or next_alias_name()
 92            if normalize and target_alias:
 93                new_alias_name = normalize_identifiers(new_alias_name, dialect=dialect).name
 94        else:
 95            return
 96
 97        alias.set("this", exp.to_identifier(new_alias_name))
 98
 99        if columns:
100            alias.set("columns", [exp.to_identifier(c) for c in columns])
101
102        expression.set("alias", alias)
103
104        if scope:
105            scope.rename_source(None, new_alias_name)
106
107    for scope in traverse_scope(expression):
108        local_columns = scope.local_columns
109        canonical_aliases: dict[str, str] = {}
110
111        for query in scope.subqueries:
112            subquery = query.parent
113            if isinstance(subquery, exp.Subquery):
114                subquery.unwrap().replace(subquery)
115
116        for derived_table in scope.derived_tables:
117            unnested = derived_table.unnest()
118            if isinstance(unnested, exp.Table):
119                joins = unnested.args.get("joins")
120                unnested.set("joins", None)
121                derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
122                derived_table.this.set("joins", joins)
123
124            _set_alias(derived_table, canonical_aliases, scope=scope)
125            if pivot := seq_get(derived_table.args.get("pivots") or [], 0):
126                _set_alias(pivot, canonical_aliases)
127
128        table_aliases = {}
129
130        for name, source in scope.sources.items():
131            if isinstance(source, exp.Table):
132                # When the name is empty, it means that we have a non-table source, e.g. a pivoted cte
133                is_real_table_source = bool(name)
134
135                if pivot := seq_get(source.args.get("pivots") or [], 0):
136                    name = source.name
137
138                table_this = source.this
139                table_alias = source.args.get("alias")
140                function_columns: Sequence[str | exp.Identifier] | None = None
141                if isinstance(table_this, exp.Func):
142                    if not table_alias:
143                        function_columns = ensure_list(
144                            dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES.get(type(table_this))
145                        )
146                    elif columns := table_alias.columns:
147                        function_columns = columns
148                    elif type(table_this) in dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES:
149                        function_columns = ensure_list(source.alias_or_name)
150                        source.set("alias", None)
151                        name = ""
152
153                _set_alias(
154                    source,
155                    canonical_aliases,
156                    target_alias=name or source.name or None,
157                    normalize=True,
158                    columns=function_columns,
159                )
160
161                source_fqn = ".".join(p.name for p in source.parts)
162                had_explicit_alias = table_alias and table_alias.name
163                if not had_explicit_alias or source_fqn not in table_aliases:
164                    table_aliases[source_fqn] = source.args["alias"].this.copy()
165
166                if pivot:
167                    target_alias = source.alias if pivot.unpivot else None
168                    _set_alias(pivot, canonical_aliases, target_alias=target_alias, normalize=True)
169
170                    # This case corresponds to a pivoted CTE, we don't want to qualify that
171                    if isinstance(scope.sources.get(source.alias_or_name), Scope):
172                        continue
173
174                if is_real_table_source:
175                    _qualify(source)
176
177                    if on_qualify:
178                        on_qualify(source)
179            elif isinstance(source, Scope) and source.is_udtf:
180                _set_alias(udtf := source.expression, canonical_aliases)
181
182                table_alias = udtf.args["alias"]
183
184                if isinstance(udtf, exp.Values) and not table_alias.columns:
185                    column_aliases = [
186                        normalize_identifiers(i, dialect=dialect)
187                        for i in dialect.generate_values_aliases(udtf)
188                    ]
189                    table_alias.set("columns", column_aliases)
190
191        for table in scope.tables:
192            if not table.alias and isinstance(table.parent, (exp.From, exp.Join)):
193                _set_alias(table, canonical_aliases, target_alias=table.name)
194
195        for column in local_columns:
196            column_table = column.table
197
198            if column.db:
199                table_alias = table_aliases.get(".".join(p.name for p in column.parts[0:-1]))
200
201                if table_alias:
202                    for p in exp.COLUMN_PARTS[1:]:
203                        column.set(p, None)
204
205                    column.set("table", table_alias.copy())
206            elif (
207                canonical_aliases
208                and column_table
209                and (canonical_table := canonical_aliases.get(column_table, "")) != column_table
210            ):
211                # Amend existing aliases, e.g. t.c -> _0.c if t is aliased to _0
212                column.set("table", exp.to_identifier(canonical_table))
213
214    return expression