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
     11from sqlglot.dialects.dialect import Dialect
     12
     13if t.TYPE_CHECKING:
     14    from sqlglot._typing import E
     15
     16
     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    infer_csv_schemas: bool = False,
     23    dialect: DialectType = None,
     24) -> E:
     25    """
     26    Rewrite sqlglot AST to have fully qualified tables. Join constructs such as
     27    (t1 JOIN t2) AS t will be expanded into (SELECT * FROM t1 AS t1, t2 AS t2) AS t.
     28
     29    Examples:
     30        >>> import sqlglot
     31        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
     32        >>> qualify_tables(expression, db="db").sql()
     33        'SELECT 1 FROM db.tbl AS tbl'
     34        >>>
     35        >>> expression = sqlglot.parse_one("SELECT 1 FROM (t1 JOIN t2) AS t")
     36        >>> qualify_tables(expression).sql()
     37        'SELECT 1 FROM (SELECT * FROM t1 AS t1, t2 AS t2) AS t'
     38
     39    Args:
     40        expression: Expression to qualify
     41        db: Database name
     42        catalog: Catalog name
     43        schema: A schema to populate
     44        infer_csv_schemas: Whether to scan READ_CSV calls in order to infer the CSVs' schemas.
     45        dialect: The dialect to parse catalog and schema into.
     46
     47    Returns:
     48        The qualified expression.
     49    """
     50    next_alias_name = name_sequence("_q_")
     51    db = exp.parse_identifier(db, dialect=dialect) if db else None
     52    catalog = exp.parse_identifier(catalog, dialect=dialect) if catalog else None
     53    dialect = Dialect.get_or_raise(dialect)
     54
     55    def _qualify(table: exp.Table) -> None:
     56        if isinstance(table.this, exp.Identifier):
     57            if db and not table.args.get("db"):
     58                table.set("db", db.copy())
     59            if catalog and not table.args.get("catalog") and table.args.get("db"):
     60                table.set("catalog", catalog.copy())
     61
     62    if (db or catalog) and not isinstance(expression, exp.Query):
     63        for node in expression.walk(prune=lambda n: isinstance(n, exp.Query)):
     64            if isinstance(node, exp.Table):
     65                _qualify(node)
     66
     67    for scope in traverse_scope(expression):
     68        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
     69            if isinstance(derived_table, exp.Subquery):
     70                unnested = derived_table.unnest()
     71                if isinstance(unnested, exp.Table):
     72                    joins = unnested.args.pop("joins", None)
     73                    derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
     74                    derived_table.this.set("joins", joins)
     75
     76            if not derived_table.args.get("alias"):
     77                alias_ = next_alias_name()
     78                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
     79                scope.rename_source(None, alias_)
     80
     81            pivots = derived_table.args.get("pivots")
     82            if pivots and not pivots[0].alias:
     83                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
     84
     85        table_aliases = {}
     86
     87        for name, source in scope.sources.items():
     88            if isinstance(source, exp.Table):
     89                pivots = source.args.get("pivots")
     90                if not source.alias:
     91                    # Don't add the pivot's alias to the pivoted table, use the table's name instead
     92                    if pivots and pivots[0].alias == name:
     93                        name = source.name
     94
     95                    # Mutates the source by attaching an alias to it
     96                    alias(source, name or source.name or next_alias_name(), copy=False, table=True)
     97
     98                table_aliases[".".join(p.name for p in source.parts)] = exp.to_identifier(
     99                    source.alias
    100                )
    101
    102                if pivots:
    103                    if not pivots[0].alias:
    104                        pivot_alias = next_alias_name()
    105                        pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(pivot_alias)))
    106
    107                    # This case corresponds to a pivoted CTE, we don't want to qualify that
    108                    if isinstance(scope.sources.get(source.alias_or_name), Scope):
    109                        continue
    110
    111                _qualify(source)
    112
    113                if infer_csv_schemas and schema and isinstance(source.this, exp.ReadCSV):
    114                    with csv_reader(source.this) as reader:
    115                        header = next(reader)
    116                        columns = next(reader)
    117                        schema.add_table(
    118                            source,
    119                            {k: type(v).__name__ for k, v in zip(header, columns)},
    120                            match_depth=False,
    121                        )
    122            elif isinstance(source, Scope) and source.is_udtf:
    123                udtf = source.expression
    124                table_alias = udtf.args.get("alias") or exp.TableAlias(
    125                    this=exp.to_identifier(next_alias_name())
    126                )
    127                udtf.set("alias", table_alias)
    128
    129                if not table_alias.name:
    130                    table_alias.set("this", exp.to_identifier(next_alias_name()))
    131                if isinstance(udtf, exp.Values) and not table_alias.columns:
    132                    column_aliases = dialect.generate_values_aliases(udtf)
    133                    table_alias.set("columns", column_aliases)
    134            else:
    135                for node in scope.walk():
    136                    if (
    137                        isinstance(node, exp.Table)
    138                        and not node.alias
    139                        and isinstance(node.parent, (exp.From, exp.Join))
    140                    ):
    141                        # Mutates the table by attaching an alias to it
    142                        alias(node, node.name, copy=False, table=True)
    143
    144        for column in scope.columns:
    145            if column.db:
    146                table_alias = table_aliases.get(".".join(p.name for p in column.parts[0:-1]))
    147
    148                if table_alias:
    149                    for p in exp.COLUMN_PARTS[1:]:
    150                        column.set(p, None)
    151
    152                    column.set("table", table_alias.copy())
    153
    154    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, infer_csv_schemas: bool = False, dialect: Union[str, sqlglot.dialects.Dialect, Type[sqlglot.dialects.Dialect], NoneType] = None) -> ~E: View Source

     18def qualify_tables(
     19    expression: E,
     20    db: t.Optional[str | exp.Identifier] = None,
     21    catalog: t.Optional[str | exp.Identifier] = None,
     22    schema: t.Optional[Schema] = None,
     23    infer_csv_schemas: bool = False,
     24    dialect: DialectType = None,
     25) -> E:
     26    """
     27    Rewrite sqlglot AST to have fully qualified tables. Join constructs such as
     28    (t1 JOIN t2) AS t will be expanded into (SELECT * FROM t1 AS t1, t2 AS t2) AS t.
     29
     30    Examples:
     31        >>> import sqlglot
     32        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
     33        >>> qualify_tables(expression, db="db").sql()
     34        'SELECT 1 FROM db.tbl AS tbl'
     35        >>>
     36        >>> expression = sqlglot.parse_one("SELECT 1 FROM (t1 JOIN t2) AS t")
     37        >>> qualify_tables(expression).sql()
     38        'SELECT 1 FROM (SELECT * FROM t1 AS t1, t2 AS t2) AS t'
     39
     40    Args:
     41        expression: Expression to qualify
     42        db: Database name
     43        catalog: Catalog name
     44        schema: A schema to populate
     45        infer_csv_schemas: Whether to scan READ_CSV calls in order to infer the CSVs' schemas.
     46        dialect: The dialect to parse catalog and schema into.
     47
     48    Returns:
     49        The qualified expression.
     50    """
     51    next_alias_name = name_sequence("_q_")
     52    db = exp.parse_identifier(db, dialect=dialect) if db else None
     53    catalog = exp.parse_identifier(catalog, dialect=dialect) if catalog else None
     54    dialect = Dialect.get_or_raise(dialect)
     55
     56    def _qualify(table: exp.Table) -> None:
     57        if isinstance(table.this, exp.Identifier):
     58            if db and not table.args.get("db"):
     59                table.set("db", db.copy())
     60            if catalog and not table.args.get("catalog") and table.args.get("db"):
     61                table.set("catalog", catalog.copy())
     62
     63    if (db or catalog) and not isinstance(expression, exp.Query):
     64        for node in expression.walk(prune=lambda n: isinstance(n, exp.Query)):
     65            if isinstance(node, exp.Table):
     66                _qualify(node)
     67
     68    for scope in traverse_scope(expression):
     69        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
     70            if isinstance(derived_table, exp.Subquery):
     71                unnested = derived_table.unnest()
     72                if isinstance(unnested, exp.Table):
     73                    joins = unnested.args.pop("joins", None)
     74                    derived_table.this.replace(exp.select("*").from_(unnested.copy(), copy=False))
     75                    derived_table.this.set("joins", joins)
     76
     77            if not derived_table.args.get("alias"):
     78                alias_ = next_alias_name()
     79                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
     80                scope.rename_source(None, alias_)
     81
     82            pivots = derived_table.args.get("pivots")
     83            if pivots and not pivots[0].alias:
     84                pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(next_alias_name())))
     85
     86        table_aliases = {}
     87
     88        for name, source in scope.sources.items():
     89            if isinstance(source, exp.Table):
     90                pivots = source.args.get("pivots")
     91                if not source.alias:
     92                    # Don't add the pivot's alias to the pivoted table, use the table's name instead
     93                    if pivots and pivots[0].alias == name:
     94                        name = source.name
     95
     96                    # Mutates the source by attaching an alias to it
     97                    alias(source, name or source.name or next_alias_name(), copy=False, table=True)
     98
     99                table_aliases[".".join(p.name for p in source.parts)] = exp.to_identifier(
    100                    source.alias
    101                )
    102
    103                if pivots:
    104                    if not pivots[0].alias:
    105                        pivot_alias = next_alias_name()
    106                        pivots[0].set("alias", exp.TableAlias(this=exp.to_identifier(pivot_alias)))
    107
    108                    # This case corresponds to a pivoted CTE, we don't want to qualify that
    109                    if isinstance(scope.sources.get(source.alias_or_name), Scope):
    110                        continue
    111
    112                _qualify(source)
    113
    114                if infer_csv_schemas and schema and isinstance(source.this, exp.ReadCSV):
    115                    with csv_reader(source.this) as reader:
    116                        header = next(reader)
    117                        columns = next(reader)
    118                        schema.add_table(
    119                            source,
    120                            {k: type(v).__name__ for k, v in zip(header, columns)},
    121                            match_depth=False,
    122                        )
    123            elif isinstance(source, Scope) and source.is_udtf:
    124                udtf = source.expression
    125                table_alias = udtf.args.get("alias") or exp.TableAlias(
    126                    this=exp.to_identifier(next_alias_name())
    127                )
    128                udtf.set("alias", table_alias)
    129
    130                if not table_alias.name:
    131                    table_alias.set("this", exp.to_identifier(next_alias_name()))
    132                if isinstance(udtf, exp.Values) and not table_alias.columns:
    133                    column_aliases = dialect.generate_values_aliases(udtf)
    134                    table_alias.set("columns", column_aliases)
    135            else:
    136                for node in scope.walk():
    137                    if (
    138                        isinstance(node, exp.Table)
    139                        and not node.alias
    140                        and isinstance(node.parent, (exp.From, exp.Join))
    141                    ):
    142                        # Mutates the table by attaching an alias to it
    143                        alias(node, node.name, copy=False, table=True)
    144
    145        for column in scope.columns:
    146            if column.db:
    147                table_alias = table_aliases.get(".".join(p.name for p in column.parts[0:-1]))
    148
    149                if table_alias:
    150                    for p in exp.COLUMN_PARTS[1:]:
    151                        column.set(p, None)
    152
    153                    column.set("table", table_alias.copy())
    154
    155    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
- infer_csv_schemas: Whether to scan READ_CSV calls in order to infer the CSVs' schemas.
- dialect: The dialect to parse catalog and schema into.

###### Returns:

> The qualified expression.