Edit on GitHub

# sqlglot.optimizer.qualify_columns

View Source

      1import itertools
      2import typing as t
      3
      4from sqlglot import alias, exp
      5from sqlglot.errors import OptimizeError
      6from sqlglot.optimizer.scope import Scope, traverse_scope
      7from sqlglot.schema import ensure_schema
      8
      9
     10def qualify_columns(expression, schema):
     11    """
     12    Rewrite sqlglot AST to have fully qualified columns.
     13
     14    Example:
     15        >>> import sqlglot
     16        >>> schema = {"tbl": {"col": "INT"}}
     17        >>> expression = sqlglot.parse_one("SELECT col FROM tbl")
     18        >>> qualify_columns(expression, schema).sql()
     19        'SELECT tbl.col AS col FROM tbl'
     20
     21    Args:
     22        expression (sqlglot.Expression): expression to qualify
     23        schema (dict|sqlglot.optimizer.Schema): Database schema
     24    Returns:
     25        sqlglot.Expression: qualified expression
     26    """
     27    schema = ensure_schema(schema)
     28
     29    for scope in traverse_scope(expression):
     30        resolver = Resolver(scope, schema)
     31        _pop_table_column_aliases(scope.ctes)
     32        _pop_table_column_aliases(scope.derived_tables)
     33        _expand_using(scope, resolver)
     34        _qualify_columns(scope, resolver)
     35        if not isinstance(scope.expression, exp.UDTF):
     36            _expand_stars(scope, resolver)
     37            _qualify_outputs(scope)
     38        _expand_group_by(scope, resolver)
     39        _expand_order_by(scope)
     40    return expression
     41
     42
     43def validate_qualify_columns(expression):
     44    """Raise an `OptimizeError` if any columns aren't qualified"""
     45    unqualified_columns = []
     46    for scope in traverse_scope(expression):
     47        if isinstance(scope.expression, exp.Select):
     48            unqualified_columns.extend(scope.unqualified_columns)
     49            if scope.external_columns and not scope.is_correlated_subquery:
     50                column = scope.external_columns[0]
     51                raise OptimizeError(f"Unknown table: '{column.table}' for column '{column}'")
     52
     53    if unqualified_columns:
     54        raise OptimizeError(f"Ambiguous columns: {unqualified_columns}")
     55    return expression
     56
     57
     58def _pop_table_column_aliases(derived_tables):
     59    """
     60    Remove table column aliases.
     61
     62    (e.g. SELECT ... FROM (SELECT ...) AS foo(col1, col2)
     63    """
     64    for derived_table in derived_tables:
     65        table_alias = derived_table.args.get("alias")
     66        if table_alias:
     67            table_alias.args.pop("columns", None)
     68
     69
     70def _expand_using(scope, resolver):
     71    joins = list(scope.expression.find_all(exp.Join))
     72    names = {join.this.alias for join in joins}
     73    ordered = [key for key in scope.selected_sources if key not in names]
     74
     75    # Mapping of automatically joined column names to source names
     76    column_tables = {}
     77
     78    for join in joins:
     79        using = join.args.get("using")
     80
     81        if not using:
     82            continue
     83
     84        join_table = join.this.alias_or_name
     85
     86        columns = {}
     87
     88        for k in scope.selected_sources:
     89            if k in ordered:
     90                for column in resolver.get_source_columns(k):
     91                    if column not in columns:
     92                        columns[column] = k
     93
     94        ordered.append(join_table)
     95        join_columns = resolver.get_source_columns(join_table)
     96        conditions = []
     97
     98        for identifier in using:
     99            identifier = identifier.name
    100            table = columns.get(identifier)
    101
    102            if not table or identifier not in join_columns:
    103                raise OptimizeError(f"Cannot automatically join: {identifier}")
    104
    105            conditions.append(
    106                exp.condition(
    107                    exp.EQ(
    108                        this=exp.column(identifier, table=table),
    109                        expression=exp.column(identifier, table=join_table),
    110                    )
    111                )
    112            )
    113
    114            tables = column_tables.setdefault(identifier, [])
    115            if table not in tables:
    116                tables.append(table)
    117            if join_table not in tables:
    118                tables.append(join_table)
    119
    120        join.args.pop("using")
    121        join.set("on", exp.and_(*conditions))
    122
    123    if column_tables:
    124        for column in scope.columns:
    125            if not column.table and column.name in column_tables:
    126                tables = column_tables[column.name]
    127                coalesce = [exp.column(column.name, table=table) for table in tables]
    128                replacement = exp.Coalesce(this=coalesce[0], expressions=coalesce[1:])
    129
    130                # Ensure selects keep their output name
    131                if isinstance(column.parent, exp.Select):
    132                    replacement = exp.alias_(replacement, alias=column.name)
    133
    134                scope.replace(column, replacement)
    135
    136
    137def _expand_group_by(scope, resolver):
    138    group = scope.expression.args.get("group")
    139    if not group:
    140        return
    141
    142    # Replace references to select aliases
    143    def transform(node, *_):
    144        if isinstance(node, exp.Column) and not node.table:
    145            table = resolver.get_table(node.name)
    146
    147            # Source columns get priority over select aliases
    148            if table:
    149                node.set("table", table)
    150                return node
    151
    152            selects = {s.alias_or_name: s for s in scope.selects}
    153
    154            select = selects.get(node.name)
    155            if select:
    156                scope.clear_cache()
    157                if isinstance(select, exp.Alias):
    158                    select = select.this
    159                return select.copy()
    160
    161        return node
    162
    163    group.transform(transform, copy=False)
    164    group.set("expressions", _expand_positional_references(scope, group.expressions))
    165    scope.expression.set("group", group)
    166
    167
    168def _expand_order_by(scope):
    169    order = scope.expression.args.get("order")
    170    if not order:
    171        return
    172
    173    ordereds = order.expressions
    174    for ordered, new_expression in zip(
    175        ordereds,
    176        _expand_positional_references(scope, (o.this for o in ordereds)),
    177    ):
    178        ordered.set("this", new_expression)
    179
    180
    181def _expand_positional_references(scope, expressions):
    182    new_nodes = []
    183    for node in expressions:
    184        if node.is_int:
    185            try:
    186                select = scope.selects[int(node.name) - 1]
    187            except IndexError:
    188                raise OptimizeError(f"Unknown output column: {node.name}")
    189            if isinstance(select, exp.Alias):
    190                select = select.this
    191            new_nodes.append(select.copy())
    192            scope.clear_cache()
    193        else:
    194            new_nodes.append(node)
    195
    196    return new_nodes
    197
    198
    199def _qualify_columns(scope, resolver):
    200    """Disambiguate columns, ensuring each column specifies a source"""
    201    for column in scope.columns:
    202        column_table = column.table
    203        column_name = column.name
    204
    205        if column_table and column_table in scope.sources:
    206            source_columns = resolver.get_source_columns(column_table)
    207            if source_columns and column_name not in source_columns and "*" not in source_columns:
    208                raise OptimizeError(f"Unknown column: {column_name}")
    209
    210        if not column_table:
    211            column_table = resolver.get_table(column_name)
    212
    213            # column_table can be a '' because bigquery unnest has no table alias
    214            if column_table:
    215                column.set("table", column_table)
    216
    217    columns_missing_from_scope = []
    218    # Determine whether each reference in the order by clause is to a column or an alias.
    219    for ordered in scope.find_all(exp.Ordered):
    220        for column in ordered.find_all(exp.Column):
    221            if (
    222                not column.table
    223                and column.parent is not ordered
    224                and column.name in resolver.all_columns
    225            ):
    226                columns_missing_from_scope.append(column)
    227
    228    # Determine whether each reference in the having clause is to a column or an alias.
    229    for having in scope.find_all(exp.Having):
    230        for column in having.find_all(exp.Column):
    231            if (
    232                not column.table
    233                and column.find_ancestor(exp.AggFunc)
    234                and column.name in resolver.all_columns
    235            ):
    236                columns_missing_from_scope.append(column)
    237
    238    for column in columns_missing_from_scope:
    239        column_table = resolver.get_table(column.name)
    240
    241        if column_table:
    242            column.set("table", column_table)
    243
    244
    245def _expand_stars(scope, resolver):
    246    """Expand stars to lists of column selections"""
    247
    248    new_selections = []
    249    except_columns = {}
    250    replace_columns = {}
    251
    252    for expression in scope.selects:
    253        if isinstance(expression, exp.Star):
    254            tables = list(scope.selected_sources)
    255            _add_except_columns(expression, tables, except_columns)
    256            _add_replace_columns(expression, tables, replace_columns)
    257        elif expression.is_star:
    258            tables = [expression.table]
    259            _add_except_columns(expression.this, tables, except_columns)
    260            _add_replace_columns(expression.this, tables, replace_columns)
    261        else:
    262            new_selections.append(expression)
    263            continue
    264
    265        for table in tables:
    266            if table not in scope.sources:
    267                raise OptimizeError(f"Unknown table: {table}")
    268            columns = resolver.get_source_columns(table, only_visible=True)
    269
    270            if columns and "*" not in columns:
    271                table_id = id(table)
    272                for name in columns:
    273                    if name not in except_columns.get(table_id, set()):
    274                        alias_ = replace_columns.get(table_id, {}).get(name, name)
    275                        column = exp.column(name, table)
    276                        new_selections.append(alias(column, alias_) if alias_ != name else column)
    277            else:
    278                return
    279    scope.expression.set("expressions", new_selections)
    280
    281
    282def _add_except_columns(expression, tables, except_columns):
    283    except_ = expression.args.get("except")
    284
    285    if not except_:
    286        return
    287
    288    columns = {e.name for e in except_}
    289
    290    for table in tables:
    291        except_columns[id(table)] = columns
    292
    293
    294def _add_replace_columns(expression, tables, replace_columns):
    295    replace = expression.args.get("replace")
    296
    297    if not replace:
    298        return
    299
    300    columns = {e.this.name: e.alias for e in replace}
    301
    302    for table in tables:
    303        replace_columns[id(table)] = columns
    304
    305
    306def _qualify_outputs(scope):
    307    """Ensure all output columns are aliased"""
    308    new_selections = []
    309
    310    for i, (selection, aliased_column) in enumerate(
    311        itertools.zip_longest(scope.selects, scope.outer_column_list)
    312    ):
    313        if isinstance(selection, exp.Subquery):
    314            if not selection.output_name:
    315                selection.set("alias", exp.TableAlias(this=exp.to_identifier(f"_col_{i}")))
    316        elif not isinstance(selection, exp.Alias) and not selection.is_star:
    317            alias_ = alias(exp.column(""), alias=selection.output_name or f"_col_{i}")
    318            alias_.set("this", selection)
    319            selection = alias_
    320
    321        if aliased_column:
    322            selection.set("alias", exp.to_identifier(aliased_column))
    323
    324        new_selections.append(selection)
    325
    326    scope.expression.set("expressions", new_selections)
    327
    328
    329class Resolver:
    330    """
    331    Helper for resolving columns.
    332
    333    This is a class so we can lazily load some things and easily share them across functions.
    334    """
    335
    336    def __init__(self, scope, schema):
    337        self.scope = scope
    338        self.schema = schema
    339        self._source_columns = None
    340        self._unambiguous_columns = None
    341        self._all_columns = None
    342
    343    def get_table(self, column_name: str) -> t.Optional[exp.Identifier]:
    344        """
    345        Get the table for a column name.
    346
    347        Args:
    348            column_name: The column name to find the table for.
    349        Returns:
    350            The table name if it can be found/inferred.
    351        """
    352        if self._unambiguous_columns is None:
    353            self._unambiguous_columns = self._get_unambiguous_columns(
    354                self._get_all_source_columns()
    355            )
    356
    357        table_name = self._unambiguous_columns.get(column_name)
    358
    359        if not table_name:
    360            sources_without_schema = tuple(
    361                source
    362                for source, columns in self._get_all_source_columns().items()
    363                if not columns or "*" in columns
    364            )
    365            if len(sources_without_schema) == 1:
    366                table_name = sources_without_schema[0]
    367
    368        if table_name not in self.scope.selected_sources:
    369            return exp.to_identifier(table_name)
    370
    371        node, _ = self.scope.selected_sources.get(table_name)
    372
    373        if isinstance(node, exp.Subqueryable):
    374            while node and node.alias != table_name:
    375                node = node.parent
    376        node_alias = node.args.get("alias")
    377        if node_alias:
    378            return node_alias.this
    379        return exp.to_identifier(table_name)
    380
    381    @property
    382    def all_columns(self):
    383        """All available columns of all sources in this scope"""
    384        if self._all_columns is None:
    385            self._all_columns = {
    386                column for columns in self._get_all_source_columns().values() for column in columns
    387            }
    388        return self._all_columns
    389
    390    def get_source_columns(self, name, only_visible=False):
    391        """Resolve the source columns for a given source `name`"""
    392        if name not in self.scope.sources:
    393            raise OptimizeError(f"Unknown table: {name}")
    394
    395        source = self.scope.sources[name]
    396
    397        # If referencing a table, return the columns from the schema
    398        if isinstance(source, exp.Table):
    399            return self.schema.column_names(source, only_visible)
    400
    401        if isinstance(source, Scope) and isinstance(source.expression, exp.Values):
    402            return source.expression.alias_column_names
    403
    404        # Otherwise, if referencing another scope, return that scope's named selects
    405        return source.expression.named_selects
    406
    407    def _get_all_source_columns(self):
    408        if self._source_columns is None:
    409            self._source_columns = {
    410                k: self.get_source_columns(k)
    411                for k in itertools.chain(self.scope.selected_sources, self.scope.lateral_sources)
    412            }
    413        return self._source_columns
    414
    415    def _get_unambiguous_columns(self, source_columns):
    416        """
    417        Find all the unambiguous columns in sources.
    418
    419        Args:
    420            source_columns (dict): Mapping of names to source columns
    421        Returns:
    422            dict: Mapping of column name to source name
    423        """
    424        if not source_columns:
    425            return {}
    426
    427        source_columns = list(source_columns.items())
    428
    429        first_table, first_columns = source_columns[0]
    430        unambiguous_columns = {col: first_table for col in self._find_unique_columns(first_columns)}
    431        all_columns = set(unambiguous_columns)
    432
    433        for table, columns in source_columns[1:]:
    434            unique = self._find_unique_columns(columns)
    435            ambiguous = set(all_columns).intersection(unique)
    436            all_columns.update(columns)
    437            for column in ambiguous:
    438                unambiguous_columns.pop(column, None)
    439            for column in unique.difference(ambiguous):
    440                unambiguous_columns[column] = table
    441
    442        return unambiguous_columns
    443
    444    @staticmethod
    445    def _find_unique_columns(columns):
    446        """
    447        Find the unique columns in a list of columns.
    448
    449        Example:
    450            >>> sorted(Resolver._find_unique_columns(["a", "b", "b", "c"]))
    451            ['a', 'c']
    452
    453        This is necessary because duplicate column names are ambiguous.
    454        """
    455        counts = {}
    456        for column in columns:
    457            counts[column] = counts.get(column, 0) + 1
    458        return {column for column, count in counts.items() if count == 1}

