sqlglot.optimizer.eliminate_ctes

 1from sqlglot.optimizer.scope import Scope, build_scope
 2
 3
 4def eliminate_ctes(expression):
 5    """
 6    Remove unused CTEs from an expression.
 7
 8    Example:
 9        >>> import sqlglot
10        >>> sql = "WITH y AS (SELECT a FROM x) SELECT a FROM z"
11        >>> expression = sqlglot.parse_one(sql)
12        >>> eliminate_ctes(expression).sql()
13        'SELECT a FROM z'
14
15    Args:
16        expression (sqlglot.Expression): expression to optimize
17    Returns:
18        sqlglot.Expression: optimized expression
19    """
20    root = build_scope(expression)
21
22    if root:
23        ref_count = root.ref_count()
24
25        # Traverse the scope tree in reverse so we can remove chains of unused CTEs
26        for scope in reversed(list(root.traverse())):
27            if scope.is_cte:
28                count = ref_count[id(scope)]
29                if count <= 0:
30                    cte_node = scope.expression.parent
31                    with_node = cte_node.parent
32                    cte_node.pop()
33
34                    # Pop the entire WITH clause if this is the last CTE
35                    if with_node and len(with_node.expressions) <= 0:
36                        with_node.pop()
37
38                    # Decrement the ref count for all sources this CTE selects from
39                    for _, source in scope.selected_sources.values():
40                        if isinstance(source, Scope):
41                            ref_count[id(source)] -= 1
42
43    return expression