Edit on GitHub

sqlglot.optimizer.qualify_tables

  1from __future__ import annotations
  2
  3import itertools
  4import typing as t
  5
  6from sqlglot import alias, exp
  7from sqlglot._typing import E
  8from sqlglot.dialects.dialect import DialectType
  9from sqlglot.helper import csv_reader, name_sequence
 10from sqlglot.optimizer.scope import Scope, traverse_scope
 11from sqlglot.schema import Schema
 12
 13
 14def qualify_tables(
 15    expression: E,
 16    db: t.Optional[str | exp.Identifier] = None,
 17    catalog: t.Optional[str | exp.Identifier] = None,
 18    schema: t.Optional[Schema] = None,
 19    dialect: DialectType = None,
 20) -> E:
 21    """
 22    Rewrite sqlglot AST to have fully qualified tables. Join constructs such as
 23    (t1 JOIN t2) AS t will be expanded into (SELECT * FROM t1 AS t1, t2 AS t2) AS t.
 24
 25    Examples:
 26        >>> import sqlglot
 27        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
 28        >>> qualify_tables(expression, db="db").sql()
 29        'SELECT 1 FROM db.tbl AS tbl'
 30        >>>
 31        >>> expression = sqlglot.parse_one("SELECT 1 FROM (t1 JOIN t2) AS t")
 32        >>> qualify_tables(expression).sql()
 33        'SELECT 1 FROM (SELECT * FROM t1 AS t1, t2 AS t2) AS t'
 34
 35    Args:
 36        expression: Expression to qualify
 37        db: Database name
 38        catalog: Catalog name
 39        schema: A schema to populate
 40        dialect: The dialect to parse catalog and schema into.
 41
 42    Returns:
 43        The qualified expression.
 44    """
 45    next_alias_name = name_sequence("_q_")
 46    db = exp.parse_identifier(db, dialect=dialect) if db else None
 47    catalog = exp.parse_identifier(catalog, dialect=dialect) if catalog else None
 48
 49    for scope in traverse_scope(expression):
 50        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
 51            if isinstance(derived_table, exp.Subquery):
 52                unnested = derived_table.unnest()
 53                if isinstance(unnested, exp.Table):
 54                    joins = unnested.args.pop("joins", None)
 55                    derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
 56                    derived_table.this.set("joins", joins)
 57
 58            if not derived_table.args.get("alias"):
 59                alias_ = next_alias_name()
 60                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
 61                scope.rename_source(None, alias_)
 62
 63            pivots = derived_table.args.get("pivots")
 64            if pivots and not pivots[0].alias:
 65                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
 66
 67        for name, source in scope.sources.items():
 68            if isinstance(source, exp.Table):
 69                if isinstance(source.this, exp.Identifier):
 70                    if not source.args.get("db"):
 71                        source.set("db", db)
 72                    if not source.args.get("catalog") and source.args.get("db"):
 73                        source.set("catalog", catalog)
 74
 75                if not source.alias:
 76                    # Mutates the source by attaching an alias to it
 77                    alias(source, name or source.name or next_alias_name(), copy=False, table=True)
 78
 79                pivots = source.args.get("pivots")
 80                if pivots and not pivots[0].alias:
 81                    pivots[0].set(
 82                        "alias", exp.TableAlias(this=exp.to_identifier(next_alias_name()))
 83                    )
 84
 85                if schema and isinstance(source.this, exp.ReadCSV):
 86                    with csv_reader(source.this) as reader:
 87                        header = next(reader)
 88                        columns = next(reader)
 89                        schema.add_table(
 90                            source,
 91                            {k: type(v).__name__ for k, v in zip(header, columns)},
 92                            match_depth=False,
 93                        )
 94            elif isinstance(source, Scope) and source.is_udtf:
 95                udtf = source.expression
 96                table_alias = udtf.args.get("alias") or exp.TableAlias(
 97                    this=exp.to_identifier(next_alias_name())
 98                )
 99                udtf.set("alias", table_alias)