def qualify_columns(expression, schema): View Source

    11def qualify_columns(expression, schema):
    12    """
    13    Rewrite sqlglot AST to have fully qualified columns.
    14
    15    Example:
    16        >>> import sqlglot
    17        >>> schema = {"tbl": {"col": "INT"}}
    18        >>> expression = sqlglot.parse_one("SELECT col FROM tbl")
    19        >>> qualify_columns(expression, schema).sql()
    20        'SELECT tbl.col AS col FROM tbl'
    21
    22    Args:
    23        expression (sqlglot.Expression): expression to qualify
    24        schema (dict|sqlglot.optimizer.Schema): Database schema
    25    Returns:
    26        sqlglot.Expression: qualified expression
    27    """
    28    schema = ensure_schema(schema)
    29
    30    for scope in traverse_scope(expression):
    31        resolver = Resolver(scope, schema)
    32        _pop_table_column_aliases(scope.ctes)
    33        _pop_table_column_aliases(scope.derived_tables)
    34        _expand_using(scope, resolver)
    35        _qualify_columns(scope, resolver)
    36        if not isinstance(scope.expression, exp.UDTF):
    37            _expand_stars(scope, resolver)
    38            _qualify_outputs(scope)
    39        _expand_group_by(scope, resolver)
    40        _expand_order_by(scope)
    41    return expression

