sqlglot.optimizer.scope

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