sqlglot.transforms

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import expressions as exp
  6from sqlglot.helper import find_new_name, name_sequence
  7
  8if t.TYPE_CHECKING:
  9    from sqlglot.generator import Generator
 10
 11
 12def preprocess(
 13    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
 14) -> t.Callable[[Generator, exp.Expression], str]:
 15    """
 16    Creates a new transform by chaining a sequence of transformations and converts the resulting
 17    expression to SQL, using either the "_sql" method corresponding to the resulting expression,
 18    or the appropriate `Generator.TRANSFORMS` function (when applicable -- see below).
 19
 20    Args:
 21        transforms: sequence of transform functions. These will be called in order.
 22
 23    Returns:
 24        Function that can be used as a generator transform.
 25    """
 26
 27    def _to_sql(self, expression: exp.Expression) -> str:
 28        expression_type = type(expression)
 29
 30        expression = transforms[0](expression)
 31        for transform in transforms[1:]:
 32            expression = transform(expression)
 33
 34        _sql_handler = getattr(self, expression.key + "_sql", None)
 35        if _sql_handler:
 36            return _sql_handler(expression)
 37
 38        transforms_handler = self.TRANSFORMS.get(type(expression))
 39        if transforms_handler:
 40            if expression_type is type(expression):
 41                if isinstance(expression, exp.Func):
 42                    return self.function_fallback_sql(expression)
 43
 44                # Ensures we don't enter an infinite loop. This can happen when the original expression
 45                # has the same type as the final expression and there's no _sql method available for it,
 46                # because then it'd re-enter _to_sql.
 47                raise ValueError(
 48                    f"Expression type {expression.__class__.__name__} requires a _sql method in order to be transformed."
 49                )
 50
 51            return transforms_handler(self, expression)
 52
 53        raise ValueError(f"Unsupported expression type {expression.__class__.__name__}.")
 54
 55    return _to_sql
 56
 57
 58def unalias_group(expression: exp.Expression) -> exp.Expression:
 59    """
 60    Replace references to select aliases in GROUP BY clauses.
 61
 62    Example:
 63        >>> import sqlglot
 64        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
 65        'SELECT a AS b FROM x GROUP BY 1'
 66
 67    Args:
 68        expression: the expression that will be transformed.
 69
 70    Returns:
 71        The transformed expression.
 72    """
 73    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
 74        aliased_selects = {
 75            e.alias: i
 76            for i, e in enumerate(expression.parent.expressions, start=1)
 77            if isinstance(e, exp.Alias)
 78        }
 79
 80        for group_by in expression.expressions:
 81            if (
 82                isinstance(group_by, exp.Column)
 83                and not group_by.table
 84                and group_by.name in aliased_selects
 85            ):
 86                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
 87
 88    return expression
 89
 90
 91def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
 92    """
 93    Convert SELECT DISTINCT ON statements to a subquery with a window function.
 94
 95    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
 96
 97    Args:
 98        expression: the expression that will be transformed.
 99
100    Returns:
101        The transformed expression.
102    """
103    if (
104        isinstance(expression, exp.Select)
105        and expression.args.get("distinct")
106        and expression.args["distinct"].args.get("on")
107        and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
108    ):
109        distinct_cols = expression.args["distinct"].pop().args["on"].expressions
110        outer_selects = expression.selects
111        row_number = find_new_name(expression.named_selects, "_row_number")
112        window = exp.Window(this=exp.RowNumber(), partition_by=distinct_cols)
113        order = expression.args.get("order")
114
115        if order:
116            window.set("order", order.pop())
117        else:
118            window.set("order", exp.Order(expressions=[c.copy() for c in distinct_cols]))
119
120        window = exp.alias_(window, row_number)
121        expression.select(window, copy=False)
122
123        return (
124            exp.select(*outer_selects, copy=False)
125            .from_(expression.subquery("_t", copy=False), copy=False)
126            .where(exp.column(row_number).eq(1), copy=False)
127        )
128
129    return expression
130
131
132def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
133    """
134    Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
135
136    The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
137    https://docs.snowflake.com/en/sql-reference/constructs/qualify
138
139    Some dialects don't support window functions in the WHERE clause, so we need to include them as
140    projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
141    if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
142    otherwise we won't be able to refer to it in the outer query's WHERE clause. Finally, if a
143    newly aliased projection is referenced in the QUALIFY clause, it will be replaced by the
144    corresponding expression to avoid creating invalid column references.
145    """
146    if isinstance(expression, exp.Select) and expression.args.get("qualify"):
147        taken = set(expression.named_selects)
148        for select in expression.selects:
149            if not select.alias_or_name:
150                alias = find_new_name(taken, "_c")
151                select.replace(exp.alias_(select, alias))
152                taken.add(alias)
153
154        def _select_alias_or_name(select: exp.Expression) -> str | exp.Column:
155            alias_or_name = select.alias_or_name
156            identifier = select.args.get("alias") or select.this
157            if isinstance(identifier, exp.Identifier):
158                return exp.column(alias_or_name, quoted=identifier.args.get("quoted"))
159            return alias_or_name
160
161        outer_selects = exp.select(*list(map(_select_alias_or_name, expression.selects)))
162        qualify_filters = expression.args["qualify"].pop().this
163        expression_by_alias = {
164            select.alias: select.this
165            for select in expression.selects
166            if isinstance(select, exp.Alias)
167        }
168
169        select_candidates = exp.Window if expression.is_star else (exp.Window, exp.Column)
170        for select_candidate in qualify_filters.find_all(select_candidates):
171            if isinstance(select_candidate, exp.Window):
172                if expression_by_alias:
173                    for column in select_candidate.find_all(exp.Column):
174                        expr = expression_by_alias.get(column.name)
175                        if expr:
176                            column.replace(expr)
177
178                alias = find_new_name(expression.named_selects, "_w")
179                expression.select(exp.alias_(select_candidate, alias), copy=False)
180                column = exp.column(alias)
181
182                if isinstance(select_candidate.parent, exp.Qualify):
183                    qualify_filters = column
184                else:
185                    select_candidate.replace(column)
186            elif select_candidate.name not in expression.named_selects:
187                expression.select(select_candidate.copy(), copy=False)
188
189        return outer_selects.from_(expression.subquery(alias="_t", copy=False), copy=False).where(
190            qualify_filters, copy=False
191        )
192
193    return expression
194
195
196def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
197    """
198    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in
199    other expressions. This transforms removes the precision from parameterized types in expressions.
200    """
201    for node in expression.find_all(exp.DataType):
202        node.set(
203            "expressions", [e for e in node.expressions if not isinstance(e, exp.DataTypeParam)]
204        )
205
206    return expression
207
208
209def unqualify_unnest(expression: exp.Expression) -> exp.Expression:
210    """Remove references to unnest table aliases, added by the optimizer's qualify_columns step."""
211    from sqlglot.optimizer.scope import find_all_in_scope
212
213    if isinstance(expression, exp.Select):
214        unnest_aliases = {
215            unnest.alias
216            for unnest in find_all_in_scope(expression, exp.Unnest)
217            if isinstance(unnest.parent, (exp.From, exp.Join))
218        }
219        if unnest_aliases:
220            for column in expression.find_all(exp.Column):
221                if column.table in unnest_aliases:
222                    column.set("table", None)
223                elif column.db in unnest_aliases:
224                    column.set("db", None)
225
226    return expression
227
228
229def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
230    """Convert cross join unnest into lateral view explode."""
231    if isinstance(expression, exp.Select):
232        from_ = expression.args.get("from")
233
234        if from_ and isinstance(from_.this, exp.Unnest):
235            unnest = from_.this
236            alias = unnest.args.get("alias")
237            udtf = exp.Posexplode if unnest.args.get("offset") else exp.Explode
238            this, *expressions = unnest.expressions
239            unnest.replace(
240                exp.Table(
241                    this=udtf(
242                        this=this,
243                        expressions=expressions,
244                    ),
245                    alias=exp.TableAlias(this=alias.this, columns=alias.columns) if alias else None,
246                )
247            )
248
249        for join in expression.args.get("joins") or []:
250            unnest = join.this
251
252            if isinstance(unnest, exp.Unnest):
253                alias = unnest.args.get("alias")
254                udtf = exp.Posexplode if unnest.args.get("offset") else exp.Explode
255
256                expression.args["joins"].remove(join)
257
258                for e, column in zip(unnest.expressions, alias.columns if alias else []):
259                    expression.append(
260                        "laterals",
261                        exp.Lateral(
262                            this=udtf(this=e),
263                            view=True,
264                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
265                        ),
266                    )
267
268    return expression
269
270
271def explode_to_unnest(index_offset: int = 0) -> t.Callable[[exp.Expression], exp.Expression]:
272    """Convert explode/posexplode into unnest."""
273
274    def _explode_to_unnest(expression: exp.Expression) -> exp.Expression:
275        if isinstance(expression, exp.Select):
276            from sqlglot.optimizer.scope import Scope
277
278            taken_select_names = set(expression.named_selects)
279            taken_source_names = {name for name, _ in Scope(expression).references}
280
281            def new_name(names: t.Set[str], name: str) -> str:
282                name = find_new_name(names, name)
283                names.add(name)
284                return name
285
286            arrays: t.List[exp.Condition] = []
287            series_alias = new_name(taken_select_names, "pos")
288            series = exp.alias_(
289                exp.Unnest(
290                    expressions=[exp.GenerateSeries(start=exp.Literal.number(index_offset))]
291                ),
292                new_name(taken_source_names, "_u"),
293                table=[series_alias],
294            )
295
296            # we use list here because expression.selects is mutated inside the loop
297            for select in list(expression.selects):
298                explode = select.find(exp.Explode)
299
300                if explode:
301                    pos_alias = ""
302                    explode_alias = ""
303
304                    if isinstance(select, exp.Alias):
305                        explode_alias = select.args["alias"]
306                        alias = select
307                    elif isinstance(select, exp.Aliases):
308                        pos_alias = select.aliases[0]
309                        explode_alias = select.aliases[1]
310                        alias = select.replace(exp.alias_(select.this, "", copy=False))
311                    else:
312                        alias = select.replace(exp.alias_(select, ""))
313                        explode = alias.find(exp.Explode)
314                        assert explode
315
316                    is_posexplode = isinstance(explode, exp.Posexplode)
317                    explode_arg = explode.this
318
319                    if isinstance(explode, exp.ExplodeOuter):
320                        bracket = explode_arg[0]
321                        bracket.set("safe", True)
322                        bracket.set("offset", True)
323                        explode_arg = exp.func(
324                            "IF",
325                            exp.func(
326                                "ARRAY_SIZE", exp.func("COALESCE", explode_arg, exp.Array())
327                            ).eq(0),
328                            exp.array(bracket, copy=False),
329                            explode_arg,
330                        )
331
332                    # This ensures that we won't use [POS]EXPLODE's argument as a new selection
333                    if isinstance(explode_arg, exp.Column):
334                        taken_select_names.add(explode_arg.output_name)
335
336                    unnest_source_alias = new_name(taken_source_names, "_u")
337
338                    if not explode_alias:
339                        explode_alias = new_name(taken_select_names, "col")
340
341                        if is_posexplode:
342                            pos_alias = new_name(taken_select_names, "pos")
343
344                    if not pos_alias:
345                        pos_alias = new_name(taken_select_names, "pos")
346
347                    alias.set("alias", exp.to_identifier(explode_alias))
348
349                    series_table_alias = series.args["alias"].this
350                    column = exp.If(
351                        this=exp.column(series_alias, table=series_table_alias).eq(
352                            exp.column(pos_alias, table=unnest_source_alias)
353                        ),
354                        true=exp.column(explode_alias, table=unnest_source_alias),
355                    )
356
357                    explode.replace(column)
358
359                    if is_posexplode:
360                        expressions = expression.expressions
361                        expressions.insert(
362                            expressions.index(alias) + 1,
363                            exp.If(
364                                this=exp.column(series_alias, table=series_table_alias).eq(
365                                    exp.column(pos_alias, table=unnest_source_alias)
366                                ),
367                                true=exp.column(pos_alias, table=unnest_source_alias),
368                            ).as_(pos_alias),
369                        )
370                        expression.set("expressions", expressions)
371
372                    if not arrays:
373                        if expression.args.get("from"):
374                            expression.join(series, copy=False, join_type="CROSS")
375                        else:
376                            expression.from_(series, copy=False)
377
378                    size: exp.Condition = exp.ArraySize(this=explode_arg.copy())
379                    arrays.append(size)
380
381                    # trino doesn't support left join unnest with on conditions
382                    # if it did, this would be much simpler
383                    expression.join(
384                        exp.alias_(
385                            exp.Unnest(
386                                expressions=[explode_arg.copy()],
387                                offset=exp.to_identifier(pos_alias),
388                            ),
389                            unnest_source_alias,
390                            table=[explode_alias],
391                        ),
392                        join_type="CROSS",
393                        copy=False,
394                    )
395
396                    if index_offset != 1:
397                        size = size - 1
398
399                    expression.where(
400                        exp.column(series_alias, table=series_table_alias)
401                        .eq(exp.column(pos_alias, table=unnest_source_alias))
402                        .or_(
403                            (exp.column(series_alias, table=series_table_alias) > size).and_(
404                                exp.column(pos_alias, table=unnest_source_alias).eq(size)
405                            )
406                        ),
407                        copy=False,
408                    )
409
410            if arrays:
411                end: exp.Condition = exp.Greatest(this=arrays[0], expressions=arrays[1:])
412
413                if index_offset != 1:
414                    end = end - (1 - index_offset)
415                series.expressions[0].set("end", end)
416
417        return expression
418
419    return _explode_to_unnest
420
421
422def add_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
423    """Transforms percentiles by adding a WITHIN GROUP clause to them."""
424    if (
425        isinstance(expression, exp.PERCENTILES)
426        and not isinstance(expression.parent, exp.WithinGroup)
427        and expression.expression
428    ):
429        column = expression.this.pop()
430        expression.set("this", expression.expression.pop())
431        order = exp.Order(expressions=[exp.Ordered(this=column)])
432        expression = exp.WithinGroup(this=expression, expression=order)
433
434    return expression
435
436
437def remove_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
438    """Transforms percentiles by getting rid of their corresponding WITHIN GROUP clause."""
439    if (
440        isinstance(expression, exp.WithinGroup)
441        and isinstance(expression.this, exp.PERCENTILES)
442        and isinstance(expression.expression, exp.Order)
443    ):
444        quantile = expression.this.this
445        input_value = t.cast(exp.Ordered, expression.find(exp.Ordered)).this
446        return expression.replace(exp.ApproxQuantile(this=input_value, quantile=quantile))
447
448    return expression
449
450
451def add_recursive_cte_column_names(expression: exp.Expression) -> exp.Expression:
452    """Uses projection output names in recursive CTE definitions to define the CTEs' columns."""
453    if isinstance(expression, exp.With) and expression.recursive:
454        next_name = name_sequence("_c_")
455
456        for cte in expression.expressions:
457            if not cte.args["alias"].columns:
458                query = cte.this
459                if isinstance(query, exp.SetOperation):
460                    query = query.this
461
462                cte.args["alias"].set(
463                    "columns",
464                    [exp.to_identifier(s.alias_or_name or next_name()) for s in query.selects],
465                )
466
467    return expression
468
469
470def epoch_cast_to_ts(expression: exp.Expression) -> exp.Expression:
471    """Replace 'epoch' in casts by the equivalent date literal."""
472    if (
473        isinstance(expression, (exp.Cast, exp.TryCast))
474        and expression.name.lower() == "epoch"
475        and expression.to.this in exp.DataType.TEMPORAL_TYPES
476    ):
477        expression.this.replace(exp.Literal.string("1970-01-01 00:00:00"))
478
479    return expression
480
481
482def eliminate_semi_and_anti_joins(expression: exp.Expression) -> exp.Expression:
483    """Convert SEMI and ANTI joins into equivalent forms that use EXIST instead."""
484    if isinstance(expression, exp.Select):
485        for join in expression.args.get("joins") or []:
486            on = join.args.get("on")
487            if on and join.kind in ("SEMI", "ANTI"):
488                subquery = exp.select("1").from_(join.this).where(on)
489                exists = exp.Exists(this=subquery)
490                if join.kind == "ANTI":
491                    exists = exists.not_(copy=False)
492
493                join.pop()
494                expression.where(exists, copy=False)
495
496    return expression
497
498
499def eliminate_full_outer_join(expression: exp.Expression) -> exp.Expression:
500    """
501    Converts a query with a FULL OUTER join to a union of identical queries that
502    use LEFT/RIGHT OUTER joins instead. This transformation currently only works
503    for queries that have a single FULL OUTER join.
504    """
505    if isinstance(expression, exp.Select):
506        full_outer_joins = [
507            (index, join)
508            for index, join in enumerate(expression.args.get("joins") or [])
509            if join.side == "FULL"
510        ]
511
512        if len(full_outer_joins) == 1:
513            expression_copy = expression.copy()
514            expression.set("limit", None)
515            index, full_outer_join = full_outer_joins[0]
516            full_outer_join.set("side", "left")
517            expression_copy.args["joins"][index].set("side", "right")
518            expression_copy.args.pop("with", None)  # remove CTEs from RIGHT side
519
520            return exp.union(expression, expression_copy, copy=False)
521
522    return expression
523
524
525def move_ctes_to_top_level(expression: exp.Expression) -> exp.Expression:
526    """
527    Some dialects (e.g. Hive, T-SQL, Spark prior to version 3) only allow CTEs to be
528    defined at the top-level, so for example queries like:
529
530        SELECT * FROM (WITH t(c) AS (SELECT 1) SELECT * FROM t) AS subq
531
532    are invalid in those dialects. This transformation can be used to ensure all CTEs are
533    moved to the top level so that the final SQL code is valid from a syntax standpoint.
534
535    TODO: handle name clashes whilst moving CTEs (it can get quite tricky & costly).
536    """
537    top_level_with = expression.args.get("with")
538    for inner_with in expression.find_all(exp.With):
539        if inner_with.parent is expression:
540            continue
541
542        if not top_level_with:
543            top_level_with = inner_with.pop()
544            expression.set("with", top_level_with)
545        else:
546            if inner_with.recursive:
547                top_level_with.set("recursive", True)
548
549            parent_cte = inner_with.find_ancestor(exp.CTE)
550            inner_with.pop()
551
552            if parent_cte:
553                i = top_level_with.expressions.index(parent_cte)
554                top_level_with.expressions[i:i] = inner_with.expressions
555                top_level_with.set("expressions", top_level_with.expressions)
556            else:
557                top_level_with.set(
558                    "expressions", top_level_with.expressions + inner_with.expressions
559                )
560
561    return expression
562
563
564def ensure_bools(expression: exp.Expression) -> exp.Expression:
565    """Converts numeric values used in conditions into explicit boolean expressions."""
566    from sqlglot.optimizer.canonicalize import ensure_bools
567
568    def _ensure_bool(node: exp.Expression) -> None:
569        if (
570            node.is_number
571            or (
572                not isinstance(node, exp.SubqueryPredicate)
573                and node.is_type(exp.DataType.Type.UNKNOWN, *exp.DataType.NUMERIC_TYPES)
574            )
575            or (isinstance(node, exp.Column) and not node.type)
576        ):
577            node.replace(node.neq(0))
578
579    for node in expression.walk():
580        ensure_bools(node, _ensure_bool)
581
582    return expression
583
584
585def unqualify_columns(expression: exp.Expression) -> exp.Expression:
586    for column in expression.find_all(exp.Column):
587        # We only wanna pop off the table, db, catalog args
588        for part in column.parts[:-1]:
589            part.pop()
590
591    return expression
592
593
594def remove_unique_constraints(expression: exp.Expression) -> exp.Expression:
595    assert isinstance(expression, exp.Create)
596    for constraint in expression.find_all(exp.UniqueColumnConstraint):
597        if constraint.parent:
598            constraint.parent.pop()
599
600    return expression
601
602
603def ctas_with_tmp_tables_to_create_tmp_view(
604    expression: exp.Expression,
605    tmp_storage_provider: t.Callable[[exp.Expression], exp.Expression] = lambda e: e,
606) -> exp.Expression:
607    assert isinstance(expression, exp.Create)
608    properties = expression.args.get("properties")
609    temporary = any(
610        isinstance(prop, exp.TemporaryProperty)
611        for prop in (properties.expressions if properties else [])
612    )
613
614    # CTAS with temp tables map to CREATE TEMPORARY VIEW
615    if expression.kind == "TABLE" and temporary:
616        if expression.expression:
617            return exp.Create(
618                kind="TEMPORARY VIEW",
619                this=expression.this,
620                expression=expression.expression,
621            )
622        return tmp_storage_provider(expression)
623
624    return expression
625
626
627def move_schema_columns_to_partitioned_by(expression: exp.Expression) -> exp.Expression:
628    """
629    In Hive, the PARTITIONED BY property acts as an extension of a table's schema. When the
630    PARTITIONED BY value is an array of column names, they are transformed into a schema.
631    The corresponding columns are removed from the create statement.
632    """
633    assert isinstance(expression, exp.Create)
634    has_schema = isinstance(expression.this, exp.Schema)
635    is_partitionable = expression.kind in {"TABLE", "VIEW"}
636
637    if has_schema and is_partitionable:
638        prop = expression.find(exp.PartitionedByProperty)
639        if prop and prop.this and not isinstance(prop.this, exp.Schema):
640            schema = expression.this
641            columns = {v.name.upper() for v in prop.this.expressions}
642            partitions = [col for col in schema.expressions if col.name.upper() in columns]
643            schema.set("expressions", [e for e in schema.expressions if e not in partitions])
644            prop.replace(exp.PartitionedByProperty(this=exp.Schema(expressions=partitions)))
645            expression.set("this", schema)
646
647    return expression
648
649
650def move_partitioned_by_to_schema_columns(expression: exp.Expression) -> exp.Expression:
651    """
652    Spark 3 supports both "HIVEFORMAT" and "DATASOURCE" formats for CREATE TABLE.
653
654    Currently, SQLGlot uses the DATASOURCE format for Spark 3.
655    """
656    assert isinstance(expression, exp.Create)
657    prop = expression.find(exp.PartitionedByProperty)
658    if (
659        prop
660        and prop.this
661        and isinstance(prop.this, exp.Schema)
662        and all(isinstance(e, exp.ColumnDef) and e.kind for e in prop.this.expressions)
663    ):
664        prop_this = exp.Tuple(
665            expressions=[exp.to_identifier(e.this) for e in prop.this.expressions]
666        )
667        schema = expression.this
668        for e in prop.this.expressions:
669            schema.append("expressions", e)
670        prop.set("this", prop_this)
671
672    return expression
673
674
675def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
676    """Converts struct arguments to aliases, e.g. STRUCT(1 AS y)."""
677    if isinstance(expression, exp.Struct):
678        expression.set(
679            "expressions",
680            [
681                exp.alias_(e.expression, e.this) if isinstance(e, exp.PropertyEQ) else e
682                for e in expression.expressions
683            ],
684        )
685
686    return expression
687
688
689def eliminate_join_marks(expression: exp.Expression) -> exp.Expression:
690    """
691    Remove join marks from an AST. This rule assumes that all marked columns are qualified.
692    If this does not hold for a query, consider running `sqlglot.optimizer.qualify` first.
693
694    For example,
695        SELECT * FROM a, b WHERE a.id = b.id(+)    -- ... is converted to
696        SELECT * FROM a LEFT JOIN b ON a.id = b.id -- this
697
698    Args:
699        expression: The AST to remove join marks from.
700
701    Returns:
702       The AST with join marks removed.
703    """
704    from sqlglot.optimizer.scope import traverse_scope
705
706    for scope in traverse_scope(expression):
707        query = scope.expression
708
709        where = query.args.get("where")
710        joins = query.args.get("joins")
711
712        if not where or not joins:
713            continue
714
715        query_from = query.args["from"]
716
717        # These keep track of the joins to be replaced
718        new_joins: t.Dict[str, exp.Join] = {}
719        old_joins = {join.alias_or_name: join for join in joins}
720
721        for column in scope.columns:
722            if not column.args.get("join_mark"):
723                continue
724
725            predicate = column.find_ancestor(exp.Predicate, exp.Select)
726            assert isinstance(
727                predicate, exp.Binary
728            ), "Columns can only be marked with (+) when involved in a binary operation"
729
730            predicate_parent = predicate.parent
731            join_predicate = predicate.pop()
732
733            left_columns = [
734                c for c in join_predicate.left.find_all(exp.Column) if c.args.get("join_mark")
735            ]
736            right_columns = [
737                c for c in join_predicate.right.find_all(exp.Column) if c.args.get("join_mark")
738            ]
739
740            assert not (
741                left_columns and right_columns
742            ), "The (+) marker cannot appear in both sides of a binary predicate"
743
744            marked_column_tables = set()
745            for col in left_columns or right_columns:
746                table = col.table
747                assert table, f"Column {col} needs to be qualified with a table"
748
749                col.set("join_mark", False)
750                marked_column_tables.add(table)
751
752            assert (
753                len(marked_column_tables) == 1
754            ), "Columns of only a single table can be marked with (+) in a given binary predicate"
755
756            join_this = old_joins.get(col.table, query_from).this
757            new_join = exp.Join(this=join_this, on=join_predicate, kind="LEFT")
758
759            # Upsert new_join into new_joins dictionary
760            new_join_alias_or_name = new_join.alias_or_name
761            existing_join = new_joins.get(new_join_alias_or_name)
762            if existing_join:
763                existing_join.set("on", exp.and_(existing_join.args.get("on"), new_join.args["on"]))
764            else:
765                new_joins[new_join_alias_or_name] = new_join
766
767            # If the parent of the target predicate is a binary node, then it now has only one child
768            if isinstance(predicate_parent, exp.Binary):
769                if predicate_parent.left is None:
770                    predicate_parent.replace(predicate_parent.right)
771                else:
772                    predicate_parent.replace(predicate_parent.left)
773
774        if query_from.alias_or_name in new_joins:
775            only_old_joins = old_joins.keys() - new_joins.keys()
776            assert (
777                len(only_old_joins) >= 1
778            ), "Cannot determine which table to use in the new FROM clause"
779
780            new_from_name = list(only_old_joins)[0]
781            query.set("from", exp.From(this=old_joins[new_from_name].this))
782
783        query.set("joins", list(new_joins.values()))
784
785        if not where.this:
786            where.pop()
787
788    return expression