sqlglot.optimizer.qualify_tables API documentation

Edit on GitHub

# sqlglot.optimizer.qualify_tables

View Source

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

     17def qualify_tables(
     18    expression: E,
     19    db: t.Optional[str | exp.Identifier] = None,
     20    catalog: t.Optional[str | exp.Identifier] = None,
     21    schema: t.Optional[Schema] = None,
     22    dialect: DialectType = None,
     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: Expression to qualify
     40        db: Database name
     41        catalog: Catalog name
     42        schema: A schema to populate
     43        dialect: The dialect to parse catalog and schema into.
     44
     45    Returns:
     46        The qualified expression.
     47    """
     48    next_alias_name = name_sequence("_q_")
     49    db = exp.parse_identifier(db, dialect=dialect) if db else None
     50    catalog = exp.parse_identifier(catalog, dialect=dialect) if catalog else None
     51
     52    def _qualify(table: exp.Table) -> None:
     53        if isinstance(table.this, exp.Identifier):
     54            if not table.args.get("db"):
     55                table.set("db", db)
     56            if not table.args.get("catalog") and table.args.get("db"):
     57                table.set("catalog", catalog)
     58
     59    if not isinstance(expression, exp.Query):
     60        for node, *_ in expression.walk(prune=lambda n, *_: isinstance(n, exp.Query)):
     61            if isinstance(node, exp.Table):
     62                _qualify(node)
     63
     64    for scope in traverse_scope(expression):
     65        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
     66            if isinstance(derived_table, exp.Subquery):
     67                unnested = derived_table.unnest()
     68                if isinstance(unnested, exp.Table):
     69                    joins = unnested.args.pop("joins", None)
     70                    derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
     71                    derived_table.this.set("joins", joins)
     72
     73            if not derived_table.args.get("alias"):
     74                alias_ = next_alias_name()
     75                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
     76                scope.rename_source(None, alias_)
     77
     78            pivots = derived_table.args.get("pivots")
     79            if pivots and not pivots[0].alias:
     80                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
     81
     82        for name, source in scope.sources.items():
     83            if isinstance(source, exp.Table):
     84                _qualify(source)
     85
     86                pivots = pivots = source.args.get("pivots")
     87                if not source.alias:
     88                    # Don't add the pivot's alias to the pivoted table, use the table's name instead
     89                    if pivots and pivots[0].alias == name:
     90                        name = source.name
     91
     92                    # Mutates the source by attaching an alias to it
     93                    alias(source, name or source.name or next_alias_name(), copy=False, table=True)
     94
     95                if pivots and not pivots[0].alias:
     96                    pivots[0].set(
     97                        "alias", exp.TableAlias(this=exp.to_identifier(next_alias_name()))
     98                    )
     99
    100                if schema and isinstance(source.this, exp.ReadCSV):
    101                    with csv_reader(source.this) as reader:
    102                        header = next(reader)
    103                        columns = next(reader)
    104                        schema.add_table(
    105                            source,
    106                            {k: type(v).__name__ for k, v in zip(header, columns)},
    107                            match_depth=False,
    108                        )
    109            elif isinstance(source, Scope) and source.is_udtf:
    110                udtf = source.expression
    111                table_alias = udtf.args.get("alias") or exp.TableAlias(
    112                    this=exp.to_identifier(next_alias_name())
    113                )
    114                udtf.set("alias", table_alias)
    115
    116                if not table_alias.name:
    117                    table_alias.set("this", exp.to_identifier(next_alias_name()))
    118                if isinstance(udtf, exp.Values) and not table_alias.columns:
    119                    for i, e in enumerate(udtf.expressions[0].expressions):
    120                        table_alias.append("columns", exp.to_identifier(f"_col_{i}"))
    121            else:
    122                for node, parent, _ in scope.walk():
    123                    if (
    124                        isinstance(node, exp.Table)
    125                        and not node.alias
    126                        and isinstance(parent, (exp.From, exp.Join))
    127                    ):
    128                        # Mutates the table by attaching an alias to it
    129                        alias(node, node.name, copy=False, table=True)
    130
    131    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.