sqlglot.optimizer.qualify_tables

  1import itertools
  2import typing as t
  3
  4from sqlglot import alias, exp
  5from sqlglot._typing import E
  6from sqlglot.helper import csv_reader, name_sequence
  7from sqlglot.optimizer.scope import Scope, traverse_scope
  8from sqlglot.schema import Schema
  9
 10
 11def qualify_tables(
 12    expression: E,
 13    db: t.Optional[str] = None,
 14    catalog: t.Optional[str] = None,
 15    schema: t.Optional[Schema] = None,
 16) -> E:
 17    """
 18    Rewrite sqlglot AST to have fully qualified, unnested tables.
 19
 20    Examples:
 21        >>> import sqlglot
 22        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
 23        >>> qualify_tables(expression, db="db").sql()
 24        'SELECT 1 FROM db.tbl AS tbl'
 25        >>>
 26        >>> expression = sqlglot.parse_one("SELECT * FROM (tbl)")
 27        >>> qualify_tables(expression).sql()
 28        'SELECT * FROM tbl AS tbl'
 29        >>>
 30        >>> expression = sqlglot.parse_one("SELECT * FROM (tbl1 JOIN tbl2 ON id1 = id2)")
 31        >>> qualify_tables(expression).sql()
 32        'SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2'
 33
 34    Note:
 35        This rule effectively enforces a left-to-right join order, since all joins
 36        are unnested. This means that the optimizer doesn't necessarily preserve the
 37        original join order, e.g. when parentheses are used to specify it explicitly.
 38
 39    Args:
 40        expression: Expression to qualify
 41        db: Database name
 42        catalog: Catalog name
 43        schema: A schema to populate
 44
 45    Returns:
 46        The qualified expression.
 47    """
 48    next_alias_name = name_sequence("_q_")
 49
 50    for scope in traverse_scope(expression):
 51        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
 52            if not derived_table.args.get("alias"):
 53                alias_ = next_alias_name()
 54                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
 55                scope.rename_source(None, alias_)
 56
 57            pivots = derived_table.args.get("pivots")
 58            if pivots and not pivots[0].alias:
 59                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
 60
 61        for name, source in scope.sources.items():
 62            if isinstance(source, exp.Table):
 63                if isinstance(source.this, exp.Identifier):
 64                    if not source.args.get("db"):
 65                        source.set("db", exp.to_identifier(db))
 66                    if not source.args.get("catalog"):
 67                        source.set("catalog", exp.to_identifier(catalog))
 68
 69                # Unnest joins attached in tables by appending them to the closest query
 70                for join in source.args.get("joins") or []:
 71                    scope.expression.append("joins", join)
 72
 73                source.set("joins", None)
 74                source.set("wrapped", None)
 75
 76                if not source.alias:
 77                    source = source.replace(
 78                        alias(
 79                            source, name or source.name or next_alias_name(), copy=True, table=True
 80                        )
 81                    )
 82
 83                pivots = source.args.get("pivots")
 84                if pivots and not pivots[0].alias:
 85                    pivots[0].set(
 86                        "alias", exp.TableAlias(this=exp.to_identifier(next_alias_name()))
 87                    )
 88
 89                if schema and isinstance(source.this, exp.ReadCSV):
 90                    with csv_reader(source.this) as reader:
 91                        header = next(reader)
 92                        columns = next(reader)
 93                        schema.add_table(
 94                            source, {k: type(v).__name__ for k, v in zip(header, columns)}
 95                        )
 96            elif isinstance(source, Scope) and source.is_udtf:
 97                udtf = source.expression
 98                table_alias = udtf.args.get("alias") or exp.TableAlias(
 99                    this=exp.to_identifier(next_alias_name())
100                )
101                udtf.set("alias", table_alias)
102
103                if not table_alias.name:
104                    table_alias.set("this", exp.to_identifier(next_alias_name()))
105                if isinstance(udtf, exp.Values) and not table_alias.columns:
106                    for i, e in enumerate(udtf.expressions[0].expressions):
107                        table_alias.append("columns", exp.to_identifier(f"_col_{i}"))
108
109    return expression