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._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: View Source

     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.