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    try:
 323        expression = maybe_parse(name, dialect=dialect, into=Identifier)
 324    except (ParseError, TokenError):
 325        expression = to_identifier(name)
 326
 327    return expression
 328
 329
 330INTERVAL_STRING_RE = re.compile(r"\s*(-?[0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]+)\s*")
 331
 332
 333INTERVAL_DAY_TIME_RE = re.compile(
 334    r"\s*-?\s*\d+(?:\.\d+)?\s+(?:-?(?:\d+:)?\d+:\d+(?:\.\d+)?|-?(?:\d+:){1,2}|:)\s*"
 335)
 336
 337
 338def to_interval(interval: str | Expr) -> Interval:
 339    """Builds an interval expression from a string like '1 day' or '5 months'."""
 340    if isinstance(interval, Literal):
 341        if not interval.is_string:
 342            raise ValueError("Invalid interval string.")
 343
 344        interval = interval.this
 345
 346    interval = maybe_parse(f"INTERVAL {interval}")
 347    assert isinstance(interval, Interval)
 348    return interval
 349
 350
 351def to_table(
 352    sql_path: str | Table, dialect: DialectType = None, copy: bool = True, **kwargs: object
 353) -> Table:
 354    """
 355    Create a table expression from a `[catalog].[schema].[table]` sql path. Catalog and schema are optional.
 356    If a table is passed in then that table is returned.
 357
 358    Args:
 359        sql_path: a `[catalog].[schema].[table]` string.
 360        dialect: the source dialect according to which the table name will be parsed.
 361        copy: Whether to copy a table if it is passed in.
 362        kwargs: the kwargs to instantiate the resulting `Table` expression with.
 363
 364    Returns:
 365        A table expression.
 366    """
 367    if isinstance(sql_path, Table):
 368        return maybe_copy(sql_path, copy=copy)
 369
 370    try:
 371        table = maybe_parse(sql_path, into=Table, dialect=dialect)
 372    except ParseError:
 373        catalog, db, this = split_num_words(sql_path, ".", 3)
 374
 375        if not this:
 376            raise
 377
 378        table = table_(this, db=db, catalog=catalog)
 379
 380    return table.set_kwargs(kwargs)
 381
 382
 383def to_column(
 384    sql_path: str | Column,
 385    quoted: bool | None = None,
 386    dialect: DialectType = None,
 387    copy: bool = True,
 388    **kwargs: t.Any,
 389) -> Column | Dot:
 390    """
 391    Create a column from a `[table].[column]` sql path. Table is optional.
 392    If a column is passed in then that column is returned.
 393
 394    Args:
 395        sql_path: a `[table].[column]` string.
 396        quoted: Whether or not to force quote identifiers.
 397        dialect: the source dialect according to which the column name will be parsed.
 398        copy: Whether to copy a column if it is passed in.
 399        kwargs: the kwargs to instantiate the resulting `Column` expression with.
 400
 401    Returns:
 402        A column expression.
 403    """
 404    if isinstance(sql_path, Column):
 405        return maybe_copy(sql_path, copy=copy)
 406
 407    try:
 408        col = maybe_parse(sql_path, into=Column, dialect=dialect)
 409    except ParseError:
 410        return column(*reversed(sql_path.split(".")), quoted=quoted, **kwargs)
 411
 412    for k, v in kwargs.items():
 413        col.set(k, v)
 414
 415    if quoted:
 416        for i in col.find_all(Identifier):
 417            i.set("quoted", True)
 418
 419    return col
 420
 421
 422def subquery(
 423    expression: ExpOrStr,
 424    alias: Identifier | str | None = None,
 425    dialect: DialectType = None,
 426    copy: bool = True,
 427    **opts: Unpack[ParserNoDialectArgs],
 428) -> Select:
 429    """
 430    Build a subquery expression that's selected from.
 431
 432    Example:
 433        >>> subquery('select x from tbl', 'bar').select('x').sql()
 434        'SELECT x FROM (SELECT x FROM tbl) AS bar'
 435
 436    Args:
 437        expression: the SQL code strings to parse.
 438            If an Expr instance is passed, this is used as-is.
 439        alias: the alias name to use.
 440        dialect: the dialect used to parse the input expression.
 441        **opts: other options to use to parse the input expressions.
 442
 443    Returns:
 444        A new Select instance with the subquery expression included.
 445    """
 446    expr = (
 447        maybe_parse(expression, dialect=dialect, **opts).assert_is(Query).subquery(alias, copy=copy)
 448    )
 449    return Select().from_(expr, dialect=dialect, **opts)
 450
 451
 452def cast(
 453    expression: ExpOrStr,
 454    to: DATA_TYPE,
 455    copy: bool = True,
 456    dialect: DialectType = None,
 457    **opts: Unpack[ParserNoDialectArgs],
 458) -> Cast:
 459    """Cast an expression to a data type.
 460
 461    Example:
 462        >>> cast('x + 1', 'int').sql()
 463        'CAST(x + 1 AS INT)'
 464
 465    Args:
 466        expression: The expression to cast.
 467        to: The datatype to cast to.
 468        copy: Whether to copy the supplied expressions.
 469        dialect: The target dialect. This is used to prevent a re-cast in the following scenario:
 470            - The expression to be cast is already a exp.Cast expression
 471            - The existing cast is to a type that is logically equivalent to new type
 472
 473            For example, if :expression='CAST(x as DATETIME)' and :to=Type.TIMESTAMP,
 474            but in the target dialect DATETIME is mapped to TIMESTAMP, then we will NOT return `CAST(x (as DATETIME) as TIMESTAMP)`
 475            and instead just return the original expression `CAST(x as DATETIME)`.
 476
 477            This is to prevent it being output as a double cast `CAST(x (as TIMESTAMP) as TIMESTAMP)` once the DATETIME -> TIMESTAMP
 478            mapping is applied in the target dialect generator.
 479
 480    Returns:
 481        The new Cast instance.
 482    """
 483    expr = maybe_parse(expression, copy=copy, dialect=dialect, **opts)
 484    data_type = DataType.build(to, copy=copy, dialect=dialect, **opts)
 485
 486    # dont re-cast if the expression is already a cast to the correct type
 487    if isinstance(expr, Cast):
 488        from sqlglot.dialects.dialect import Dialect
 489
 490        target_dialect = Dialect.get_or_raise(dialect)
 491        type_mapping = target_dialect.generator_class.TYPE_MAPPING
 492
 493        existing_cast_type: DType = expr.to.this
 494        new_cast_type: DType = data_type.this
 495        types_are_equivalent = type_mapping.get(
 496            existing_cast_type, existing_cast_type.value
 497        ) == type_mapping.get(new_cast_type, new_cast_type.value)
 498
 499        if expr.is_type(data_type) or types_are_equivalent:
 500            return expr
 501
 502    expr = Cast(this=expr, to=data_type)
 503    expr.type = data_type
 504
 505    return expr
 506
 507
 508def table_(
 509    table: Identifier | str,
 510    db: Identifier | str | None = None,
 511    catalog: Identifier | str | None = None,
 512    quoted: bool | None = None,
 513    alias: Identifier | str | None = None,
 514) -> Table:
 515    """Build a Table.
 516
 517    Args:
 518        table: Table name.
 519        db: Database name.
 520        catalog: Catalog name.
 521        quote: Whether to force quotes on the table's identifiers.
 522        alias: Table's alias.
 523
 524    Returns:
 525        The new Table instance.
 526    """
 527    return Table(
 528        this=to_identifier(table, quoted=quoted) if table else None,
 529        db=to_identifier(db, quoted=quoted) if db else None,
 530        catalog=to_identifier(catalog, quoted=quoted) if catalog else None,
 531        alias=TableAlias(this=to_identifier(alias)) if alias else None,
 532    )
 533
 534
 535def values(
 536    values: Iterable[tuple[object, ...] | Tuple],
 537    alias: str | None = None,
 538    columns: Iterable[str] | dict[str, DataType] | None = None,
 539) -> Values:
 540    """Build VALUES statement.
 541
 542    Example:
 543        >>> values([(1, '2')]).sql()
 544        "VALUES (1, '2')"
 545
 546    Args:
 547        values: values statements that will be converted to SQL
 548        alias: optional alias
 549        columns: Optional list of ordered column names or ordered dictionary of column names to types.
 550         If either are provided then an alias is also required.
 551
 552    Returns:
 553        Values: the Values expression object
 554    """
 555    if columns and not alias:
 556        raise ValueError("Alias is required when providing columns")
 557
 558    return Values(
 559        expressions=[convert(tup) for tup in values],
 560        alias=(
 561            TableAlias(this=to_identifier(alias), columns=[to_identifier(x) for x in columns])
 562            if columns
 563            else (TableAlias(this=to_identifier(alias)) if alias else None)
 564        ),
 565    )
 566
 567
 568def var(name: ExpOrStr | None) -> Var:
 569    """Build a SQL variable.
 570
 571    Example:
 572        >>> repr(var('x'))
 573        'Var(this=x)'
 574
 575        >>> repr(var(column('x', table='y')))
 576        'Var(this=x)'
 577
 578    Args:
 579        name: The name of the var or an expression who's name will become the var.
 580
 581    Returns:
 582        The new variable node.
 583    """
 584    if not name:
 585        raise ValueError("Cannot convert empty name into var.")
 586
 587    if isinstance(name, Expr):
 588        name = name.name
 589    return Var(this=name)
 590
 591
 592def rename_table(
 593    old_name: str | Table,
 594    new_name: str | Table,
 595    dialect: DialectType = None,
 596) -> Alter:
 597    """Build ALTER TABLE... RENAME... expression
 598
 599    Args:
 600        old_name: The old name of the table
 601        new_name: The new name of the table
 602        dialect: The dialect to parse the table.
 603
 604    Returns:
 605        Alter table expression
 606    """
 607    old_table = to_table(old_name, dialect=dialect)
 608    new_table = to_table(new_name, dialect=dialect)
 609    return Alter(
 610        this=old_table,
 611        kind="TABLE",
 612        actions=[
 613            AlterRename(this=new_table),
 614        ],
 615    )
 616
 617
 618def rename_column(
 619    table_name: str | Table,
 620    old_column_name: str | Column,
 621    new_column_name: str | Column,
 622    exists: bool | None = None,
 623    dialect: DialectType = None,
 624) -> Alter:
 625    """Build ALTER TABLE... RENAME COLUMN... expression
 626
 627    Args:
 628        table_name: Name of the table
 629        old_column: The old name of the column
 630        new_column: The new name of the column
 631        exists: Whether to add the `IF EXISTS` clause
 632        dialect: The dialect to parse the table/column.
 633
 634    Returns:
 635        Alter table expression
 636    """
 637    table = to_table(table_name, dialect=dialect)
 638    old_column = to_column(old_column_name, dialect=dialect)
 639    new_column = to_column(new_column_name, dialect=dialect)
 640    return Alter(
 641        this=table,
 642        kind="TABLE",
 643        actions=[
 644            RenameColumn(this=old_column, to=new_column, exists=exists),
 645        ],
 646    )
 647
 648
 649def replace_children(
 650    expression: Expr,
 651    fun: t.Callable[Concatenate[Expr, P], object],
 652    *args: P.args,
 653    **kwargs: P.kwargs,
 654) -> None:
 655    """
 656    Replace children of an expression with the result of a lambda fun(child) -> exp.
 657    """
 658    for k, v in tuple(expression.args.items()):
 659        is_list_arg = type(v) is list
 660
 661        child_nodes = v if is_list_arg else [v]
 662        new_child_nodes = []
 663
 664        for cn in child_nodes:
 665            if isinstance(cn, Expr):
 666                for child_node in ensure_collection(fun(cn, *args, **kwargs)):
 667                    new_child_nodes.append(child_node)
 668            else:
 669                new_child_nodes.append(cn)
 670
 671        if is_list_arg:
 672            expression.set(k, new_child_nodes)
 673        else:
 674            expression.set(k, seq_get(new_child_nodes, 0))
 675
 676
 677def replace_tree(
 678    expression: Expr,
 679    fun: t.Callable[[Expr], Expr],
 680    prune: t.Callable[[Expr], bool] | None = None,
 681) -> Expr:
 682    """
 683    Replace an entire tree with the result of function calls on each node.
 684
 685    This will be traversed in reverse dfs, so leaves first.
 686    If new nodes are created as a result of function calls, they will also be traversed.
 687    """
 688    stack = list(expression.dfs(prune=prune))
 689
 690    while stack:
 691        node = stack.pop()
 692        new_node = fun(node)
 693
 694        if new_node is not node:
 695            node.replace(new_node)
 696
 697            if isinstance(new_node, Expr):
 698                stack.append(new_node)
 699
 700    return new_node
 701
 702
 703def find_tables(expression: Expr) -> set[Table]:
 704    """
 705    Find all tables referenced in a query.
 706
 707    Args:
 708        expressions: The query to find the tables in.
 709
 710    Returns:
 711        A set of all the tables.
 712    """
 713    from sqlglot.optimizer.scope import traverse_scope
 714
 715    return {
 716        table
 717        for scope in traverse_scope(expression)
 718        for table in scope.tables
 719        if isinstance(table, Table) and table.name and table.name not in scope.cte_sources
 720    }
 721
 722
 723def column_table_names(expression: Expr, exclude: str = "") -> set[str]:
 724    """
 725    Return all table names referenced through columns in an expression.
 726
 727    Example:
 728        >>> import sqlglot
 729        >>> sorted(column_table_names(sqlglot.parse_one("a.b AND c.d AND c.e")))
 730        ['a', 'c']
 731
 732    Args:
 733        expression: expression to find table names.
 734        exclude: a table name to exclude
 735
 736    Returns:
 737        A list of unique names.
 738    """
 739    return {
 740        table
 741        for table in (column.table for column in expression.find_all(Column))
 742        if table and table != exclude
 743    }
 744
 745
 746def table_name(table: Table | str, dialect: DialectType = None, identify: bool = False) -> str:
 747    """Get the full name of a table as a string.
 748
 749    Args:
 750        table: Table expression node or string.
 751        dialect: The dialect to generate the table name for.
 752        identify: Determines when an identifier should be quoted. Possible values are:
 753            False (default): Never quote, except in cases where it's mandatory by the dialect.
 754            True: Always quote.
 755
 756    Examples:
 757        >>> from sqlglot import exp, parse_one
 758        >>> table_name(parse_one("select * from a.b.c").find(exp.Table))
 759        'a.b.c'
 760
 761    Returns:
 762        The table name.
 763    """
 764
 765    expr = maybe_parse(table, into=Table, dialect=dialect)
 766
 767    if not expr:
 768        raise ValueError(f"Cannot parse {table}")
 769
 770    return ".".join(
 771        (
 772            part.sql(dialect=dialect, identify=True, copy=False, comments=False)
 773            if identify or not SAFE_IDENTIFIER_RE.match(part.name)
 774            else part.name
 775        )
 776        for part in expr.parts
 777    )
 778
 779
 780def normalize_table_name(table: str | Table, dialect: DialectType = None, copy: bool = True) -> str:
 781    """Returns a case normalized table name without quotes.
 782
 783    Args:
 784        table: the table to normalize
 785        dialect: the dialect to use for normalization rules
 786        copy: whether to copy the expression.
 787
 788    Examples:
 789        >>> normalize_table_name("`A-B`.c", dialect="bigquery")
 790        'A-B.c'
 791    """
 792    from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 793
 794    return ".".join(
 795        p.name
 796        for p in normalize_identifiers(
 797            to_table(table, dialect=dialect, copy=copy), dialect=dialect
 798        ).parts
 799    )
 800
 801
 802def replace_tables(
 803    expression: E, mapping: dict[str, str], dialect: DialectType = None, copy: bool = True
 804) -> E:
 805    """Replace all tables in expression according to the mapping.
 806
 807    Args:
 808        expression: expression node to be transformed and replaced.
 809        mapping: mapping of table names.
 810        dialect: the dialect of the mapping table
 811        copy: whether to copy the expression.
 812
 813    Examples:
 814        >>> from sqlglot import exp, parse_one
 815        >>> replace_tables(parse_one("select * from a.b"), {"a.b": "c"}).sql()
 816        'SELECT * FROM c /* a.b */'
 817
 818    Returns:
 819        The mapped expression.
 820    """
 821
 822    mapping = {normalize_table_name(k, dialect=dialect): v for k, v in mapping.items()}
 823
 824    def _replace_tables(node: Expr) -> Expr:
 825        if isinstance(node, Table) and node.meta.get("replace") is not False:
 826            original = normalize_table_name(node, dialect=dialect)
 827            new_name = mapping.get(original)
 828
 829            if new_name:
 830                table = to_table(
 831                    new_name,
 832                    **{k: v for k, v in node.args.items() if k not in TABLE_PARTS},
 833                    dialect=dialect,
 834                )
 835                table.add_comments([original])
 836                return table
 837        return node
 838
 839    return expression.transform(_replace_tables, copy=copy)  # type: ignore
 840
 841
 842def replace_placeholders(expression: Expr, *args: object, **kwargs: t.Any) -> Expr:
 843    """Replace placeholders in an expression.
 844
 845    Args:
 846        expression: expression node to be transformed and replaced.
 847        args: positional names that will substitute unnamed placeholders in the given order.
 848        kwargs: keyword arguments that will substitute named placeholders.
 849
 850    Examples:
 851        >>> from sqlglot import exp, parse_one
 852        >>> replace_placeholders(
 853        ...     parse_one("select * from :tbl where ? = ?"),
 854        ...     exp.to_identifier("str_col"), "b", tbl=exp.to_identifier("foo")
 855        ... ).sql()
 856        "SELECT * FROM foo WHERE str_col = 'b'"
 857
 858    Returns:
 859        The mapped expression.
 860    """
 861
 862    def _replace_placeholders(node: Expr, args: Iterator[object], **kwargs: object) -> Expr:
 863        if isinstance(node, Placeholder):
 864            if node.this:
 865                new_name = kwargs.get(node.this)
 866                if new_name is not None:
 867                    return convert(new_name)
 868            else:
 869                try:
 870                    return convert(next(args))
 871                except StopIteration:
 872                    pass
 873        return node
 874
 875    return expression.transform(_replace_placeholders, iter(args), **kwargs)
 876
 877
 878def expand(
 879    expression: Expr,
 880    sources: dict[str, Query | t.Callable[[], Query]],
 881    dialect: DialectType = None,
 882    copy: bool = True,
 883) -> Expr:
 884    """Transforms an expression by expanding all referenced sources into subqueries.
 885
 886    Examples:
 887        >>> from sqlglot import parse_one
 888        >>> expand(parse_one("select * from x AS z"), {"x": parse_one("select * from y")}).sql()
 889        'SELECT * FROM (SELECT * FROM y) AS z /* source: x */'
 890
 891        >>> expand(parse_one("select * from x AS z"), {"x": parse_one("select * from y"), "y": parse_one("select * from z")}).sql()
 892        'SELECT * FROM (SELECT * FROM (SELECT * FROM z) AS y /* source: y */) AS z /* source: x */'
 893
 894    Args:
 895        expression: The expression to expand.
 896        sources: A dict of name to query or a callable that provides a query on demand.
 897        dialect: The dialect of the sources dict or the callable.
 898        copy: Whether to copy the expression during transformation. Defaults to True.
 899
 900    Returns:
 901        The transformed expression.
 902    """
 903    normalized_sources = {normalize_table_name(k, dialect=dialect): v for k, v in sources.items()}
 904
 905    def _expand(node: Expr):
 906        if isinstance(node, Table):
 907            name = normalize_table_name(node, dialect=dialect)
 908            source = normalized_sources.get(name)
 909
 910            if source:
 911                # Create a subquery with the same alias (or table name if no alias)
 912                parsed_source = source() if callable(source) else source
 913                subquery = parsed_source.subquery(node.alias or name)
 914                subquery.comments = [f"source: {name}"]
 915
 916                # Continue expanding within the subquery
 917                return subquery.transform(_expand, copy=False)
 918
 919        return node
 920
 921    return expression.transform(_expand, copy=copy)
 922
 923
 924def func(
 925    name: str, *args: t.Any, copy: bool = True, dialect: DialectType = None, **kwargs: t.Any
 926) -> Func:
 927    """
 928    Returns a Func expression.
 929
 930    Examples:
 931        >>> func("abs", 5).sql()
 932        'ABS(5)'
 933
 934        >>> func("cast", this=5, to=DataType.build("DOUBLE")).sql()
 935        'CAST(5 AS DOUBLE)'
 936
 937    Args:
 938        name: the name of the function to build.
 939        args: the args used to instantiate the function of interest.
 940        copy: whether to copy the argument expressions.
 941        dialect: the source dialect.
 942        kwargs: the kwargs used to instantiate the function of interest.
 943
 944    Note:
 945        The arguments `args` and `kwargs` are mutually exclusive.
 946
 947    Returns:
 948        An instance of the function of interest, or an anonymous function, if `name` doesn't
 949        correspond to an existing `sqlglot.expressions.Func` class.
 950    """
 951    if args and kwargs:
 952        raise ValueError("Can't use both args and kwargs to instantiate a function.")
 953
 954    from sqlglot.dialects.dialect import Dialect
 955
 956    dialect = Dialect.get_or_raise(dialect)
 957
 958    converted: list[Expr] = [maybe_parse(arg, dialect=dialect, copy=copy) for arg in args]
 959    kwargs = {key: maybe_parse(value, dialect=dialect, copy=copy) for key, value in kwargs.items()}
 960
 961    constructor = dialect.parser_class.FUNCTIONS.get(name.upper())
 962    if constructor:
 963        if converted:
 964            try:
 965                function = constructor(converted)
 966            except TypeError:
 967                function = constructor(converted, dialect=dialect)
 968        elif constructor.__name__ == "from_arg_list":
 969            function = constructor.__self__(**kwargs)  # type: ignore
 970        else:
 971            from sqlglot.expressions import FUNCTION_BY_NAME as _FUNCTION_BY_NAME
 972
 973            constructor = _FUNCTION_BY_NAME.get(name.upper())
 974            if constructor:
 975                function = constructor(**kwargs)
 976            else:
 977                raise ValueError(
 978                    f"Unable to convert '{name}' into a Func. Either manually construct "
 979                    "the Func expression of interest or parse the function call."
 980                )
 981    else:
 982        kwargs = kwargs or {"expressions": converted}
 983        function = Anonymous(this=name, **kwargs)
 984
 985    for error_message in function.error_messages(converted):
 986        raise ValueError(error_message)
 987
 988    return function
 989
 990
 991def case(
 992    expression: ExpOrStr | None = None,
 993    copy: bool = True,
 994    **opts: Unpack[ParserArgs],
 995) -> Case:
 996    """
 997    Initialize a CASE statement.
 998
 999    Example:
1000        case().when("a = 1", "foo").else_("bar")
1001
1002    Args:
1003        expression: Optionally, the input expression (not all dialects support this)
1004        copy: whether to copy the argument expressions.
1005        **opts: Extra keyword arguments for parsing `expression`
1006    """
1007    if expression is not None:
1008        this = maybe_parse(expression, copy=copy, **opts)
1009    else:
1010        this = None
1011    return Case(this=this, ifs=[])
1012
1013
1014def array(
1015    *expressions: ExpOrStr,
1016    copy: bool = True,
1017    dialect: DialectType = None,
1018    **kwargs: Unpack[ParserNoDialectArgs],
1019) -> Array:
1020    """
1021    Returns an array.
1022
1023    Examples:
1024        >>> array(1, 'x').sql()
1025        'ARRAY(1, x)'
1026
1027    Args:
1028        expressions: the expressions to add to the array.
1029        copy: whether to copy the argument expressions.
1030        dialect: the source dialect.
1031        kwargs: the kwargs used to instantiate the function of interest.
1032
1033    Returns:
1034        An array expression.
1035    """
1036    return Array(
1037        expressions=[
1038            maybe_parse(expression, copy=copy, dialect=dialect, **kwargs)
1039            for expression in expressions
1040        ]
1041    )
1042
1043
1044def tuple_(
1045    *expressions: ExpOrStr,
1046    copy: bool = True,
1047    dialect: DialectType = None,
1048    **kwargs: Unpack[ParserNoDialectArgs],
1049) -> Tuple:
1050    """
1051    Returns an tuple.
1052
1053    Examples:
1054        >>> tuple_(1, 'x').sql()
1055        '(1, x)'
1056
1057    Args:
1058        expressions: the expressions to add to the tuple.
1059        copy: whether to copy the argument expressions.
1060        dialect: the source dialect.
1061        kwargs: the kwargs used to instantiate the function of interest.
1062
1063    Returns:
1064        A tuple expression.
1065    """
1066    return Tuple(
1067        expressions=[
1068            maybe_parse(expression, copy=copy, dialect=dialect, **kwargs)
1069            for expression in expressions
1070        ]
1071    )
1072
1073
1074def true() -> Boolean:
1075    """
1076    Returns a true Boolean expression.
1077    """
1078    return Boolean(this=True)
1079
1080
1081def false() -> Boolean:
1082    """
1083    Returns a false Boolean expression.
1084    """
1085    return Boolean(this=False)
1086
1087
1088def null() -> Null:
1089    """
1090    Returns a Null expression.
1091    """
1092    return Null()
1093
1094
1095def apply_index_offset(
1096    this: Expr,
1097    expressions: list[E],
1098    offset: int,
1099    dialect: DialectType = None,
1100) -> list[E]:
1101    if not offset or len(expressions) != 1:
1102        return expressions
1103
1104    expression = expressions[0]
1105
1106    from sqlglot.optimizer.annotate_types import annotate_types
1107    from sqlglot.optimizer.simplify import simplify
1108
1109    if not this.type:
1110        annotate_types(this, dialect=dialect)
1111
1112    if t.cast(DataType, this.type).this not in (
1113        DType.UNKNOWN,
1114        DType.ARRAY,
1115    ):
1116        return expressions
1117
1118    if not expression.type:
1119        annotate_types(expression, dialect=dialect)
1120
1121    if t.cast(DataType, expression.type).this in DataType.INTEGER_TYPES:
1122        logger.info("Applying array index offset (%s)", offset)
1123        expression = simplify(expression + offset)
1124        return [expression]
1125
1126    return expressions
1127
1128
1129NONNULL_CONSTANTS = (
1130    Literal,
1131    Boolean,
1132)
1133
1134CONSTANTS = (
1135    Literal,
1136    Boolean,
1137    Null,
1138)