Rewrite sqlglot AST to have fully qualified columns.

###### Example:

>     >>> import sqlglot
>     >>> schema = {"tbl": {"col": "INT"}}
>     >>> expression = sqlglot.parse_one("SELECT col FROM tbl")
>     >>> qualify_columns(expression, schema).sql()
>     'SELECT tbl.col AS col FROM tbl'

###### Arguments:

- expression (sqlglot.Expression): expression to qualify
- schema (dict|sqlglot.optimizer.Schema): Database schema

###### Returns:

> sqlglot.Expression: qualified expression

def validate_qualify_columns(expression): View Source

    44def validate_qualify_columns(expression):
    45    """Raise an `OptimizeError` if any columns aren't qualified"""
    46    unqualified_columns = []
    47    for scope in traverse_scope(expression):
    48        if isinstance(scope.expression, exp.Select):
    49            unqualified_columns.extend(scope.unqualified_columns)
    50            if scope.external_columns and not scope.is_correlated_subquery:
    51                column = scope.external_columns[0]
    52                raise OptimizeError(f"Unknown table: '{column.table}' for column '{column}'")
    53
    54    if unqualified_columns:
    55        raise OptimizeError(f"Ambiguous columns: {unqualified_columns}")
    56    return expression

Raise an `OptimizeError` if any columns aren't qualified

