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