sqlglot.optimizer.qualify_columns

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