class Resolver: View Source

    330class Resolver:
    331    """
    332    Helper for resolving columns.
    333
    334    This is a class so we can lazily load some things and easily share them across functions.
    335    """
    336
    337    def __init__(self, scope, schema):
    338        self.scope = scope
    339        self.schema = schema
    340        self._source_columns = None
    341        self._unambiguous_columns = None
    342        self._all_columns = None
    343
    344    def get_table(self, column_name: str) -> t.Optional[exp.Identifier]:
    345        """
    346        Get the table for a column name.
    347
    348        Args:
    349            column_name: The column name to find the table for.
    350        Returns:
    351            The table name if it can be found/inferred.
    352        """
    353        if self._unambiguous_columns is None:
    354            self._unambiguous_columns = self._get_unambiguous_columns(
    355                self._get_all_source_columns()
    356            )
    357
    358        table_name = self._unambiguous_columns.get(column_name)
    359
    360        if not table_name:
    361            sources_without_schema = tuple(
    362                source
    363                for source, columns in self._get_all_source_columns().items()
    364                if not columns or "*" in columns
    365            )
    366            if len(sources_without_schema) == 1:
    367                table_name = sources_without_schema[0]
    368
    369        if table_name not in self.scope.selected_sources:
    370            return exp.to_identifier(table_name)
    371
    372        node, _ = self.scope.selected_sources.get(table_name)
    373
    374        if isinstance(node, exp.Subqueryable):
    375            while node and node.alias != table_name:
    376                node = node.parent
    377        node_alias = node.args.get("alias")
    378        if node_alias:
    379            return node_alias.this
    380        return exp.to_identifier(table_name)
    381
    382    @property
    383    def all_columns(self):
    384        """All available columns of all sources in this scope"""
    385        if self._all_columns is None:
    386            self._all_columns = {
    387                column for columns in self._get_all_source_columns().values() for column in columns
    388            }
    389        return self._all_columns
    390
    391    def get_source_columns(self, name, only_visible=False):
    392        """Resolve the source columns for a given source `name`"""
    393        if name not in self.scope.sources:
    394            raise OptimizeError(f"Unknown table: {name}")
    395
    396        source = self.scope.sources[name]
    397
    398        # If referencing a table, return the columns from the schema
    399        if isinstance(source, exp.Table):
    400            return self.schema.column_names(source, only_visible)
    401
    402        if isinstance(source, Scope) and isinstance(source.expression, exp.Values):
    403            return source.expression.alias_column_names
    404
    405        # Otherwise, if referencing another scope, return that scope's named selects
    406        return source.expression.named_selects
    407
    408    def _get_all_source_columns(self):
    409        if self._source_columns is None:
    410            self._source_columns = {
    411                k: self.get_source_columns(k)
    412                for k in itertools.chain(self.scope.selected_sources, self.scope.lateral_sources)
    413            }
    414        return self._source_columns
    415
    416    def _get_unambiguous_columns(self, source_columns):
    417        """
    418        Find all the unambiguous columns in sources.
    419
    420        Args:
    421            source_columns (dict): Mapping of names to source columns
    422        Returns:
    423            dict: Mapping of column name to source name
    424        """
    425        if not source_columns:
    426            return {}
    427
    428        source_columns = list(source_columns.items())
    429
    430        first_table, first_columns = source_columns[0]
    431        unambiguous_columns = {col: first_table for col in self._find_unique_columns(first_columns)}
    432        all_columns = set(unambiguous_columns)
    433
    434        for table, columns in source_columns[1:]:
    435            unique = self._find_unique_columns(columns)
    436            ambiguous = set(all_columns).intersection(unique)
    437            all_columns.update(columns)
    438            for column in ambiguous:
    439                unambiguous_columns.pop(column, None)
    440            for column in unique.difference(ambiguous):
    441                unambiguous_columns[column] = table
    442
    443        return unambiguous_columns
    444
    445    @staticmethod
    446    def _find_unique_columns(columns):
    447        """
    448        Find the unique columns in a list of columns.
    449
    450        Example:
    451            >>> sorted(Resolver._find_unique_columns(["a", "b", "b", "c"]))
    452            ['a', 'c']
    453
    454        This is necessary because duplicate column names are ambiguous.
    455        """
    456        counts = {}
    457        for column in columns:
    458            counts[column] = counts.get(column, 0) + 1
    459        return {column for column, count in counts.items() if count == 1}

