sqlglot.optimizer.qualify_tables

 1import itertools
 2
 3from sqlglot import alias, exp
 4from sqlglot.helper import csv_reader
 5from sqlglot.optimizer.scope import Scope, traverse_scope
 6
 7
 8def qualify_tables(expression, db=None, catalog=None, schema=None):
 9    """
10    Rewrite sqlglot AST to have fully qualified tables.
11
12    Example:
13        >>> import sqlglot
14        >>> expression = sqlglot.parse_one("SELECT 1 FROM tbl")
15        >>> qualify_tables(expression, db="db").sql()
16        'SELECT 1 FROM db.tbl AS tbl'
17
18    Args:
19        expression (sqlglot.Expression): expression to qualify
20        db (str): Database name
21        catalog (str): Catalog name
22        schema: A schema to populate
23    Returns:
24        sqlglot.Expression: qualified expression
25    """
26    sequence = itertools.count()
27
28    next_name = lambda: f"_q_{next(sequence)}"
29
30    for scope in traverse_scope(expression):
31        for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
32            if not derived_table.args.get("alias"):
33                alias_ = f"_q_{next(sequence)}"
34                derived_table.set("alias", exp.TableAlias(this=exp.to_identifier(alias_)))
35                scope.rename_source(None, alias_)
36
37        for source in scope.sources.values():
38            if isinstance(source, exp.Table):
39                identifier = isinstance(source.this, exp.Identifier)
40
41                if identifier:
42                    if not source.args.get("db"):
43                        source.set("db", exp.to_identifier(db))
44                    if not source.args.get("catalog"):
45                        source.set("catalog", exp.to_identifier(catalog))
46
47                if not source.alias:
48                    source = source.replace(
49                        alias(
50                            source.copy(),
51                            source.this if identifier else next_name(),
52                            table=True,
53                        )
54                    )
55
56                if schema and isinstance(source.this, exp.ReadCSV):
57                    with csv_reader(source.this) as reader:
58                        header = next(reader)
59                        columns = next(reader)
60                        schema.add_table(
61                            source, {k: type(v).__name__ for k, v in zip(header, columns)}
62                        )
63            elif isinstance(source, Scope) and source.is_udtf:
64                udtf = source.expression
65                table_alias = udtf.args.get("alias") or exp.TableAlias(this=next_name())
66                udtf.set("alias", table_alias)
67
68                if not table_alias.name:
69                    table_alias.set("this", next_name())
70
71    return expression