1"""sqlglot expressions builders."""
   2
   3from __future__ import annotations
   4
   5import re
   6import typing as t
   7
   8from sqlglot.helper import seq_get, ensure_collection, split_num_words
   9from sqlglot.errors import ParseError, TokenError
  10from sqlglot.expressions.core import (
  11    Alias,
  12    Anonymous,
  13    AtTimeZone,
  14    Between,
  15    Binary,
  16    BitwiseNot,
  17    Boolean,
  18    Bracket,
  19    Column,
  20    Condition,
  21    Dot,
  22    EQ,
  23    Expr,
  24    Func,
  25    Identifier,
  26    In,
  27    Literal,
  28    Neg,
  29    Not,
  30    Null,
  31    Paren,
  32    Placeholder,
  33    TABLE_PARTS,
  34    Var,
  35    logger,
  36    SAFE_IDENTIFIER_RE,
  37    maybe_parse,
  38    maybe_copy,
  39    to_identifier,
  40    convert,
  41    alias_,
  42    column,
  43)
  44from sqlglot.expressions.datatypes import DataType, DType, Interval
  45from sqlglot.expressions.query import (
  46    CTE,
  47    From,
  48    Schema,
  49    Select,
  50    SetOperation,
  51    Subquery,
  52    Table,
  53    TableAlias,
  54    Tuple,
  55    Values,
  56    Where,
  57    With,
  58    Query,
  59)
  60from sqlglot.expressions.ddl import Alter, AlterRename, Create, RenameColumn
  61from sqlglot.expressions.dml import Delete, Insert, Merge, Update, When, Whens
  62from sqlglot.expressions.functions import Case, Cast
  63from sqlglot.expressions.array import Array
  64
  65
  66if t.TYPE_CHECKING:
  67    from collections.abc import Sequence, Iterable, Iterator
  68    from sqlglot.dialects.dialect import DialectType
  69    from sqlglot.expressions.core import ExpOrStr
  70    from sqlglot.expressions.datatypes import DATA_TYPE
  71    from sqlglot._typing import ParserArgs, ParserNoDialectArgs, E, P
  72    from typing_extensions import Unpack, Concatenate
  73
  74
  75def select(
  76    *expressions: ExpOrStr,
  77    dialect: DialectType = None,
  78    copy: bool = True,
  79    **opts: Unpack[ParserNoDialectArgs],
  80) -> Select:
  81    """
  82    Initializes a syntax tree from one or multiple SELECT expressions.
  83
  84    Example:
  85        >>> select("col1", "col2").from_("tbl").sql()
  86        'SELECT col1, col2 FROM tbl'
  87
  88    Args:
  89        *expressions: the SQL code string to parse as the expressions of a
  90            SELECT statement. If an Expr instance is passed, this is used as-is.
  91        dialect: the dialect used to parse the input expressions (in the case that an
  92            input expression is a SQL string).
  93        **opts: other options to use to parse the input expressions (again, in the case
  94            that an input expression is a SQL string).
  95
  96    Returns:
  97        Select: the syntax tree for the SELECT statement.
  98    """
  99    return Select().select(*expressions, dialect=dialect, copy=copy, **opts)
 100
 101
 102def from_(
 103    expression: ExpOrStr,
 104    dialect: DialectType = None,
 105    copy: bool = True,
 106    **opts: Unpack[ParserNoDialectArgs],
 107) -> Select:
 108    """
 109    Initializes a syntax tree from a FROM expression.
 110
 111    Example:
 112        >>> from_("tbl").select("col1", "col2").sql()
 113        'SELECT col1, col2 FROM tbl'
 114
 115    Args:
 116        *expression: the SQL code string to parse as the FROM expressions of a
 117            SELECT statement. If an Expr instance is passed, this is used as-is.
 118        dialect: the dialect used to parse the input expression (in the case that the
 119            input expression is a SQL string).
 120        **opts: other options to use to parse the input expressions (again, in the case
 121            that the input expression is a SQL string).
 122
 123    Returns:
 124        Select: the syntax tree for the SELECT statement.
 125    """
 126    return Select().from_(expression, dialect=dialect, copy=copy, **opts)
 127
 128
 129def update(
 130    table: str | Table,
 131    properties: dict[str, object] | None = None,
 132    where: ExpOrStr | None = None,
 133    from_: ExpOrStr | None = None,
 134    with_: dict[str, ExpOrStr] | None = None,
 135    dialect: DialectType = None,
 136    copy: bool = True,
 137    **opts: Unpack[ParserNoDialectArgs],
 138) -> Update:
 139    """
 140    Creates an update statement.
 141
 142    Example:
 143        >>> update("my_table", {"x": 1, "y": "2", "z": None}, from_="baz_cte", where="baz_cte.id > 1 and my_table.id = baz_cte.id", with_={"baz_cte": "SELECT id FROM foo"}).sql()
 144        "WITH baz_cte AS (SELECT id FROM foo) UPDATE my_table SET x = 1, y = '2', z = NULL FROM baz_cte WHERE baz_cte.id > 1 AND my_table.id = baz_cte.id"
 145
 146    Args:
 147        properties: dictionary of properties to SET which are
 148            auto converted to sql objects eg None -> NULL
 149        where: sql conditional parsed into a WHERE statement
 150        from_: sql statement parsed into a FROM statement
 151        with_: dictionary of CTE aliases / select statements to include in a WITH clause.
 152        dialect: the dialect used to parse the input expressions.
 153        copy: whether to copy the input expressions.
 154        **opts: other options to use to parse the input expressions.
 155
 156    Returns:
 157        Update: the syntax tree for the UPDATE statement.
 158    """
 159    update_expr = Update(this=maybe_parse(table, into=Table, dialect=dialect, copy=copy))
 160    if properties:
 161        update_expr.set(
 162            "expressions",
 163            [
 164                EQ(this=maybe_parse(k, dialect=dialect, copy=copy, **opts), expression=convert(v))
 165                for k, v in properties.items()
 166            ],
 167        )
 168    if from_:
 169        update_expr.set(
 170            "from_",
 171            maybe_parse(from_, into=From, dialect=dialect, prefix="FROM", copy=copy, **opts),
 172        )
 173    if isinstance(where, Condition):
 174        where = Where(this=where)
 175    if where:
 176        update_expr.set(
 177            "where",
 178            maybe_parse(where, into=Where, dialect=dialect, prefix="WHERE", copy=copy, **opts),
 179        )
 180    if with_:
 181        cte_list = [
 182            alias_(
 183                CTE(this=maybe_parse(qry, dialect=dialect, copy=copy, **opts)), alias, table=True
 184            )
 185            for alias, qry in with_.items()
 186        ]
 187        update_expr.set(
 188            "with_",
 189            With(expressions=cte_list),
 190        )
 191    return update_expr
 192
 193
 194def delete(
 195    table: ExpOrStr,
 196    where: ExpOrStr | None = None,
 197    returning: ExpOrStr | None = None,
 198    dialect: DialectType = None,
 199    **opts: Unpack[ParserNoDialectArgs],
 200) -> Delete:
 201    """
 202    Builds a delete statement.
 203
 204    Example:
 205        >>> delete("my_table", where="id > 1").sql()
 206        'DELETE FROM my_table WHERE id > 1'
 207
 208    Args:
 209        where: sql conditional parsed into a WHERE statement
 210        returning: sql conditional parsed into a RETURNING statement
 211        dialect: the dialect used to parse the input expressions.
 212        **opts: other options to use to parse the input expressions.
 213
 214    Returns:
 215        Delete: the syntax tree for the DELETE statement.
 216    """
 217    delete_expr = Delete().delete(table, dialect=dialect, copy=False, **opts)
 218    if where:
 219        delete_expr = delete_expr.where(where, dialect=dialect, copy=False, **opts)
 220    if returning:
 221        delete_expr = delete_expr.returning(returning, dialect=dialect, copy=False, **opts)
 222    return delete_expr
 223
 224
 225def insert(
 226    expression: ExpOrStr,
 227    into: str | Table,
 228    columns: Sequence[str | Identifier] | None = None,
 229    overwrite: bool | None = None,
 230    returning: ExpOrStr | None = None,
 231    dialect: DialectType = None,
 232    copy: bool = True,
 233    **opts: Unpack[ParserNoDialectArgs],
 234) -> Insert:
 235    """
 236    Builds an INSERT statement.
 237
 238    Example:
 239        >>> insert("VALUES (1, 2, 3)", "tbl").sql()
 240        'INSERT INTO tbl VALUES (1, 2, 3)'
 241
 242    Args:
 243        expression: the sql string or expression of the INSERT statement
 244        into: the tbl to insert data to.
 245        columns: optionally the table's column names.
 246        overwrite: whether to INSERT OVERWRITE or not.
 247        returning: sql conditional parsed into a RETURNING statement
 248        dialect: the dialect used to parse the input expressions.
 249        copy: whether to copy the expression.
 250        **opts: other options to use to parse the input expressions.
 251
 252    Returns:
 253        Insert: the syntax tree for the INSERT statement.
 254    """
 255    expr = maybe_parse(expression, dialect=dialect, copy=copy, **opts)
 256    this: Table | Schema = maybe_parse(into, into=Table, dialect=dialect, copy=copy, **opts)
 257
 258    if columns:
 259        this = Schema(this=this, expressions=[to_identifier(c, copy=copy) for c in columns])
 260
 261    insert = Insert(this=this, expression=expr, overwrite=overwrite)
 262
 263    if returning:
 264        insert = insert.returning(returning, dialect=dialect, copy=False, **opts)
 265
 266    return insert
 267
 268
 269def merge(
 270    *when_exprs: ExpOrStr,
 271    into: ExpOrStr,
 272    using: ExpOrStr,
 273    on: ExpOrStr,
 274    returning: ExpOrStr | None = None,
 275    dialect: DialectType = None,
 276    copy: bool = True,
 277    **opts: Unpack[ParserNoDialectArgs],
 278) -> Merge:
 279    """
 280    Builds a MERGE statement.
 281
 282    Example:
 283        >>> merge("WHEN MATCHED THEN UPDATE SET col1 = source_table.col1",
 284        ...       "WHEN NOT MATCHED THEN INSERT (col1) VALUES (source_table.col1)",
 285        ...       into="my_table",
 286        ...       using="source_table",
 287        ...       on="my_table.id = source_table.id").sql()
 288        'MERGE INTO my_table USING source_table ON my_table.id = source_table.id WHEN MATCHED THEN UPDATE SET col1 = source_table.col1 WHEN NOT MATCHED THEN INSERT (col1) VALUES (source_table.col1)'
 289
 290    Args:
 291        *when_exprs: The WHEN clauses specifying actions for matched and unmatched rows.
 292        into: The target table to merge data into.
 293        using: The source table to merge data from.
 294        on: The join condition for the merge.
 295        returning: The columns to return from the merge.
 296        dialect: The dialect used to parse the input expressions.
 297        copy: Whether to copy the expression.
 298        **opts: Other options to use to parse the input expressions.
 299
 300    Returns:
 301        Merge: The syntax tree for the MERGE statement.
 302    """
 303    expressions: list[Expr] = []
 304    for when_expr in when_exprs:
 305        expression = maybe_parse(when_expr, dialect=dialect, copy=copy, into=Whens, **opts)
 306        expressions.extend([expression] if isinstance(expression, When) else expression.expressions)
 307
 308    merge = Merge(
 309        this=maybe_parse(into, dialect=dialect, copy=copy, **opts),
 310        using=maybe_parse(using, dialect=dialect, copy=copy, **opts),
 311        on=maybe_parse(on, dialect=dialect, copy=copy, **opts),
 312        whens=Whens(expressions=expressions),
 313    )
 314    if returning:
 315        merge = merge.returning(returning, dialect=dialect, copy=False, **opts)
 316
 317    if isinstance(using_clause := merge.args.get("using"), Alias):
 318        using_clause.replace(alias_(using_clause.this, using_clause.args["alias"], table=True))
 319
 320    return merge
 321
 322
 323def parse_identifier(name: str | Identifier, dialect: DialectType = None) -> Identifier:
 324    """
 325    Parses a given string into an identifier.
 326
 327    Args:
 328        name: The name to parse into an identifier.
 329        dialect: The dialect to parse against.
 330
 331    Returns:
 332        The identifier ast node.
 333    """
 334    if isinstance(name, str) and SAFE_IDENTIFIER_RE.match(name):
 335        # Simple names parse to a single unquoted identifier in all dialects, so we can
 336        # avoid the tokenizer/parser round-trip for them.
 337        return Identifier(this=name, quoted=False)
 338
 339    try:
 340        expression = maybe_parse(name, dialect=dialect, into=Identifier)
 341    except (ParseError, TokenError):
 342        expression = to_identifier(name)
 343
 344    return expression
 345
 346
 347INTERVAL_STRING_RE = re.compile(r"\s*(-?[0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]+)\s*")
 348
 349
 350INTERVAL_DAY_TIME_RE = re.compile(
 351    r"\s*-?\s*\d+(?:\.\d+)?\s+(?:-?(?:\d+:)?\d+:\d+(?:\.\d+)?|-?(?:\d+:){1,2}|:)\s*"
 352)
 353
 354
 355def to_interval(interval: str | Expr) -> Interval:
 356    """Builds an interval expression from a string like '1 day' or '5 months'."""
 357    if isinstance(interval, Literal):
 358        if not interval.is_string:
 359            raise ValueError("Invalid interval string.")
 360
 361        interval = interval.this
 362
 363    interval = maybe_parse(f"INTERVAL {interval}")
 364    assert isinstance(interval, Interval)
 365    return interval
 366
 367
 368def to_table(
 369    sql_path: str | Table, dialect: DialectType = None, copy: bool = True, **kwargs: object
 370) -> Table:
 371    """
 372    Create a table expression from a `[catalog].[schema].[table]` sql path. Catalog and schema are optional.
 373    If a table is passed in then that table is returned.
 374
 375    Args:
 376        sql_path: a `[catalog].[schema].[table]` string.
 377        dialect: the source dialect according to which the table name will be parsed.
 378        copy: Whether to copy a table if it is passed in.
 379        kwargs: the kwargs to instantiate the resulting `Table` expression with.
 380
 381    Returns:
 382        A table expression.
 383    """
 384    if isinstance(sql_path, Table):
 385        return maybe_copy(sql_path, copy=copy)
 386
 387    try:
 388        table = maybe_parse(sql_path, into=Table, dialect=dialect)
 389    except ParseError:
 390        catalog, db, this = split_num_words(sql_path, ".", 3)
 391
 392        if not this:
 393            raise
 394
 395        table = table_(this, db=db, catalog=catalog)
 396
 397    return table.set_kwargs(kwargs)
 398
 399
 400def to_column(
 401    sql_path: str | Column,
 402    quoted: bool | None = None,
 403    dialect: DialectType = None,
 404    copy: bool = True,
 405    **kwargs: t.Any,
 406) -> Column | Dot:
 407    """
 408    Create a column from a `[table].[column]` sql path. Table is optional.
 409    If a column is passed in then that column is returned.
 410
 411    Args:
 412        sql_path: a `[table].[column]` string.
 413        quoted: Whether or not to force quote identifiers.
 414        dialect: the source dialect according to which the column name will be parsed.
 415        copy: Whether to copy a column if it is passed in.
 416        kwargs: the kwargs to instantiate the resulting `Column` expression with.
 417
 418    Returns:
 419        A column expression.
 420    """
 421    if isinstance(sql_path, Column):
 422        return maybe_copy(sql_path, copy=copy)
 423
 424    try:
 425        col = maybe_parse(sql_path, into=Column, dialect=dialect)
 426    except ParseError:
 427        return column(*reversed(sql_path.split(".")), quoted=quoted, **kwargs)
 428
 429    for k, v in kwargs.items():
 430        col.set(k, v)
 431
 432    if quoted:
 433        for i in col.find_all(Identifier):
 434            i.set("quoted", True)
 435
 436    return col
 437
 438
 439def subquery(
 440    expression: ExpOrStr,
 441    alias: Identifier | str | None = None,
 442    dialect: DialectType = None,
 443    copy: bool = True,
 444    **opts: Unpack[ParserNoDialectArgs],
 445) -> Select:
 446    """
 447    Build a subquery expression that's selected from.
 448
 449    Example:
 450        >>> subquery('select x from tbl', 'bar').select('x').sql()
 451        'SELECT x FROM (SELECT x FROM tbl) AS bar'
 452
 453    Args:
 454        expression: the SQL code strings to parse.
 455            If an Expr instance is passed, this is used as-is.
 456        alias: the alias name to use.
 457        dialect: the dialect used to parse the input expression.
 458        **opts: other options to use to parse the input expressions.
 459
 460    Returns:
 461        A new Select instance with the subquery expression included.
 462    """
 463    expr = (
 464        maybe_parse(expression, dialect=dialect, **opts).assert_is(Query).subquery(alias, copy=copy)
 465    )
 466    return Select().from_(expr, dialect=dialect, **opts)
 467
 468
 469def cast(
 470    expression: ExpOrStr,
 471    to: DATA_TYPE,
 472    copy: bool = True,
 473    dialect: DialectType = None,
 474    **opts: Unpack[ParserNoDialectArgs],
 475) -> Cast:
 476    """Cast an expression to a data type.
 477
 478    Example:
 479        >>> cast('x + 1', 'int').sql()
 480        'CAST(x + 1 AS INT)'
 481
 482    Args:
 483        expression: The expression to cast.
 484        to: The datatype to cast to.
 485        copy: Whether to copy the supplied expressions.
 486        dialect: The target dialect. This is used to prevent a re-cast in the following scenario:
 487            - The expression to be cast is already a exp.Cast expression
 488            - The existing cast is to a type that is logically equivalent to new type
 489
 490            For example, if :expression='CAST(x as DATETIME)' and :to=Type.TIMESTAMP,
 491            but in the target dialect DATETIME is mapped to TIMESTAMP, then we will NOT return `CAST(x (as DATETIME) as TIMESTAMP)`
 492            and instead just return the original expression `CAST(x as DATETIME)`.
 493
 494            This is to prevent it being output as a double cast `CAST(x (as TIMESTAMP) as TIMESTAMP)` once the DATETIME -> TIMESTAMP
 495            mapping is applied in the target dialect generator.
 496
 497    Returns:
 498        The new Cast instance.
 499    """
 500    expr = maybe_parse(expression, copy=copy, dialect=dialect, **opts)
 501    data_type = DataType.build(to, copy=copy, dialect=dialect, **opts)
 502
 503    # dont re-cast if the expression is already a cast to the correct type
 504    if isinstance(expr, Cast):
 505        from sqlglot.dialects.dialect import Dialect
 506
 507        target_dialect = Dialect.get_or_raise(dialect)
 508        type_mapping = target_dialect.generator_class.TYPE_MAPPING
 509
 510        existing_cast_type = expr.to.this
 511        new_cast_type: DType = data_type.this
 512        # `this` is only a plain type enum for simple types; complex ones such as
 513        # INTERVAL nest another expression there, so the equivalence check is skipped.
 514        types_are_equivalent = (
 515            isinstance(existing_cast_type, DType)
 516            and isinstance(new_cast_type, DType)
 517            and type_mapping.get(existing_cast_type, existing_cast_type.value)
 518            == type_mapping.get(new_cast_type, new_cast_type.value)
 519        )
 520
 521        if expr.is_type(data_type) or types_are_equivalent:
 522            return expr
 523
 524    expr = Cast(this=expr, to=data_type)
 525    expr.type = data_type
 526
 527    return expr
 528
 529
 530def table_(
 531    table: Identifier | str | None,
 532    db: Identifier | str | None = None,
 533    catalog: Identifier | str | None = None,
 534    quoted: bool | None = None,
 535    alias: Identifier | str | None = None,
 536) -> Table:
 537    """Build a Table.
 538
 539    Args:
 540        table: Table name.
 541        db: Database name.
 542        catalog: Catalog name.
 543        quote: Whether to force quotes on the table's identifiers.
 544        alias: Table's alias.
 545
 546    Returns:
 547        The new Table instance.
 548    """
 549    return Table(
 550        this=to_identifier(table, quoted=quoted) if table else None,
 551        db=to_identifier(db, quoted=quoted) if db else None,
 552        catalog=to_identifier(catalog, quoted=quoted) if catalog else None,
 553        alias=TableAlias(this=to_identifier(alias)) if alias else None,
 554    )
 555
 556
 557def values(
 558    values: Iterable[tuple[object, ...] | Tuple],
 559    alias: str | None = None,
 560    columns: Iterable[str] | dict[str, DataType] | None = None,
 561) -> Values:
 562    """Build VALUES statement.
 563
 564    Example:
 565        >>> values([(1, '2')]).sql()
 566        "VALUES (1, '2')"
 567
 568    Args:
 569        values: values statements that will be converted to SQL
 570        alias: optional alias
 571        columns: Optional list of ordered column names or ordered dictionary of column names to types.
 572         If either are provided then an alias is also required.
 573
 574    Returns:
 575        Values: the Values expression object
 576    """
 577    if columns and not alias:
 578        raise ValueError("Alias is required when providing columns")
 579
 580    return Values(
 581        expressions=[convert(tup) for tup in values],
 582        alias=(
 583            TableAlias(this=to_identifier(alias), columns=[to_identifier(x) for x in columns])
 584            if columns
 585            else (TableAlias(this=to_identifier(alias)) if alias else None)
 586        ),
 587    )
 588
 589
 590def var(name: ExpOrStr | None) -> Var:
 591    """Build a SQL variable.
 592
 593    Example:
 594        >>> repr(var('x'))
 595        'Var(this=x)'
 596
 597        >>> repr(var(column('x', table='y')))
 598        'Var(this=x)'
 599
 600    Args:
 601        name: The name of the var or an expression who's name will become the var.
 602
 603    Returns:
 604        The new variable node.
 605    """
 606    if not name:
 607        raise ValueError("Cannot convert empty name into var.")
 608
 609    if isinstance(name, Expr):
 610        name = name.name
 611    return Var(this=name)
 612
 613
 614def rename_table(
 615    old_name: str | Table,
 616    new_name: str | Table,
 617    dialect: DialectType = None,
 618) -> Alter:
 619    """Build ALTER TABLE... RENAME... expression
 620
 621    Args:
 622        old_name: The old name of the table
 623        new_name: The new name of the table
 624        dialect: The dialect to parse the table.
 625
 626    Returns:
 627        Alter table expression
 628    """
 629    old_table = to_table(old_name, dialect=dialect)
 630    new_table = to_table(new_name, dialect=dialect)
 631    return Alter(
 632        this=old_table,
 633        kind="TABLE",
 634        actions=[
 635            AlterRename(this=new_table),
 636        ],
 637    )
 638
 639
 640def rename_column(
 641    table_name: str | Table,
 642    old_column_name: str | Column,
 643    new_column_name: str | Column,
 644    exists: bool | None = None,
 645    dialect: DialectType = None,
 646) -> Alter:
 647    """Build ALTER TABLE... RENAME COLUMN... expression
 648
 649    Args:
 650        table_name: Name of the table
 651        old_column: The old name of the column
 652        new_column: The new name of the column
 653        exists: Whether to add the `IF EXISTS` clause
 654        dialect: The dialect to parse the table/column.
 655
 656    Returns:
 657        Alter table expression
 658    """
 659    table = to_table(table_name, dialect=dialect)
 660    old_column = to_column(old_column_name, dialect=dialect)
 661    new_column = to_column(new_column_name, dialect=dialect)
 662    return Alter(
 663        this=table,
 664        kind="TABLE",
 665        actions=[
 666            RenameColumn(this=old_column, to=new_column, exists=exists),
 667        ],
 668    )
 669
 670
 671def replace_children(
 672    expression: Expr,
 673    fun: t.Callable[Concatenate[Expr, P], object],
 674    *args: P.args,
 675    **kwargs: P.kwargs,
 676) -> None:
 677    """
 678    Replace children of an expression with the result of a lambda fun(child) -> exp.
 679    """
 680    for k, v in tuple(expression.args.items()):
 681        is_list_arg = type(v) is list
 682
 683        child_nodes = v if is_list_arg else [v]
 684        new_child_nodes = []
 685
 686        for cn in child_nodes:
 687            if isinstance(cn, Expr):
 688                for child_node in ensure_collection(fun(cn, *args, **kwargs)):
 689                    new_child_nodes.append(child_node)
 690            else:
 691                new_child_nodes.append(cn)
 692
 693        if is_list_arg:
 694            expression.set(k, new_child_nodes)
 695        else:
 696            expression.set(k, seq_get(new_child_nodes, 0))
 697
 698
 699def replace_tree(
 700    expression: Expr,
 701    fun: t.Callable[[Expr], Expr],
 702    prune: t.Callable[[Expr], bool] | None = None,
 703) -> Expr:
 704    """
 705    Replace an entire tree with the result of function calls on each node.
 706
 707    This will be traversed in reverse dfs, so leaves first.
 708    If new nodes are created as a result of function calls, they will also be traversed.
 709    """
 710    stack = list(expression.dfs(prune=prune))
 711
 712    while stack:
 713        node = stack.pop()
 714        new_node = fun(node)
 715
 716        if new_node is not node:
 717            node.replace(new_node)
 718
 719            if isinstance(new_node, Expr):
 720                stack.append(new_node)
 721
 722    return new_node
 723
 724
 725def find_tables(expression: Expr) -> set[Table]:
 726    """
 727    Find all tables referenced in a query.
 728
 729    Args:
 730        expressions: The query to find the tables in.
 731
 732    Returns:
 733        A set of all the tables.
 734    """
 735    from sqlglot.optimizer.scope import traverse_scope
 736
 737    return {
 738        table
 739        for scope in traverse_scope(expression)
 740        for table in scope.tables
 741        if isinstance(table, Table) and table.name and table.name not in scope.cte_sources
 742    }
 743
 744
 745def column_table_names(expression: Expr, exclude: str = "") -> set[str]:
 746    """
 747    Return all table names referenced through columns in an expression.
 748
 749    Example:
 750        >>> import sqlglot
 751        >>> sorted(column_table_names(sqlglot.parse_one("a.b AND c.d AND c.e")))
 752        ['a', 'c']
 753
 754    Args:
 755        expression: expression to find table names.
 756        exclude: a table name to exclude
 757
 758    Returns:
 759        A list of unique names.
 760    """
 761    return {
 762        table
 763        for table in (column.table for column in expression.find_all(Column))
 764        if table and table != exclude
 765    }
 766
 767
 768def table_name(table: Table | str, dialect: DialectType = None, identify: bool = False) -> str:
 769    """Get the full name of a table as a string.
 770
 771    Args:
 772        table: Table expression node or string.
 773        dialect: The dialect to generate the table name for.
 774        identify: Determines when an identifier should be quoted. Possible values are:
 775            False (default): Never quote, except in cases where it's mandatory by the dialect.
 776            True: Always quote.
 777
 778    Examples:
 779        >>> from sqlglot import exp, parse_one
 780        >>> table_name(parse_one("select * from a.b.c").find(exp.Table))
 781        'a.b.c'
 782
 783    Returns:
 784        The table name.
 785    """
 786
 787    expr = maybe_parse(table, into=Table, dialect=dialect)
 788
 789    if not expr:
 790        raise ValueError(f"Cannot parse {table}")
 791
 792    return ".".join(
 793        (
 794            part.sql(dialect=dialect, identify=True, copy=False, comments=False)
 795            if identify or not SAFE_IDENTIFIER_RE.match(part.name)
 796            else part.name
 797        )
 798        for part in expr.parts
 799    )
 800
 801
 802def normalize_table_name(table: str | Table, dialect: DialectType = None, copy: bool = True) -> str:
 803    """Returns a case normalized table name without quotes.
 804
 805    Args:
 806        table: the table to normalize
 807        dialect: the dialect to use for normalization rules
 808        copy: whether to copy the expression.
 809
 810    Examples:
 811        >>> normalize_table_name("`A-B`.c", dialect="bigquery")
 812        'A-B.c'
 813    """
 814    from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 815
 816    return ".".join(
 817        p.name
 818        for p in normalize_identifiers(
 819            to_table(table, dialect=dialect, copy=copy), dialect=dialect
 820        ).parts
 821    )
 822
 823
 824def replace_tables(
 825    expression: E, mapping: dict[str, str], dialect: DialectType = None, copy: bool = True
 826) -> E:
 827    """Replace all tables in expression according to the mapping.
 828
 829    Args:
 830        expression: expression node to be transformed and replaced.
 831        mapping: mapping of table names.
 832        dialect: the dialect of the mapping table
 833        copy: whether to copy the expression.
 834
 835    Examples:
 836        >>> from sqlglot import exp, parse_one
 837        >>> replace_tables(parse_one("select * from a.b"), {"a.b": "c"}).sql()
 838        'SELECT * FROM c /* a.b */'
 839
 840    Returns:
 841        The mapped expression.
 842    """
 843
 844    mapping = {normalize_table_name(k, dialect=dialect): v for k, v in mapping.items()}
 845
 846    def _replace_tables(node: Expr) -> Expr:
 847        if isinstance(node, Table) and node.meta_get("replace") is not False:
 848            original = normalize_table_name(node, dialect=dialect)
 849            new_name = mapping.get(original)
 850
 851            if new_name:
 852                table = to_table(
 853                    new_name,
 854                    **{k: v for k, v in node.args.items() if k not in TABLE_PARTS},
 855                    dialect=dialect,
 856                )
 857                table.add_comments([original])
 858                return table
 859        return node
 860
 861    return expression.transform(_replace_tables, copy=copy)  # type: ignore
 862
 863
 864OPERATORS = (AtTimeZone, Between, Binary, BitwiseNot, In, Neg, Not)
 865
 866
 867def _is_operand(placeholder: Placeholder) -> bool:
 868    parent = placeholder.parent
 869    if isinstance(parent, (Bracket, In)):
 870        return placeholder.arg_key == "this"
 871    return isinstance(parent, OPERATORS)
 872
 873
 874def _takes_query(placeholder: Placeholder) -> bool:
 875    parent = placeholder.parent
 876    if isinstance(parent, In):
 877        return placeholder.arg_key == "expressions" and len(parent.expressions) == 1
 878    if isinstance(parent, Func):
 879        return len(list(parent.iter_expressions())) == 1
 880    if isinstance(parent, Table):
 881        return isinstance(parent.parent, (Create, Subquery))
 882    return not parent or isinstance(parent, (CTE, Paren, SetOperation))
 883
 884
 885def _replace_placeholder(placeholder: Placeholder, expression: Expr) -> Expr:
 886    if (
 887        isinstance(expression, Query)
 888        and not isinstance(expression, Subquery)
 889        and not _takes_query(placeholder)
 890    ):
 891        return expression.subquery(copy=False)
 892    if (
 893        _is_operand(placeholder)
 894        and isinstance(expression, OPERATORS)
 895        and not isinstance(expression, Dot)
 896    ):
 897        return Paren(this=expression)
 898    return expression
 899
 900
 901def replace_placeholders(expression: Expr, *args: object, **kwargs: t.Any) -> Expr:
 902    """Replace placeholders in an expression.
 903
 904    Args:
 905        expression: expression node to be transformed and replaced.
 906        args: positional names that will substitute unnamed placeholders in the given order.
 907        kwargs: keyword arguments that will substitute named placeholders.
 908
 909    Examples:
 910        >>> from sqlglot import exp, parse_one
 911        >>> replace_placeholders(
 912        ...     parse_one("select * from :tbl where ? = ?"),
 913        ...     exp.to_identifier("str_col"), "b", tbl=exp.to_identifier("foo")
 914        ... ).sql()
 915        "SELECT * FROM foo WHERE str_col = 'b'"
 916
 917    Returns:
 918        The mapped expression.
 919    """
 920
 921    def _replace_placeholders(node: Expr, args: Iterator[object], **kwargs: object) -> Expr:
 922        if isinstance(node, Placeholder):
 923            if node.this:
 924                new_name = kwargs.get(node.this)
 925                if new_name is not None:
 926                    return _replace_placeholder(node, convert(new_name))
 927            else:
 928                try:
 929                    return _replace_placeholder(node, convert(next(args)))
 930                except StopIteration:
 931                    pass
 932        return node
 933
 934    return expression.transform(_replace_placeholders, iter(args), **kwargs)
 935
 936
 937def expand(
 938    expression: Expr,
 939    sources: dict[str, Query | t.Callable[[], Query]],
 940    dialect: DialectType = None,
 941    copy: bool = True,
 942) -> Expr:
 943    """Transforms an expression by expanding all referenced sources into subqueries.
 944
 945    Examples:
 946        >>> from sqlglot import parse_one
 947        >>> expand(parse_one("select * from x AS z"), {"x": parse_one("select * from y")}).sql()
 948        'SELECT * FROM (SELECT * FROM y) AS z /* source: x */'
 949
 950        >>> expand(parse_one("select * from x AS z"), {"x": parse_one("select * from y"), "y": parse_one("select * from z")}).sql()
 951        'SELECT * FROM (SELECT * FROM (SELECT * FROM z) AS y /* source: y */) AS z /* source: x */'
 952
 953    Args:
 954        expression: The expression to expand.
 955        sources: A dict of name to query or a callable that provides a query on demand.
 956        dialect: The dialect of the sources dict or the callable.
 957        copy: Whether to copy the expression during transformation. Defaults to True.
 958
 959    Returns:
 960        The transformed expression.
 961    """
 962    normalized_sources = {normalize_table_name(k, dialect=dialect): v for k, v in sources.items()}
 963
 964    def _expand(node: Expr):
 965        if isinstance(node, Table):
 966            name = normalize_table_name(node, dialect=dialect)
 967            source = normalized_sources.get(name)
 968
 969            if source:
 970                # Create a subquery with the same alias (or table name if no alias)
 971                parsed_source = source() if callable(source) else source
 972                subquery = parsed_source.subquery(node.alias or name)
 973                subquery.comments = [f"source: {name}"]
 974
 975                # Continue expanding within the subquery
 976                return subquery.transform(_expand, copy=False)
 977
 978        return node
 979
 980    return expression.transform(_expand, copy=copy)
 981
 982
 983def func(
 984    name: str, *args: t.Any, copy: bool = True, dialect: DialectType = None, **kwargs: t.Any
 985) -> Func:
 986    """
 987    Returns a Func expression.
 988
 989    Examples:
 990        >>> func("abs", 5).sql()
 991        'ABS(5)'
 992
 993        >>> func("cast", this=5, to=DataType.build("DOUBLE")).sql()
 994        'CAST(5 AS DOUBLE)'
 995
 996    Args:
 997        name: the name of the function to build.
 998        args: the args used to instantiate the function of interest.
 999        copy: whether to copy the argument expressions.
1000        dialect: the source dialect.
1001        kwargs: the kwargs used to instantiate the function of interest.
1002
1003    Note:
1004        The arguments `args` and `kwargs` are mutually exclusive.
1005
1006    Returns:
1007        An instance of the function of interest, or an anonymous function, if `name` doesn't
1008        correspond to an existing `sqlglot.expressions.Func` class.
1009    """
1010    if args and kwargs:
1011        raise ValueError("Can't use both args and kwargs to instantiate a function.")
1012
1013    from sqlglot.dialects.dialect import Dialect
1014
1015    dialect = Dialect.get_or_raise(dialect)
1016
1017    converted: list[Expr] = [maybe_parse(arg, dialect=dialect, copy=copy) for arg in args]
1018    kwargs = {key: maybe_parse(value, dialect=dialect, copy=copy) for key, value in kwargs.items()}
1019
1020    constructor = dialect.parser_class.FUNCTIONS.get(name.upper())
1021    if constructor:
1022        if converted:
1023            try:
1024                function = constructor(converted)
1025            except TypeError:
1026                function = constructor(converted, dialect=dialect)
1027        elif constructor.__name__ == "from_arg_list":
1028            function = constructor.__self__(**kwargs)  # type: ignore
1029        else:
1030            from sqlglot.expressions import FUNCTION_BY_NAME as _FUNCTION_BY_NAME
1031
1032            constructor = _FUNCTION_BY_NAME.get(name.upper())
1033            if constructor:
1034                function = constructor(**kwargs)
1035            else:
1036                raise ValueError(
1037                    f"Unable to convert '{name}' into a Func. Either manually construct "
1038                    "the Func expression of interest or parse the function call."
1039                )
1040    else:
1041        kwargs = kwargs or {"expressions": converted}
1042        function = Anonymous(this=name, **kwargs)
1043
1044    for error_message in function.error_messages(converted):
1045        raise ValueError(error_message)
1046
1047    return function
1048
1049
1050def case(
1051    expression: ExpOrStr | None = None,
1052    copy: bool = True,
1053    **opts: Unpack[ParserArgs],
1054) -> Case:
1055    """
1056    Initialize a CASE statement.
1057
1058    Example:
1059        case().when("a = 1", "foo").else_("bar")
1060
1061    Args:
1062        expression: Optionally, the input expression (not all dialects support this)
1063        copy: whether to copy the argument expressions.
1064        **opts: Extra keyword arguments for parsing `expression`
1065    """
1066    if expression is not None:
1067        this = maybe_parse(expression, copy=copy, **opts)
1068    else:
1069        this = None
1070    return Case(this=this, ifs=[])
1071
1072
1073def array(
1074    *expressions: ExpOrStr,
1075    copy: bool = True,
1076    dialect: DialectType = None,
1077    **kwargs: Unpack[ParserNoDialectArgs],
1078) -> Array:
1079    """
1080    Returns an array.
1081
1082    Examples:
1083        >>> array(1, 'x').sql()
1084        'ARRAY(1, x)'
1085
1086    Args:
1087        expressions: the expressions to add to the array.
1088        copy: whether to copy the argument expressions.
1089        dialect: the source dialect.
1090        kwargs: the kwargs used to instantiate the function of interest.
1091
1092    Returns:
1093        An array expression.
1094    """
1095    return Array(
1096        expressions=[
1097            maybe_parse(expression, copy=copy, dialect=dialect, **kwargs)
1098            for expression in expressions
1099        ]
1100    )
1101
1102
1103def tuple_(
1104    *expressions: ExpOrStr,
1105    copy: bool = True,
1106    dialect: DialectType = None,
1107    **kwargs: Unpack[ParserNoDialectArgs],
1108) -> Tuple:
1109    """
1110    Returns an tuple.
1111
1112    Examples:
1113        >>> tuple_(1, 'x').sql()
1114        '(1, x)'
1115
1116    Args:
1117        expressions: the expressions to add to the tuple.
1118        copy: whether to copy the argument expressions.
1119        dialect: the source dialect.
1120        kwargs: the kwargs used to instantiate the function of interest.
1121
1122    Returns:
1123        A tuple expression.
1124    """
1125    return Tuple(
1126        expressions=[
1127            maybe_parse(expression, copy=copy, dialect=dialect, **kwargs)
1128            for expression in expressions
1129        ]
1130    )
1131
1132
1133def true() -> Boolean:
1134    """
1135    Returns a true Boolean expression.
1136    """
1137    return Boolean(this=True)
1138
1139
1140def false() -> Boolean:
1141    """
1142    Returns a false Boolean expression.
1143    """
1144    return Boolean(this=False)
1145
1146
1147def null() -> Null:
1148    """
1149    Returns a Null expression.
1150    """
1151    return Null()
1152
1153
1154def apply_index_offset(
1155    this: Expr,
1156    expressions: list[E],
1157    offset: int,
1158    dialect: DialectType = None,
1159) -> list[E]:
1160    if not offset or len(expressions) != 1:
1161        return expressions
1162
1163    expression = expressions[0]
1164
1165    from sqlglot.optimizer.annotate_types import annotate_types
1166    from sqlglot.optimizer.simplify import simplify
1167
1168    if not this.type:
1169        annotate_types(this, dialect=dialect)
1170
1171    if t.cast(DataType, this.type).this not in (
1172        DType.UNKNOWN,
1173        DType.ARRAY,
1174    ):
1175        return expressions
1176
1177    if not expression.type:
1178        annotate_types(expression, dialect=dialect)
1179
1180    if t.cast(DataType, expression.type).this in DataType.INTEGER_TYPES:
1181        logger.info("Applying array index offset (%s)", offset)
1182        expression = simplify(expression + offset)
1183        return [expression]
1184
1185    return expressions
1186
1187
1188NONNULL_CONSTANTS = (
1189    Literal,
1190    Boolean,
1191)
1192
1193CONSTANTS = (
1194    Literal,
1195    Boolean,
1196    Null,
1197)