sqlglot.optimizer.scope

  1import itertools
  2import logging
  3import typing as t
  4from collections import defaultdict
  5from enum import Enum, auto
  6
  7from sqlglot import exp
  8from sqlglot.errors import OptimizeError
  9from sqlglot.helper import find_new_name
 10
 11logger = logging.getLogger("sqlglot")
 12
 13
 14class ScopeType(Enum):
 15    ROOT = auto()
 16    SUBQUERY = auto()
 17    DERIVED_TABLE = auto()
 18    CTE = auto()
 19    UNION = auto()
 20    UDTF = auto()
 21
 22
 23class Scope:
 24    """
 25    Selection scope.
 26
 27    Attributes:
 28        expression (exp.Select|exp.Union): Root expression of this scope
 29        sources (dict[str, exp.Table|Scope]): Mapping of source name to either
 30            a Table expression or another Scope instance. For example:
 31                SELECT * FROM x                     {"x": Table(this="x")}
 32                SELECT * FROM x AS y                {"y": Table(this="x")}
 33                SELECT * FROM (SELECT ...) AS y     {"y": Scope(...)}
 34        lateral_sources (dict[str, exp.Table|Scope]): Sources from laterals
 35            For example:
 36                SELECT c FROM x LATERAL VIEW EXPLODE (a) AS c;
 37            The LATERAL VIEW EXPLODE gets x as a source.
 38        outer_column_list (list[str]): If this is a derived table or CTE, and the outer query
 39            defines a column list of it's alias of this scope, this is that list of columns.
 40            For example:
 41                SELECT * FROM (SELECT ...) AS y(col1, col2)
 42            The inner query would have `["col1", "col2"]` for its `outer_column_list`
 43        parent (Scope): Parent scope
 44        scope_type (ScopeType): Type of this scope, relative to it's parent
 45        subquery_scopes (list[Scope]): List of all child scopes for subqueries
 46        cte_scopes (list[Scope]): List of all child scopes for CTEs
 47        derived_table_scopes (list[Scope]): List of all child scopes for derived_tables
 48        udtf_scopes (list[Scope]): List of all child scopes for user defined tabular functions
 49        table_scopes (list[Scope]): derived_table_scopes + udtf_scopes, in the order that they're defined
 50        union_scopes (list[Scope, Scope]): If this Scope is for a Union expression, this will be
 51            a list of the left and right child scopes.
 52    """
 53
 54    def __init__(
 55        self,
 56        expression,
 57        sources=None,
 58        outer_column_list=None,
 59        parent=None,
 60        scope_type=ScopeType.ROOT,
 61        lateral_sources=None,
 62    ):
 63        self.expression = expression
 64        self.sources = sources or {}
 65        self.lateral_sources = lateral_sources.copy() if lateral_sources else {}
 66        self.sources.update(self.lateral_sources)
 67        self.outer_column_list = outer_column_list or []
 68        self.parent = parent
 69        self.scope_type = scope_type
 70        self.subquery_scopes = []
 71        self.derived_table_scopes = []
 72        self.table_scopes = []
 73        self.cte_scopes = []
 74        self.union_scopes = []
 75        self.udtf_scopes = []
 76        self.clear_cache()
 77
 78    def clear_cache(self):
 79        self._collected = False
 80        self._raw_columns = None
 81        self._derived_tables = None
 82        self._udtfs = None
 83        self._tables = None
 84        self._ctes = None
 85        self._subqueries = None
 86        self._selected_sources = None
 87        self._columns = None
 88        self._external_columns = None
 89        self._join_hints = None
 90        self._pivots = None
 91        self._references = None
 92
 93    def branch(self, expression, scope_type, chain_sources=None, **kwargs):
 94        """Branch from the current scope to a new, inner scope"""
 95        return Scope(
 96            expression=expression.unnest(),
 97            sources={**self.cte_sources, **(chain_sources or {})},
 98            parent=self,
 99            scope_type=scope_type,
100            **kwargs,
101        )
102
103    def _collect(self):
104        self._tables = []
105        self._ctes = []
106        self._subqueries = []
107        self._derived_tables = []
108        self._udtfs = []
109        self._raw_columns = []
110        self._join_hints = []
111
112        for node, parent, _ in self.walk(bfs=False):
113            if node is self.expression:
114                continue
115            elif isinstance(node, exp.Column) and not isinstance(node.this, exp.Star):
116                self._raw_columns.append(node)
117            elif isinstance(node, exp.Table) and not isinstance(node.parent, exp.JoinHint):
118                self._tables.append(node)
119            elif isinstance(node, exp.JoinHint):
120                self._join_hints.append(node)
121            elif isinstance(node, exp.UDTF):
122                self._udtfs.append(node)
123            elif isinstance(node, exp.CTE):
124                self._ctes.append(node)
125            elif (
126                isinstance(node, exp.Subquery)
127                and isinstance(parent, (exp.From, exp.Join, exp.Subquery))
128                and _is_derived_table(node)
129            ):
130                self._derived_tables.append(node)
131            elif isinstance(node, exp.Subqueryable):
132                self._subqueries.append(node)
133
134        self._collected = True
135
136    def _ensure_collected(self):
137        if not self._collected:
138            self._collect()
139
140    def walk(self, bfs=True):
141        return walk_in_scope(self.expression, bfs=bfs)
142
143    def find(self, *expression_types, bfs=True):
144        """
145        Returns the first node in this scope which matches at least one of the specified types.
146
147        This does NOT traverse into subscopes.
148
149        Args:
150            expression_types (type): the expression type(s) to match.
151            bfs (bool): True to use breadth-first search, False to use depth-first.
152
153        Returns:
154            exp.Expression: the node which matches the criteria or None if no node matching
155            the criteria was found.
156        """
157        return next(self.find_all(*expression_types, bfs=bfs), None)
158
159    def find_all(self, *expression_types, bfs=True):
160        """
161        Returns a generator object which visits all nodes in this scope and only yields those that
162        match at least one of the specified expression types.
163
164        This does NOT traverse into subscopes.
165
166        Args:
167            expression_types (type): the expression type(s) to match.
168            bfs (bool): True to use breadth-first search, False to use depth-first.
169
170        Yields:
171            exp.Expression: nodes
172        """
173        for expression, *_ in self.walk(bfs=bfs):
174            if isinstance(expression, expression_types):
175                yield expression
176
177    def replace(self, old, new):
178        """
179        Replace `old` with `new`.
180
181        This can be used instead of `exp.Expression.replace` to ensure the `Scope` is kept up-to-date.
182
183        Args:
184            old (exp.Expression): old node
185            new (exp.Expression): new node
186        """
187        old.replace(new)
188        self.clear_cache()
189
190    @property
191    def tables(self):
192        """
193        List of tables in this scope.
194
195        Returns:
196            list[exp.Table]: tables
197        """
198        self._ensure_collected()
199        return self._tables
200
201    @property
202    def ctes(self):
203        """
204        List of CTEs in this scope.
205
206        Returns:
207            list[exp.CTE]: ctes
208        """
209        self._ensure_collected()
210        return self._ctes
211
212    @property
213    def derived_tables(self):
214        """
215        List of derived tables in this scope.
216
217        For example:
218            SELECT * FROM (SELECT ...) <- that's a derived table
219
220        Returns:
221            list[exp.Subquery]: derived tables
222        """
223        self._ensure_collected()
224        return self._derived_tables
225
226    @property
227    def udtfs(self):
228        """
229        List of "User Defined Tabular Functions" in this scope.
230
231        Returns:
232            list[exp.UDTF]: UDTFs
233        """
234        self._ensure_collected()
235        return self._udtfs
236
237    @property
238    def subqueries(self):
239        """
240        List of subqueries in this scope.
241
242        For example:
243            SELECT * FROM x WHERE a IN (SELECT ...) <- that's a subquery
244
245        Returns:
246            list[exp.Subqueryable]: subqueries
247        """
248        self._ensure_collected()
249        return self._subqueries
250
251    @property
252    def columns(self):
253        """
254        List of columns in this scope.
255
256        Returns:
257            list[exp.Column]: Column instances in this scope, plus any
258                Columns that reference this scope from correlated subqueries.
259        """
260        if self._columns is None:
261            self._ensure_collected()
262            columns = self._raw_columns
263
264            external_columns = [
265                column
266                for scope in itertools.chain(self.subquery_scopes, self.udtf_scopes)
267                for column in scope.external_columns
268            ]
269
270            named_selects = set(self.expression.named_selects)
271
272            self._columns = []
273            for column in columns + external_columns:
274                ancestor = column.find_ancestor(
275                    exp.Select, exp.Qualify, exp.Order, exp.Having, exp.Hint, exp.Table
276                )
277                if (
278                    not ancestor
279                    or column.table
280                    or isinstance(ancestor, exp.Select)
281                    or (isinstance(ancestor, exp.Table) and not isinstance(ancestor.this, exp.Func))
282                    or (
283                        isinstance(ancestor, exp.Order)
284                        and (
285                            isinstance(ancestor.parent, exp.Window)
286                            or column.name not in named_selects
287                        )
288                    )
289                ):
290                    self._columns.append(column)
291
292        return self._columns
293
294    @property
295    def selected_sources(self):
296        """
297        Mapping of nodes and sources that are actually selected from in this scope.
298
299        That is, all tables in a schema are selectable at any point. But a
300        table only becomes a selected source if it's included in a FROM or JOIN clause.
301
302        Returns:
303            dict[str, (exp.Table|exp.Select, exp.Table|Scope)]: selected sources and nodes
304        """
305        if self._selected_sources is None:
306            result = {}
307
308            for name, node in self.references:
309                if name in result:
310                    raise OptimizeError(f"Alias already used: {name}")
311                if name in self.sources:
312                    result[name] = (node, self.sources[name])
313
314            self._selected_sources = result
315        return self._selected_sources
316
317    @property
318    def references(self) -> t.List[t.Tuple[str, exp.Expression]]:
319        if self._references is None:
320            self._references = []
321
322            for table in self.tables:
323                self._references.append((table.alias_or_name, table))
324            for expression in itertools.chain(self.derived_tables, self.udtfs):
325                self._references.append(
326                    (
327                        expression.alias,
328                        expression if expression.args.get("pivots") else expression.unnest(),
329                    )
330                )
331
332        return self._references
333
334    @property
335    def cte_sources(self):
336        """
337        Sources that are CTEs.
338
339        Returns:
340            dict[str, Scope]: Mapping of source alias to Scope
341        """
342        return {
343            alias: scope
344            for alias, scope in self.sources.items()
345            if isinstance(scope, Scope) and scope.is_cte
346        }
347
348    @property
349    def external_columns(self):
350        """
351        Columns that appear to reference sources in outer scopes.
352
353        Returns:
354            list[exp.Column]: Column instances that don't reference
355                sources in the current scope.
356        """
357        if self._external_columns is None:
358            self._external_columns = [
359                c for c in self.columns if c.table not in self.selected_sources
360            ]
361        return self._external_columns
362
363    @property
364    def unqualified_columns(self):
365        """
366        Unqualified columns in the current scope.
367
368        Returns:
369             list[exp.Column]: Unqualified columns
370        """
371        return [c for c in self.columns if not c.table]
372
373    @property
374    def join_hints(self):
375        """
376        Hints that exist in the scope that reference tables
377
378        Returns:
379            list[exp.JoinHint]: Join hints that are referenced within the scope
380        """
381        if self._join_hints is None:
382            return []
383        return self._join_hints
384
385    @property
386    def pivots(self):
387        if not self._pivots:
388            self._pivots = [
389                pivot for _, node in self.references for pivot in node.args.get("pivots") or []
390            ]
391
392        return self._pivots
393
394    def source_columns(self, source_name):
395        """
396        Get all columns in the current scope for a particular source.
397
398        Args:
399            source_name (str): Name of the source
400        Returns:
401            list[exp.Column]: Column instances that reference `source_name`
402        """
403        return [column for column in self.columns if column.table == source_name]
404
405    @property
406    def is_subquery(self):
407        """Determine if this scope is a subquery"""
408        return self.scope_type == ScopeType.SUBQUERY
409
410    @property
411    def is_derived_table(self):
412        """Determine if this scope is a derived table"""
413        return self.scope_type == ScopeType.DERIVED_TABLE
414
415    @property
416    def is_union(self):
417        """Determine if this scope is a union"""
418        return self.scope_type == ScopeType.UNION
419
420    @property
421    def is_cte(self):
422        """Determine if this scope is a common table expression"""
423        return self.scope_type == ScopeType.CTE
424
425    @property
426    def is_root(self):
427        """Determine if this is the root scope"""
428        return self.scope_type == ScopeType.ROOT
429
430    @property
431    def is_udtf(self):
432        """Determine if this scope is a UDTF (User Defined Table Function)"""
433        return self.scope_type == ScopeType.UDTF
434
435    @property
436    def is_correlated_subquery(self):
437        """Determine if this scope is a correlated subquery"""
438        return bool(
439            (self.is_subquery or (self.parent and isinstance(self.parent.expression, exp.Lateral)))
440            and self.external_columns
441        )
442
443    def rename_source(self, old_name, new_name):
444        """Rename a source in this scope"""
445        columns = self.sources.pop(old_name or "", [])
446        self.sources[new_name] = columns
447
448    def add_source(self, name, source):
449        """Add a source to this scope"""
450        self.sources[name] = source
451        self.clear_cache()
452
453    def remove_source(self, name):
454        """Remove a source from this scope"""
455        self.sources.pop(name, None)
456        self.clear_cache()
457
458    def __repr__(self):
459        return f"Scope<{self.expression.sql()}>"
460
461    def traverse(self):
462        """
463        Traverse the scope tree from this node.
464
465        Yields:
466            Scope: scope instances in depth-first-search post-order
467        """
468        for child_scope in itertools.chain(
469            self.cte_scopes, self.union_scopes, self.table_scopes, self.subquery_scopes
470        ):
471            yield from child_scope.traverse()
472        yield self
473
474    def ref_count(self):
475        """
476        Count the number of times each scope in this tree is referenced.
477
478        Returns:
479            dict[int, int]: Mapping of Scope instance ID to reference count
480        """
481        scope_ref_count = defaultdict(lambda: 0)
482
483        for scope in self.traverse():
484            for _, source in scope.selected_sources.values():
485                scope_ref_count[id(source)] += 1
486
487        return scope_ref_count
488
489
490def traverse_scope(expression: exp.Expression) -> t.List[Scope]:
491    """
492    Traverse an expression by its "scopes".
493
494    "Scope" represents the current context of a Select statement.
495
496    This is helpful for optimizing queries, where we need more information than
497    the expression tree itself. For example, we might care about the source
498    names within a subquery. Returns a list because a generator could result in
499    incomplete properties which is confusing.
500
501    Examples:
502        >>> import sqlglot
503        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT a FROM x) AS y")
504        >>> scopes = traverse_scope(expression)
505        >>> scopes[0].expression.sql(), list(scopes[0].sources)
506        ('SELECT a FROM x', ['x'])
507        >>> scopes[1].expression.sql(), list(scopes[1].sources)
508        ('SELECT a FROM (SELECT a FROM x) AS y', ['y'])
509
510    Args:
511        expression (exp.Expression): expression to traverse
512    Returns:
513        list[Scope]: scope instances
514    """
515    if isinstance(expression, exp.Unionable) or (
516        isinstance(expression, exp.DDL) and isinstance(expression.expression, exp.Subqueryable)
517    ):
518        return list(_traverse_scope(Scope(expression)))
519
520    return []
521
522
523def build_scope(expression: exp.Expression) -> t.Optional[Scope]:
524    """
525    Build a scope tree.
526
527    Args:
528        expression (exp.Expression): expression to build the scope tree for
529    Returns:
530        Scope: root scope
531    """
532    scopes = traverse_scope(expression)
533    if scopes:
534        return scopes[-1]
535    return None
536
537
538def _traverse_scope(scope):
539    if isinstance(scope.expression, exp.Select):
540        yield from _traverse_select(scope)
541    elif isinstance(scope.expression, exp.Union):
542        yield from _traverse_union(scope)
543    elif isinstance(scope.expression, exp.Subquery):
544        yield from _traverse_subqueries(scope)
545    elif isinstance(scope.expression, exp.Table):
546        yield from _traverse_tables(scope)
547    elif isinstance(scope.expression, exp.UDTF):
548        yield from _traverse_udtfs(scope)
549    elif isinstance(scope.expression, exp.DDL):
550        yield from _traverse_ddl(scope)
551    else:
552        logger.warning(
553            "Cannot traverse scope %s with type '%s'", scope.expression, type(scope.expression)
554        )
555        return
556
557    yield scope
558
559
560def _traverse_select(scope):
561    yield from _traverse_ctes(scope)
562    yield from _traverse_tables(scope)
563    yield from _traverse_subqueries(scope)
564
565
566def _traverse_union(scope):
567    yield from _traverse_ctes(scope)
568
569    # The last scope to be yield should be the top most scope
570    left = None
571    for left in _traverse_scope(scope.branch(scope.expression.left, scope_type=ScopeType.UNION)):
572        yield left
573
574    right = None
575    for right in _traverse_scope(scope.branch(scope.expression.right, scope_type=ScopeType.UNION)):
576        yield right
577
578    scope.union_scopes = [left, right]
579
580
581def _traverse_ctes(scope):
582    sources = {}
583
584    for cte in scope.ctes:
585        recursive_scope = None
586
587        # if the scope is a recursive cte, it must be in the form of base_case UNION recursive.
588        # thus the recursive scope is the first section of the union.
589        with_ = scope.expression.args.get("with")
590        if with_ and with_.recursive:
591            union = cte.this
592
593            if isinstance(union, exp.Union):
594                recursive_scope = scope.branch(union.this, scope_type=ScopeType.CTE)
595
596        child_scope = None
597
598        for child_scope in _traverse_scope(
599            scope.branch(
600                cte.this,
601                chain_sources=sources,
602                outer_column_list=cte.alias_column_names,
603                scope_type=ScopeType.CTE,
604            )
605        ):
606            yield child_scope
607
608            alias = cte.alias
609            sources[alias] = child_scope
610
611            if recursive_scope:
612                child_scope.add_source(alias, recursive_scope)
613
614        # append the final child_scope yielded
615        if child_scope:
616            scope.cte_scopes.append(child_scope)
617
618    scope.sources.update(sources)
619
620
621def _is_derived_table(expression: exp.Subquery) -> bool:
622    """
623    We represent (tbl1 JOIN tbl2) as a Subquery, but it's not really a "derived table",
624    as it doesn't introduce a new scope. If an alias is present, it shadows all names
625    under the Subquery, so that's one exception to this rule.
626    """
627    return bool(expression.alias or isinstance(expression.this, exp.Subqueryable))
628
629
630def _traverse_tables(scope):
631    sources = {}
632
633    # Traverse FROMs, JOINs, and LATERALs in the order they are defined
634    expressions = []
635    from_ = scope.expression.args.get("from")
636    if from_:
637        expressions.append(from_.this)
638
639    for join in scope.expression.args.get("joins") or []:
640        expressions.append(join.this)
641
642    if isinstance(scope.expression, exp.Table):
643        expressions.append(scope.expression)
644
645    expressions.extend(scope.expression.args.get("laterals") or [])
646
647    for expression in expressions:
648        if isinstance(expression, exp.Table):
649            table_name = expression.name
650            source_name = expression.alias_or_name
651
652            if table_name in scope.sources and not expression.db:
653                # This is a reference to a parent source (e.g. a CTE), not an actual table, unless
654                # it is pivoted, because then we get back a new table and hence a new source.
655                pivots = expression.args.get("pivots")
656                if pivots:
657                    sources[pivots[0].alias] = expression
658                else:
659                    sources[source_name] = scope.sources[table_name]
660            elif source_name in sources:
661                sources[find_new_name(sources, table_name)] = expression
662            else:
663                sources[source_name] = expression
664
665            # Make sure to not include the joins twice
666            if expression is not scope.expression:
667                expressions.extend(join.this for join in expression.args.get("joins") or [])
668
669            continue
670
671        if not isinstance(expression, exp.DerivedTable):
672            continue
673
674        if isinstance(expression, exp.UDTF):
675            lateral_sources = sources
676            scope_type = ScopeType.UDTF
677            scopes = scope.udtf_scopes
678        elif _is_derived_table(expression):
679            lateral_sources = None
680            scope_type = ScopeType.DERIVED_TABLE
681            scopes = scope.derived_table_scopes
682            expressions.extend(join.this for join in expression.args.get("joins") or [])
683        else:
684            # Makes sure we check for possible sources in nested table constructs
685            expressions.append(expression.this)
686            expressions.extend(join.this for join in expression.args.get("joins") or [])
687            continue
688
689        for child_scope in _traverse_scope(
690            scope.branch(
691                expression,
692                lateral_sources=lateral_sources,
693                outer_column_list=expression.alias_column_names,
694                scope_type=scope_type,
695            )
696        ):
697            yield child_scope
698
699            # Tables without aliases will be set as ""
700            # This shouldn't be a problem once qualify_columns runs, as it adds aliases on everything.
701            # Until then, this means that only a single, unaliased derived table is allowed (rather,
702            # the latest one wins.
703            sources[expression.alias] = child_scope
704
705        # append the final child_scope yielded
706        scopes.append(child_scope)
707        scope.table_scopes.append(child_scope)
708
709    scope.sources.update(sources)
710
711
712def _traverse_subqueries(scope):
713    for subquery in scope.subqueries:
714        top = None
715        for child_scope in _traverse_scope(scope.branch(subquery, scope_type=ScopeType.SUBQUERY)):
716            yield child_scope
717            top = child_scope
718        scope.subquery_scopes.append(top)
719
720
721def _traverse_udtfs(scope):
722    if isinstance(scope.expression, exp.Unnest):
723        expressions = scope.expression.expressions
724    elif isinstance(scope.expression, exp.Lateral):
725        expressions = [scope.expression.this]
726    else:
727        expressions = []
728
729    sources = {}
730    for expression in expressions:
731        if isinstance(expression, exp.Subquery) and _is_derived_table(expression):
732            top = None
733            for child_scope in _traverse_scope(
734                scope.branch(
735                    expression,
736                    scope_type=ScopeType.DERIVED_TABLE,
737                    outer_column_list=expression.alias_column_names,
738                )
739            ):
740                yield child_scope
741                top = child_scope
742                sources[expression.alias] = child_scope
743
744            scope.derived_table_scopes.append(top)
745            scope.table_scopes.append(top)
746
747    scope.sources.update(sources)
748
749
750def _traverse_ddl(scope):
751    yield from _traverse_ctes(scope)
752
753    query_scope = scope.branch(
754        scope.expression.expression, scope_type=ScopeType.DERIVED_TABLE, chain_sources=scope.sources
755    )
756    query_scope._collect()
757    query_scope._ctes = scope.ctes + query_scope._ctes
758
759    yield from _traverse_scope(query_scope)
760
761
762def walk_in_scope(expression, bfs=True):
763    """
764    Returns a generator object which visits all nodes in the syntrax tree, stopping at
765    nodes that start child scopes.
766
767    Args:
768        expression (exp.Expression):
769        bfs (bool): if set to True the BFS traversal order will be applied,
770            otherwise the DFS traversal will be used instead.
771
772    Yields:
773        tuple[exp.Expression, Optional[exp.Expression], str]: node, parent, arg key
774    """
775    # We'll use this variable to pass state into the dfs generator.
776    # Whenever we set it to True, we exclude a subtree from traversal.
777    prune = False
778
779    for node, parent, key in expression.walk(bfs=bfs, prune=lambda *_: prune):
780        prune = False
781
782        yield node, parent, key
783
784        if node is expression:
785            continue
786        if (
787            isinstance(node, exp.CTE)
788            or (
789                isinstance(node, exp.Subquery)
790                and isinstance(parent, (exp.From, exp.Join, exp.Subquery))
791                and _is_derived_table(node)
792            )
793            or isinstance(node, exp.UDTF)
794            or isinstance(node, exp.Subqueryable)
795        ):
796            prune = True
797
798            if isinstance(node, (exp.Subquery, exp.UDTF)):
799                # The following args are not actually in the inner scope, so we should visit them
800                for key in ("joins", "laterals", "pivots"):
801                    for arg in node.args.get(key) or []:
802                        yield from walk_in_scope(arg, bfs=bfs)