100
101                if not table_alias.name:
102                    table_alias.set("this", exp.to_identifier(next_alias_name()))
103                if isinstance(udtf, exp.Values) and not table_alias.columns:
104                    for i, e in enumerate(udtf.expressions[0].expressions):
105                        table_alias.append("columns", exp.to_identifier(f"_col_{i}"))
106
107    return expression
def qualify_tables( expression: ~E, db: Union[sqlglot.expressions.Identifier, str, NoneType] = None, catalog: Union[sqlglot.expressions.Identifier, str, NoneType] = None, schema: Optional[sqlglot.schema.Schema] = None, dialect: Union[str, sqlglot.dialects.dialect.Dialect, Type[sqlglot.dialects.dialect.Dialect], NoneType] = None) -> ~E:
 15def qualify_tables(
 16    expression: E,
 17    db: t.Optional[str | exp.Identifier] = None,
 18    catalog: t.Optional[str | exp.Identifier] = None,
 19    schema: t.Optional[Schema] = None,
 20    dialect: DialectType = None,
 21) -> E:
 22    """
 23    Rewrite sqlglot AST to have fully qualified tables. Join constructs such as
 24    (t1 JOIN t2) AS t will be expanded into (SELECT * FROM t1 AS t1, t2 AS t2) AS t.
 25
 26    Examples:
 27        >>> import sqlglot
 28        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
 29        >>> qualify_tables(expression, db="db").sql()
 30        'SELECT 1 FROM db.tbl AS tbl'
 31        >>>
 32        >>> expression = sqlglot.parse_one("SELECT 1 FROM (t1 JOIN t2) AS t")
 33        >>> qualify_tables(expression).sql()
 34        'SELECT 1 FROM (SELECT * FROM t1 AS t1, t2 AS t2) AS t'
 35
 36    Args:
 37        expression: Expression to qualify
 38        db: Database name
 39        catalog: Catalog name
 40        schema: A schema to populate
 41        dialect: The dialect to parse catalog and schema into.
 42
 43    Returns:
 44        The qualified expression.
 45    """
 46    next_alias_name = name_sequence("_q_")
 47    db = exp.parse_identifier(db, dialect=dialect) if db else None
 48    catalog = exp.parse_identifier(catalog, dialect=dialect) if catalog else None
 49
 50    for scope in traverse_scope(expression):
 51        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
 52            if isinstance(derived_table, exp.Subquery):
 53                unnested = derived_table.unnest()
 54                if isinstance(unnested, exp.Table):
 55                    joins = unnested.args.pop("joins", None)
 56                    derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
 57                    derived_table.this.set("joins", joins)
 58
 59            if not derived_table.args.get("alias"):
 60                alias_ = next_alias_name()
 61                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
 62                scope.rename_source(None, alias_)
 63
 64            pivots = derived_table.args.get("pivots")
 65            if pivots and not pivots[0].alias:
 66                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
 67
 68        for name, source in scope.sources.items():
 69            if isinstance(source, exp.Table):
 70                if isinstance(source.this, exp.Identifier):
 71                    if not source.args.get("db"):
 72                        source.set("db", db)
 73                    if not source.args.get("catalog") and source.args.get("db"):
 74                        source.set("catalog", catalog)
 75
 76                if not source.alias:
 77                    # Mutates the source by attaching an alias to it
 78                    alias(source, name or source.name or next_alias_name(), copy=False, table=True)
 79
 80                pivots = source.args.get("pivots")
 81                if pivots and not pivots[0].alias:
 82                    pivots[0].set(
 83                        "alias", exp.TableAlias(this=exp.to_identifier(next_alias_name()))
 84                    )
 85
 86                if schema and isinstance(source.this, exp.ReadCSV):
 87                    with csv_reader(source.this) as reader:
 88                        header = next(reader)
 89                        columns = next(reader)
 90                        schema.add_table(
 91                            source,
 92                            {k: type(v).__name__ for k, v in zip(header, columns)},
 93                            match_depth=False,
 94                        )
 95            elif isinstance(source, Scope) and source.is_udtf:
 96                udtf = source.expression
 97                table_alias = udtf.args.get("alias") or exp.TableAlias(
 98                    this=exp.to_identifier(next_alias_name())
 99                )
100                udtf.set("alias", table_alias)
101
102                if not table_alias.name:
103                    table_alias.set("this", exp.to_identifier(next_alias_name()))
104                if isinstance(udtf, exp.Values) and not table_alias.columns:
105                    for i, e in enumerate(udtf.expressions[0].expressions):
106                        table_alias.append("columns", exp.to_identifier(f"_col_{i}"))
107
108    return expression

Rewrite sqlglot AST to have fully qualified tables. Join constructs such as (t1 JOIN t2) AS t will be expanded into (SELECT * FROM t1 AS t1, t2 AS t2) AS t.

Examples:
>>> import sqlglot
>>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
>>> qualify_tables(expression, db="db").sql()
'SELECT 1 FROM db.tbl AS tbl'
>>>
>>> expression = sqlglot.parse_one("SELECT 1 FROM (t1 JOIN t2) AS t")
>>> qualify_tables(expression).sql()
'SELECT 1 FROM (SELECT * FROM t1 AS t1, t2 AS t2) AS t'
Arguments:
  • expression: Expression to qualify
  • db: Database name
  • catalog: Catalog name
  • schema: A schema to populate
  • dialect: The dialect to parse catalog and schema into.
Returns:

The qualified expression.