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