Helper for resolving columns.

This is a class so we can lazily load some things and easily share them across functions.

Resolver(scope, schema) View Source

    337    def __init__(self, scope, schema):
    338        self.scope = scope
    339        self.schema = schema
    340        self._source_columns = None
    341        self._unambiguous_columns = None
    342        self._all_columns = None

def get_table(self, column_name: str) -> Optional[sqlglot.expressions.Identifier]: View Source

    344    def get_table(self, column_name: str) -> t.Optional[exp.Identifier]:
    345        """
    346        Get the table for a column name.
    347
    348        Args:
    349            column_name: The column name to find the table for.
    350        Returns:
    351            The table name if it can be found/inferred.
    352        """
    353        if self._unambiguous_columns is None:
    354            self._unambiguous_columns = self._get_unambiguous_columns(
    355                self._get_all_source_columns()
    356            )
    357
    358        table_name = self._unambiguous_columns.get(column_name)
    359
    360        if not table_name:
    361            sources_without_schema = tuple(
    362                source
    363                for source, columns in self._get_all_source_columns().items()
    364                if not columns or "*" in columns
    365            )
    366            if len(sources_without_schema) == 1:
    367                table_name = sources_without_schema[0]
    368
    369        if table_name not in self.scope.selected_sources:
    370            return exp.to_identifier(table_name)
    371
    372        node, _ = self.scope.selected_sources.get(table_name)
    373
    374        if isinstance(node, exp.Subqueryable):
    375            while node and node.alias != table_name:
    376                node = node.parent
    377        node_alias = node.args.get("alias")
    378        if node_alias:
    379            return node_alias.this
    380        return exp.to_identifier(table_name)

Get the table for a column name.

###### Arguments:

- column_name: The column name to find the table for.

###### Returns:

> The table name if it can be found/inferred.

all_columns

All available columns of all sources in this scope

def get_source_columns(self, name, only_visible=False): View Source

    391    def get_source_columns(self, name, only_visible=False):
    392        """Resolve the source columns for a given source `name`"""
    393        if name not in self.scope.sources:
    394            raise OptimizeError(f"Unknown table: {name}")
    395
    396        source = self.scope.sources[name]
    397
    398        # If referencing a table, return the columns from the schema
    399        if isinstance(source, exp.Table):
    400            return self.schema.column_names(source, only_visible)
    401
    402        if isinstance(source, Scope) and isinstance(source.expression, exp.Values):
    403            return source.expression.alias_column_names
    404
    405        # Otherwise, if referencing another scope, return that scope's named selects
    406        return source.expression.named_selects

Resolve the source columns for a given source `name`