sqlmesh.core.dialect

   1from __future__ import annotations
   2
   3import functools
   4import logging
   5import re
   6import sys
   7import typing as t
   8from contextlib import contextmanager
   9from difflib import unified_diff
  10from enum import Enum, auto
  11from functools import lru_cache
  12
  13from sqlglot import Dialect, Generator, ParseError, Parser, Tokenizer, TokenType, exp
  14from sqlglot.dialects.dialect import DialectType
  15from sqlglot.dialects import DuckDB, Snowflake, TSQL
  16import sqlglot.dialects.athena as athena
  17import sqlglot.generators.athena as athena_generators
  18from sqlglot.parsers.athena import AthenaTrinoParser
  19from sqlglot.helper import seq_get
  20from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
  21from sqlglot.optimizer.qualify_columns import quote_identifiers
  22from sqlglot.optimizer.qualify_tables import qualify_tables
  23from sqlglot.optimizer.scope import traverse_scope
  24from sqlglot.schema import MappingSchema
  25from sqlglot.tokens import Token
  26
  27from sqlmesh.core.constants import LIQUID_CLUSTERING_KEYWORDS, MAX_MODEL_DEFINITION_SIZE
  28from sqlmesh.utils import get_source_columns_to_types
  29from sqlmesh.utils.errors import SQLMeshError, ConfigError
  30from sqlmesh.utils.pandas import columns_to_types_from_df
  31
  32if t.TYPE_CHECKING:
  33    import pandas as pd
  34
  35    from sqlglot._typing import E
  36
  37
  38SQLMESH_MACRO_PREFIX = "@"
  39
  40TABLES_META = "sqlmesh.tables"
  41
  42logger = logging.getLogger(__name__)
  43
  44
  45class Model(exp.Expression):
  46    arg_types = {"expressions": True}
  47
  48
  49class Audit(exp.Expression):
  50    arg_types = {"expressions": True}
  51
  52
  53class Metric(exp.Expression):
  54    arg_types = {"expressions": True}
  55
  56
  57class Jinja(exp.Expression, exp.Func):
  58    arg_types = {"this": True}
  59
  60
  61class JinjaQuery(Jinja):
  62    pass
  63
  64
  65class JinjaStatement(Jinja):
  66    pass
  67
  68
  69class VirtualUpdateStatement(exp.Expression):
  70    arg_types = {"expressions": True}
  71
  72
  73class ModelKind(exp.Expression):
  74    arg_types = {"this": True, "expressions": False}
  75
  76
  77class MacroVar(exp.Var):
  78    pass
  79
  80
  81class MacroFunc(exp.Expression, exp.Func):
  82    @property
  83    def name(self) -> str:
  84        return self.this.name
  85
  86
  87class MacroDef(MacroFunc):
  88    arg_types = {"this": True, "expression": True}
  89
  90
  91class MacroSQL(MacroFunc):
  92    arg_types = {"this": True, "into": False}
  93
  94
  95class MacroStrReplace(MacroFunc):
  96    pass
  97
  98
  99class PythonCode(exp.Expression):
 100    arg_types = {"expressions": True}
 101
 102
 103class DColonCast(exp.Cast):
 104    pass
 105
 106
 107class MetricAgg(exp.Expression, exp.AggFunc):
 108    """Used for computing metrics."""
 109
 110    arg_types = {"this": True}
 111
 112    @property
 113    def output_name(self) -> str:
 114        return self.this.name
 115
 116
 117class StagedFilePath(exp.Expression):
 118    """Represents paths to "staged files" in Snowflake."""
 119
 120    arg_types = exp.Table.arg_types.copy()
 121
 122
 123def _parse_statement(self: Parser) -> t.Optional[exp.Expr]:
 124    if self._curr is None:
 125        return None
 126
 127    parser = PARSERS.get(self._curr.text.upper())
 128    error_msg = None
 129
 130    if parser:
 131        # Capture any available description in the form of a comment
 132        comments = self._curr.comments
 133
 134        index = self._index
 135        try:
 136            self._advance()
 137            meta = self._parse_wrapped(lambda: t.cast(t.Callable, parser)(self))
 138        except ParseError as parse_error:
 139            error_msg = parse_error.args[0]
 140            self._retreat(index)
 141
 142        # Only return the DDL expression if we actually managed to parse one. This is
 143        # done in order to allow parsing standalone identifiers / function calls like
 144        # "metric", or "model(1, 2, 3)", which collide with SQLMesh's DDL syntax.
 145        if self._index != index:
 146            meta.comments = comments
 147            return meta
 148
 149    try:
 150        return self.__parse_statement()  # type: ignore
 151    except ParseError:
 152        if error_msg:
 153            raise ParseError(error_msg)
 154        raise
 155
 156
 157def _parse_lambda(self: Parser, alias: bool = False) -> t.Optional[exp.Expr]:
 158    node = self.__parse_lambda(alias=alias)  # type: ignore
 159    if isinstance(node, exp.Lambda):
 160        node.set("this", self._parse_alias(node.this))
 161    return node
 162
 163
 164def _parse_id_var(
 165    self: Parser,
 166    any_token: bool = True,
 167    tokens: t.Optional[t.Collection[TokenType]] = None,
 168) -> t.Optional[exp.Expr]:
 169    if self._prev and self._prev.text == SQLMESH_MACRO_PREFIX and self._match(TokenType.L_BRACE):
 170        identifier = self.__parse_id_var(any_token=any_token, tokens=tokens)  # type: ignore
 171        if not self._match(TokenType.R_BRACE):
 172            self.raise_error("Expecting }")
 173        identifier.args["this"] = f"@{{{identifier.name}}}"
 174    else:
 175        identifier = self.__parse_id_var(any_token=any_token, tokens=tokens)  # type: ignore
 176
 177    while (
 178        identifier
 179        and not identifier.args.get("quoted")
 180        and self._is_connected()
 181        and (
 182            self._match_texts(("{", SQLMESH_MACRO_PREFIX))
 183            or self._curr.token_type not in self.RESERVED_TOKENS
 184        )
 185    ):
 186        this = identifier.name
 187        brace = False
 188
 189        if self._prev.text == "{":
 190            this += "{"
 191            brace = True
 192        else:
 193            if self._prev.text == SQLMESH_MACRO_PREFIX:
 194                this += "@"
 195            if self._match(TokenType.L_BRACE):
 196                this += "{"
 197                brace = True
 198
 199        next_id = self._parse_id_var(any_token=False)
 200
 201        if next_id:
 202            this += next_id.name
 203        else:
 204            return identifier
 205
 206        if brace:
 207            if self._match(TokenType.R_BRACE):
 208                this += "}"
 209            else:
 210                self.raise_error("Expecting }")
 211
 212        identifier = self.expression(exp.Identifier(this=this, quoted=identifier.quoted))
 213
 214    return identifier
 215
 216
 217def _parse_macro(self: Parser, keyword_macro: str = "") -> t.Optional[exp.Expr]:
 218    if self._prev.text != SQLMESH_MACRO_PREFIX:
 219        return self._parse_parameter()
 220
 221    comments = self._prev.comments
 222    index = self._index
 223    field = self._parse_primary() or self._parse_function(functions={}) or self._parse_id_var()
 224
 225    def _build_macro(field: t.Optional[exp.Expr]) -> t.Optional[exp.Expr]:
 226        if isinstance(field, exp.Func):
 227            macro_name = field.name.upper()
 228            if macro_name != keyword_macro and macro_name in KEYWORD_MACROS:
 229                self._retreat(index)
 230                return None
 231
 232            if isinstance(field, exp.Anonymous):
 233                if macro_name == "DEF":
 234                    return self.expression(
 235                        MacroDef(
 236                            this=field.expressions[0],
 237                            expression=field.expressions[1],
 238                        ),
 239                        comments=comments,
 240                    )
 241                if macro_name == "SQL":
 242                    into = field.expressions[1].this.lower() if len(field.expressions) > 1 else None
 243                    return self.expression(
 244                        MacroSQL(this=field.expressions[0], into=into), comments=comments
 245                    )
 246            else:
 247                field = self.expression(
 248                    exp.Anonymous(
 249                        this=field.sql_name(),
 250                        expressions=list(field.args.values()),
 251                    ),
 252                    comments=comments,
 253                )
 254
 255            return self.expression(MacroFunc(this=field), comments=comments)
 256
 257        if field is None:
 258            return None
 259
 260        if field.is_string or (isinstance(field, exp.Identifier) and field.quoted):
 261            return self.expression(
 262                MacroStrReplace(this=exp.Literal.string(field.this)), comments=comments
 263            )
 264
 265        if "@" in field.this:
 266            return field  # type: ignore[return-value]
 267        return self.expression(MacroVar(this=field.this), comments=comments)
 268
 269    if isinstance(field, (exp.Window, exp.IgnoreNulls, exp.RespectNulls)):
 270        field.set("this", _build_macro(field.this))
 271    else:
 272        field = _build_macro(field)
 273
 274    return field
 275
 276
 277KEYWORD_MACROS = {"WITH", "JOIN", "WHERE", "GROUP_BY", "HAVING", "ORDER_BY", "LIMIT"}
 278
 279
 280def _parse_matching_macro(self: Parser, name: str) -> t.Optional[exp.Expr]:
 281    if not self._match_pair(TokenType.PARAMETER, TokenType.VAR, advance=False) or (
 282        self._next and self._next.text.upper() != name.upper()
 283    ):
 284        return None
 285
 286    self._advance()
 287    return _parse_macro(self, keyword_macro=name)
 288
 289
 290def _parse_body_macro(self: Parser) -> t.Tuple[str, t.Optional[exp.Expr]]:
 291    name = self._next and self._next.text.upper()
 292
 293    if name == "JOIN":
 294        return ("joins", self._parse_join())
 295    if name == "WHERE":
 296        return ("where", self._parse_where())
 297    if name == "GROUP_BY":
 298        return ("group", self._parse_group())
 299    if name == "HAVING":
 300        return ("having", self._parse_having())
 301    if name == "ORDER_BY":
 302        return ("order", self._parse_order())
 303    if name == "LIMIT":
 304        return ("limit", self._parse_limit())
 305    return ("", None)
 306
 307
 308def _parse_with(self: Parser, skip_with_token: bool = False) -> t.Optional[exp.Expr]:
 309    macro = _parse_matching_macro(self, "WITH")
 310    if not macro:
 311        return self.__parse_with(skip_with_token=skip_with_token)  # type: ignore
 312
 313    macro.this.append("expressions", self.__parse_with(skip_with_token=True))  # type: ignore
 314    return macro
 315
 316
 317def _parse_join(
 318    self: Parser, skip_join_token: bool = False, parse_bracket: bool = False
 319) -> t.Optional[exp.Expr]:
 320    index = self._index
 321    method, side, kind = self._parse_join_parts()
 322    macro = _parse_matching_macro(self, "JOIN")
 323    if not macro:
 324        self._retreat(index)
 325        return self.__parse_join(skip_join_token=skip_join_token, parse_bracket=parse_bracket)  # type: ignore
 326
 327    join = self.__parse_join(skip_join_token=True)  # type: ignore
 328    if method:
 329        join.set("method", method.text)
 330    if side:
 331        join.set("side", side.text)
 332    if kind:
 333        join.set("kind", kind.text)
 334
 335    macro.this.append("expressions", join)
 336    return macro
 337
 338
 339def _warn_unsupported(self: Parser) -> None:
 340    from sqlmesh.core.console import get_console
 341
 342    sql = self._find_sql(self._tokens[0], self._tokens[-1])[: self.error_message_context]
 343
 344    get_console().log_warning(
 345        f"'{sql}' could not be semantically understood as it contains unsupported syntax, SQLMesh will treat the command as is. Note that any references to the model's "
 346        "underlying physical table can't be resolved in this case, consider using Jinja as explained here https://sqlmesh.readthedocs.io/en/stable/concepts/macros/macro_variables/#audit-only-variables"
 347    )
 348
 349
 350def _parse_select(
 351    self: Parser,
 352    nested: bool = False,
 353    table: bool = False,
 354    parse_subquery_alias: bool = True,
 355    parse_set_operation: bool = True,
 356    consume_pipe: bool = True,
 357    from_: t.Optional[exp.From] = None,
 358) -> t.Optional[exp.Expr]:
 359    select = self.__parse_select(  # type: ignore
 360        nested=nested,
 361        table=table,
 362        parse_subquery_alias=parse_subquery_alias,
 363        parse_set_operation=parse_set_operation,
 364        consume_pipe=consume_pipe,
 365        from_=from_,
 366    )
 367
 368    if (
 369        not select
 370        and not parse_set_operation
 371        and self._match_pair(TokenType.PARAMETER, TokenType.VAR, advance=False)
 372    ):
 373        self._advance()
 374        return _parse_macro(self)
 375
 376    return select
 377
 378
 379def _parse_where(self: Parser, skip_where_token: bool = False) -> t.Optional[exp.Expr]:
 380    macro = _parse_matching_macro(self, "WHERE")
 381    if not macro:
 382        return self.__parse_where(skip_where_token=skip_where_token)  # type: ignore
 383
 384    macro.this.append("expressions", self.__parse_where(skip_where_token=True))  # type: ignore
 385    return macro
 386
 387
 388def _parse_group(self: Parser, skip_group_by_token: bool = False) -> t.Optional[exp.Expr]:
 389    macro = _parse_matching_macro(self, "GROUP_BY")
 390    if not macro:
 391        return self.__parse_group(skip_group_by_token=skip_group_by_token)  # type: ignore
 392
 393    macro.this.append("expressions", self.__parse_group(skip_group_by_token=True))  # type: ignore
 394    return macro
 395
 396
 397def _parse_having(self: Parser, skip_having_token: bool = False) -> t.Optional[exp.Expr]:
 398    macro = _parse_matching_macro(self, "HAVING")
 399    if not macro:
 400        return self.__parse_having(skip_having_token=skip_having_token)  # type: ignore
 401
 402    macro.this.append("expressions", self.__parse_having(skip_having_token=True))  # type: ignore
 403    return macro
 404
 405
 406def _parse_order(
 407    self: Parser, this: t.Optional[exp.Expr] = None, skip_order_token: bool = False
 408) -> t.Optional[exp.Expr]:
 409    macro = _parse_matching_macro(self, "ORDER_BY")
 410    if not macro:
 411        return self.__parse_order(this, skip_order_token=skip_order_token)  # type: ignore
 412
 413    macro.this.append("expressions", self.__parse_order(this, skip_order_token=True))  # type: ignore
 414    return macro
 415
 416
 417def _parse_limit(
 418    self: Parser,
 419    this: t.Optional[exp.Expr] = None,
 420    top: bool = False,
 421    skip_limit_token: bool = False,
 422) -> t.Optional[exp.Expr]:
 423    macro = _parse_matching_macro(self, "TOP" if top else "LIMIT")
 424    if not macro:
 425        return self.__parse_limit(this, top=top, skip_limit_token=skip_limit_token)  # type: ignore
 426
 427    macro.this.append("expressions", self.__parse_limit(this, top=top, skip_limit_token=True))  # type: ignore
 428    return macro
 429
 430
 431def _parse_value(self: Parser, values: bool = True) -> t.Optional[exp.Expr]:
 432    wrapped = self._match(TokenType.L_PAREN, advance=False)
 433
 434    # The base _parse_value method always constructs a Tuple instance. This is problematic when
 435    # generating values with a macro function, because it's impossible to tell whether the user's
 436    # intention was to construct a row or a column with the VALUES expression. To avoid this, we
 437    # amend the AST such that the Tuple is replaced by the macro function call itself.
 438    expr = self.__parse_value()  # type: ignore
 439    if expr and not wrapped and isinstance(seq_get(expr.expressions, 0), MacroFunc):
 440        return expr.expressions[0]
 441
 442    return expr
 443
 444
 445def _parse_macro_or_clause(self: Parser, parser: t.Callable) -> t.Optional[exp.Expr]:
 446    return _parse_macro(self) if self._match(TokenType.PARAMETER) else parser()
 447
 448
 449def _parse_props(self: Parser) -> t.Optional[exp.Expr]:
 450    key = self._parse_id_var(any_token=True)
 451    if not key:
 452        return None
 453
 454    name = key.name.lower()
 455    if name == "time_data_type":
 456        # TODO: if we make *_data_type a convention to parse things into exp.DataType, we could make this more generic
 457        value = self._parse_types(schema=True)
 458    elif name == "when_matched":
 459        # Parentheses around the WHEN clauses can be used to disambiguate them from other properties
 460        value = self._parse_wrapped(
 461            lambda: _parse_macro_or_clause(self, self._parse_when_matched),
 462            optional=True,
 463        )
 464    elif name == "merge_filter":
 465        value = self._parse_conjunction()
 466    elif self._match(TokenType.L_PAREN):
 467        value = self.expression(exp.Tuple(expressions=self._parse_csv(self._parse_equality)))
 468        self._match_r_paren()
 469    else:
 470        value = self._parse_bracket(self._parse_field(any_token=True))
 471
 472    if name == "path" and value:
 473        # Make sure if we get a windows path that it is converted to posix
 474        value = exp.Literal.string(value.this.replace("\\", "/"))  # type: ignore
 475
 476    return self.expression(exp.Property(this=name, value=value))
 477
 478
 479def _parse_types(
 480    self: Parser,
 481    check_func: bool = False,
 482    schema: bool = False,
 483    allow_identifiers: bool = True,
 484    with_collation: bool = False,
 485) -> t.Optional[exp.Expr]:
 486    start = self._curr
 487    parsed_type = self.__parse_types(  # type: ignore
 488        check_func=check_func,
 489        schema=schema,
 490        allow_identifiers=allow_identifiers,
 491        with_collation=with_collation,
 492    )
 493
 494    if schema and parsed_type:
 495        parsed_type.meta["sql"] = self._find_sql(start, self._prev)
 496
 497    return parsed_type
 498
 499
 500# Only needed for Snowflake: its "staged file" syntax (@<path>) clashes with our macro
 501# var syntax. By converting the Var representation to a MacroVar, we should be able to
 502# handle both use cases: if there's no value in the MacroEvaluator's context for that
 503# MacroVar, it'll render into @<path>, so it won't break staged file path references.
 504#
 505# See: https://docs.snowflake.com/en/user-guide/querying-stage
 506def _parse_table_parts(
 507    self: Parser,
 508    schema: bool = False,
 509    is_db_reference: bool = False,
 510    wildcard: bool = False,
 511    fast: bool = False,
 512) -> exp.Table | StagedFilePath:
 513    index = self._index
 514    table = self.__parse_table_parts(  # type: ignore
 515        schema=schema, is_db_reference=is_db_reference, wildcard=wildcard, fast=fast
 516    )
 517
 518    if table is None:
 519        return table  # type: ignore[return-value]
 520
 521    table_arg = table.this
 522    name = table_arg.name if isinstance(table_arg, exp.Var) else ""
 523
 524    if name.startswith(SQLMESH_MACRO_PREFIX):
 525        # In these cases, we don't want to produce a `StagedFilePath` node:
 526        #
 527        # - @'...' needs to parsed as a string template
 528        # - @{foo}.bar needs to be parsed as a table with a macro var part
 529        # - @name(arg1 [, arg2 ...]) needs to be parsed as a macro function call
 530        #
 531        # These cases can unambiguously be parsed using the base `_parse_table_parts`, as there
 532        # is no overlap with staged files https://docs.snowflake.com/en/user-guide/querying-stage
 533        if (
 534            self._prev.token_type == TokenType.STRING
 535            or "{" in name
 536            or (
 537                self._curr
 538                and self._prev.token_type in (TokenType.L_PAREN, TokenType.R_PAREN)
 539                and self._curr.text.upper() not in ("FILE_FORMAT", "PATTERN")
 540                and not (table.args.get("format") or table.args.get("pattern"))
 541            )
 542        ):
 543            self._retreat(index)
 544            return Parser._parse_table_parts(
 545                self, schema=schema, is_db_reference=is_db_reference, fast=fast
 546            )  # type: ignore[return-value]
 547
 548        table_arg.replace(MacroVar(this=name[1:]))
 549        return StagedFilePath(**table.args)
 550
 551    return table
 552
 553
 554def _parse_if(self: Parser) -> t.Optional[exp.Expr]:
 555    # If we fail to parse an IF function with expressions as arguments, we then try
 556    # to parse a statement / command to support the macro @IF(condition, statement)
 557    index = self._index
 558    try:
 559        if self.dialect == "tsql":
 560            if not (self._index >= 2 and self._tokens[self._index - 2].text == "@"):
 561                return self.__parse_if()  # type: ignore
 562            return Parser.__parse_if(self)  # type: ignore
 563        return self.__parse_if()  # type: ignore
 564    except ParseError:
 565        self._retreat(index)
 566        self._match_l_paren()
 567
 568        cond = self._parse_conjunction()
 569        self._match(TokenType.COMMA)
 570
 571        # Try to parse a known statement, otherwise fall back to parsing a command
 572        # Since the trailing `)` token is not expected by the statement parsers, we
 573        # remove it from the token stream before trying to parse the statement.
 574        last_token = self._tokens[-1]
 575        if last_token.token_type == TokenType.R_PAREN:
 576            self._tokens[-2].comments.extend(last_token.comments)
 577            self._tokens.pop()
 578            if hasattr(self, "_tokens_size"):
 579                # keep _tokens_size in sync sqlglot 30.0.3 caches len(_tokens)
 580                # _advance() tries to read tokens[index + 1] past the new end
 581                self._tokens_size -= 1
 582        else:
 583            self.raise_error("Expecting )")
 584
 585        index = self._index
 586        stmt = self._parse_statement()
 587        if self._curr:
 588            self._retreat(index)
 589            stmt = self._parse_as_command(self._tokens[index])
 590
 591        return exp.Anonymous(this="IF", expressions=[cond, stmt])
 592
 593
 594def _create_parser(expression_type: t.Type[exp.Expr], table_keys: t.List[str]) -> t.Callable:
 595    def parse(self: Parser) -> t.Optional[exp.Expr]:
 596        from sqlmesh.core.model.kind import ModelKindName
 597
 598        expressions: t.List[exp.Expr] = []
 599
 600        while True:
 601            prev_property = seq_get(expressions, -1)
 602            if not self._match(TokenType.COMMA, expression=prev_property) and expressions:
 603                break
 604
 605            key_expression = self._parse_id_var(any_token=True)
 606            if not key_expression:
 607                break
 608
 609            # This allows macro functions that programmaticaly generate the property key-value pair
 610            if isinstance(key_expression, MacroFunc):
 611                expressions.append(key_expression)
 612                continue
 613
 614            key = key_expression.name.lower()
 615
 616            start = self._curr
 617            value: t.Optional[exp.Expr | str]
 618
 619            if key in table_keys:
 620                value = self._parse_table_parts()
 621                if value and self._prev.token_type == TokenType.STRING:
 622                    self.raise_error(
 623                        f"'{key}' property cannot be a string value: {value}. "
 624                        "Please use the identifier syntax instead, e.g. foo.bar instead of 'foo.bar'"
 625                    )
 626            elif key == "columns":
 627                value = self._parse_schema()
 628            elif key == "kind":
 629                field = _parse_macro_or_clause(self, lambda: self._parse_id_var(any_token=True))
 630
 631                if not field or isinstance(field, (MacroVar, MacroFunc)):
 632                    value = field
 633                else:
 634                    try:
 635                        kind = ModelKindName[field.name.upper()]
 636                    except KeyError:
 637                        raise SQLMeshError(
 638                            f"Model kind specified as '{field.name}', but that is not a valid model kind.\n\nPlease specify one of {', '.join(ModelKindName)}."
 639                        )
 640
 641                    if kind in (
 642                        ModelKindName.INCREMENTAL_BY_TIME_RANGE,
 643                        ModelKindName.INCREMENTAL_BY_UNIQUE_KEY,
 644                        ModelKindName.INCREMENTAL_BY_PARTITION,
 645                        ModelKindName.INCREMENTAL_UNMANAGED,
 646                        ModelKindName.SEED,
 647                        ModelKindName.VIEW,
 648                        ModelKindName.SCD_TYPE_2,
 649                        ModelKindName.SCD_TYPE_2_BY_TIME,
 650                        ModelKindName.SCD_TYPE_2_BY_COLUMN,
 651                        ModelKindName.CUSTOM,
 652                    ) and self._match(TokenType.L_PAREN, advance=False):
 653                        props = self._parse_wrapped_csv(functools.partial(_parse_props, self))
 654                    else:
 655                        props = None
 656
 657                    value = self.expression(ModelKind(this=kind.value, expressions=props))
 658            elif key == "expression":
 659                value = self._parse_conjunction()
 660            elif key == "partitioned_by":
 661                partitioned_by = self._parse_partitioned_by()
 662                if isinstance(partitioned_by.this, exp.Schema):
 663                    value = exp.tuple_(*partitioned_by.this.expressions)
 664                else:
 665                    value = partitioned_by.this
 666            elif key == "clustered_by":
 667                # Bare AUTO / NONE are Databricks liquid clustering keywords, not column refs.
 668                # Detect keywords by token type: unquoted bare identifiers arrive as VAR tokens.
 669                # Backtick-quoted identifiers (e.g. `auto`) have IDENTIFIER token type and are
 670                # treated as real column names.
 671                if (
 672                    self._curr is not None
 673                    and self._curr.token_type == TokenType.VAR
 674                    and self._curr.text.upper() in LIQUID_CLUSTERING_KEYWORDS
 675                ):
 676                    value = exp.Var(this=self._curr.text.upper())
 677                    self._advance()
 678                else:
 679                    parsed = self._parse_bracket(self._parse_field(any_token=True))
 680                    # Unwrap Paren wrapping a bare column to match partitioned_by normalisation:
 681                    # clustered_by (a) → stored as Column(a), not Paren(Column(a)).
 682                    # Preserve parens around function expressions: (TO_DATE(col)) stays as-is.
 683                    if isinstance(parsed, exp.Paren) and isinstance(parsed.this, exp.Column):
 684                        value = parsed.unnest()
 685                    else:
 686                        value = parsed
 687            else:
 688                value = self._parse_bracket(self._parse_field(any_token=True))
 689
 690            if isinstance(value, exp.Expr):
 691                value.meta["sql"] = self._find_sql(start, self._prev)
 692
 693            expressions.append(self.expression(exp.Property(this=key, value=value)))
 694
 695        return self.expression(expression_type(expressions=expressions))
 696
 697    return parse
 698
 699
 700PARSERS = {
 701    "MODEL": _create_parser(Model, ["name"]),
 702    "AUDIT": _create_parser(Audit, ["model"]),
 703    "METRIC": _create_parser(Metric, ["name"]),
 704}
 705
 706
 707def _props_sql(self: Generator, expressions: t.List[exp.Expr]) -> str:
 708    props = []
 709    size = len(expressions)
 710
 711    for i, prop in enumerate(expressions):
 712        if isinstance(prop, MacroFunc):
 713            sql = self.indent(self.sql(prop, comment=False))
 714        else:
 715            sql = self.indent(f"{prop.name} {self.sql(prop, 'value')}")
 716
 717        if i < size - 1:
 718            sql += ","
 719
 720        props.append(self.maybe_comment(sql, expression=prop))
 721
 722    return "\n".join(props)
 723
 724
 725def _on_virtual_update_sql(self: Generator, expressions: t.List[exp.Expr]) -> str:
 726    statements = "\n".join(
 727        self.sql(expression)
 728        if isinstance(expression, JinjaStatement)
 729        else f"{self.sql(expression)};"
 730        for expression in expressions
 731    )
 732    return f"{ON_VIRTUAL_UPDATE_BEGIN};\n{statements}\n{ON_VIRTUAL_UPDATE_END};"
 733
 734
 735def _sqlmesh_ddl_sql(self: Generator, expression: Model | Audit | Metric, name: str) -> str:
 736    return "\n".join([f"{name} (", _props_sql(self, expression.expressions), ")"])
 737
 738
 739def _model_kind_sql(self: Generator, expression: ModelKind) -> str:
 740    props = _props_sql(self, expression.expressions)
 741    if props:
 742        return "\n".join([f"{expression.this} (", props, ")"])
 743    return expression.name.upper()
 744
 745
 746def _macro_keyword_func_sql(self: Generator, expression: exp.Expr) -> str:
 747    name = expression.name
 748    keyword = name.replace("_", " ")
 749    *args, clause = expression.expressions
 750    macro = f"@{name}({self.format_args(*args)})"
 751    return self.sql(clause).replace(keyword, macro, 1)
 752
 753
 754def _macro_func_sql(self: Generator, expression: MacroFunc) -> str:
 755    expression = expression.this
 756    name = expression.name
 757    if name in KEYWORD_MACROS:
 758        sql = _macro_keyword_func_sql(self, expression)
 759    else:
 760        sql = f"@{name}({self.format_args(*expression.expressions)})"
 761    return self.maybe_comment(sql, expression)
 762
 763
 764def _whens_sql(self: Generator, expression: exp.Whens) -> str:
 765    if isinstance(expression.parent, exp.Merge):
 766        return self.whens_sql(expression)
 767
 768    # If the `WHEN` clauses aren't part of a MERGE statement (e.g. they
 769    # appear in the `MODEL` DDL), then we will wrap them with parentheses.
 770    return self.wrap(self.expressions(expression, sep=" ", indent=False))
 771
 772
 773def _parse_interval_span(self: Parser, this: exp.Expr) -> exp.Interval:
 774    interval = self.__parse_interval_span(this)  # type: ignore
 775    # Without this, @unit in `INTERVAL @value @unit` is misread as an alias.
 776    if not interval.args.get("unit") and self._match(TokenType.PARAMETER):
 777        macro = _parse_macro(self)
 778        if macro is not None:
 779            interval.set("unit", macro)
 780    return interval
 781
 782
 783def _override(klass: t.Type[Tokenizer | Parser], func: t.Callable) -> None:
 784    name = func.__name__
 785    setattr(klass, f"_{name}", getattr(klass, name))
 786    setattr(klass, name, func)
 787
 788
 789def format_model_expressions(
 790    expressions: t.List[exp.Expr],
 791    dialect: t.Optional[str] = None,
 792    rewrite_casts: bool = True,
 793    **kwargs: t.Any,
 794) -> str:
 795    """Format a model's expressions into a standardized format.
 796
 797    Args:
 798        expressions: The model's expressions, must be at least model def + query.
 799        dialect: The dialect to render the expressions as.
 800        rewrite_casts: Whether to rewrite all casts to use the :: syntax.
 801        **kwargs: Additional keyword arguments to pass to the sql generator.
 802
 803    Returns:
 804        A string representing the formatted model.
 805    """
 806    if len(expressions) == 1 and is_meta_expression(expressions[0]):
 807        # Meta expressions (MODEL/AUDIT/METRIC) are SQLMesh DDL, not standard SQL,
 808        # so they must never be transpiled to the target dialect (e.g. tsql would
 809        # rewrite a boolean property like `allow_partials TRUE` to `(1 = 1)`).
 810        return expressions[0].sql(pretty=True, dialect=None)
 811
 812    if rewrite_casts:
 813
 814        def cast_to_colon(node: exp.Expr) -> exp.Expr:
 815            # Directly check type instead of isinstance to avoid rewriting subclasses of CAST, e.g. JSONCast
 816            if type(node) is exp.Cast and not any(
 817                # Only convert CAST into :: if it doesn't have additional args set, otherwise this
 818                # conversion could alter the semantics (eg. changing SAFE_CAST in BigQuery to CAST)
 819                arg
 820                for name, arg in node.args.items()
 821                if name not in ("this", "to")
 822            ):
 823                this = node.this
 824
 825                if not isinstance(this, (exp.Binary, exp.Unary)) or isinstance(this, exp.Paren):
 826                    cast = DColonCast(this=this, to=node.to)
 827                    cast.comments = node.comments
 828                    node = cast
 829
 830            exp.replace_children(node, cast_to_colon)
 831            return node
 832
 833        new_expressions = []
 834        for expression in expressions:
 835            expression = expression.copy()
 836            exp.replace_children(expression, cast_to_colon)
 837            new_expressions.append(expression)
 838
 839        expressions = new_expressions
 840
 841    return ";\n\n".join(
 842        # Meta expressions (MODEL/AUDIT/METRIC) are SQLMesh DDL and must stay
 843        # dialect-agnostic; only the actual query/statement expressions transpile.
 844        expression.sql(
 845            pretty=True,
 846            dialect=None if is_meta_expression(expression) else dialect,
 847            **kwargs,
 848        )
 849        for expression in expressions
 850    ).strip()
 851
 852
 853def text_diff(
 854    a: t.List[exp.Expr],
 855    b: t.List[exp.Expr],
 856    a_dialect: t.Optional[str] = None,
 857    b_dialect: t.Optional[str] = None,
 858) -> str:
 859    """Find the unified text diff between two expressions."""
 860    a_sql = [
 861        line
 862        for expr in a
 863        for line in expr.sql(pretty=True, comments=False, dialect=a_dialect).split("\n")
 864    ]
 865    b_sql = [
 866        line
 867        for expr in b
 868        for line in expr.sql(pretty=True, comments=False, dialect=b_dialect).split("\n")
 869    ]
 870    return "\n".join(unified_diff(a_sql, b_sql))
 871
 872
 873WS_OR_COMMENT = r"(?:\s|--[^\n]*\n|/\*.*?\*/)"
 874HEADER = r"\b(?:model|audit)\b(?=\s*\()"
 875KEY_BOUNDARY = r"(?:\(|,)"  # key is preceded by either '(' or ','
 876DIALECT_VALUE = r"['\"]?(?P<dialect>[a-z][a-z0-9]*)['\"]?"
 877VALUE_BOUNDARY = r"(?=,|\))"  # value is followed by comma or closing paren
 878
 879DIALECT_PATTERN = re.compile(
 880    rf"{HEADER}.*?{KEY_BOUNDARY}{WS_OR_COMMENT}*dialect{WS_OR_COMMENT}+{DIALECT_VALUE}{WS_OR_COMMENT}*{VALUE_BOUNDARY}",
 881    re.IGNORECASE | re.DOTALL,
 882)
 883
 884
 885def _is_command_statement(command: str, tokens: t.List[Token], pos: int) -> bool:
 886    try:
 887        return (
 888            tokens[pos].text.upper() == command.upper()
 889            and tokens[pos + 1].token_type == TokenType.SEMICOLON
 890        )
 891    except IndexError:
 892        return False
 893
 894
 895JINJA_QUERY_BEGIN = "JINJA_QUERY_BEGIN"
 896JINJA_STATEMENT_BEGIN = "JINJA_STATEMENT_BEGIN"
 897JINJA_END = "JINJA_END"
 898ON_VIRTUAL_UPDATE_BEGIN = "ON_VIRTUAL_UPDATE_BEGIN"
 899ON_VIRTUAL_UPDATE_END = "ON_VIRTUAL_UPDATE_END"
 900
 901
 902def _is_jinja_statement_begin(tokens: t.List[Token], pos: int) -> bool:
 903    return _is_command_statement(JINJA_STATEMENT_BEGIN, tokens, pos)
 904
 905
 906def _is_jinja_query_begin(tokens: t.List[Token], pos: int) -> bool:
 907    return _is_command_statement(JINJA_QUERY_BEGIN, tokens, pos)
 908
 909
 910def _is_jinja_end(tokens: t.List[Token], pos: int) -> bool:
 911    return _is_command_statement(JINJA_END, tokens, pos)
 912
 913
 914def jinja_query(query: str) -> JinjaQuery:
 915    return JinjaQuery(this=exp.Literal.string(query.strip()))
 916
 917
 918def jinja_statement(statement: str) -> JinjaStatement:
 919    return JinjaStatement(this=exp.Literal.string(statement.strip()))
 920
 921
 922def _is_virtual_statement_begin(tokens: t.List[Token], pos: int) -> bool:
 923    return _is_command_statement(ON_VIRTUAL_UPDATE_BEGIN, tokens, pos)
 924
 925
 926def _is_virtual_statement_end(tokens: t.List[Token], pos: int) -> bool:
 927    return _is_command_statement(ON_VIRTUAL_UPDATE_END, tokens, pos)
 928
 929
 930def virtual_statement(statements: t.List[exp.Expr]) -> VirtualUpdateStatement:
 931    return VirtualUpdateStatement(expressions=statements)
 932
 933
 934class ChunkType(Enum):
 935    JINJA_QUERY = auto()
 936    JINJA_STATEMENT = auto()
 937    SQL = auto()
 938    VIRTUAL_STATEMENT = auto()
 939    VIRTUAL_JINJA_STATEMENT = auto()
 940
 941
 942def parse_one(
 943    sql: str, dialect: t.Optional[str] = None, into: t.Optional[exp.IntoType] = None
 944) -> exp.Expr:
 945    expressions = parse(sql, default_dialect=dialect, match_dialect=False, into=into)
 946    if not expressions:
 947        raise SQLMeshError(f"No expressions found in '{sql}'")
 948    elif len(expressions) > 1:
 949        raise SQLMeshError(f"Multiple expressions found in '{sql}'")
 950    return expressions[0]
 951
 952
 953def parse(
 954    sql: str,
 955    default_dialect: t.Optional[str] = None,
 956    match_dialect: bool = True,
 957    into: t.Optional[exp.IntoType] = None,
 958) -> t.List[exp.Expr]:
 959    """Parse a sql string.
 960
 961    Supports parsing model definition.
 962    If a jinja block is detected, the query is stored as raw string in a Jinja node.
 963
 964    Args:
 965        sql: The sql based definition.
 966        default_dialect: The dialect to use if the model does not specify one.
 967
 968    Returns:
 969        A list of the parsed expressions: [Model, *Statements, Query, *Statements]
 970    """
 971    match = match_dialect and DIALECT_PATTERN.search(sql[:MAX_MODEL_DEFINITION_SIZE])
 972    dialect_str = match.group("dialect") if match else None
 973    dialect = Dialect.get_or_raise(dialect_str or default_dialect)
 974
 975    tokens = dialect.tokenize(sql)
 976    chunks: t.List[t.Tuple[t.List[Token], ChunkType]] = [([], ChunkType.SQL)]
 977    total = len(tokens)
 978
 979    pos = 0
 980    virtual = False
 981    while pos < total:
 982        token = tokens[pos]
 983        if _is_virtual_statement_end(tokens, pos):
 984            chunks[-1][0].append(token)
 985            virtual = False
 986            chunks.append(([], ChunkType.SQL))
 987            pos += 2
 988        elif _is_jinja_end(tokens, pos) or (
 989            chunks[-1][1] == ChunkType.SQL
 990            and token.token_type == TokenType.SEMICOLON
 991            and pos < total - 1
 992        ):
 993            if token.token_type == TokenType.SEMICOLON:
 994                pos += 1
 995            else:
 996                # Jinja end statement
 997                chunks[-1][0].append(token)
 998                pos += 2
 999            chunks.append(
1000                (
1001                    [],
1002                    ChunkType.VIRTUAL_STATEMENT
1003                    if virtual and tokens[pos] != ON_VIRTUAL_UPDATE_END
1004                    else ChunkType.SQL,
1005                )
1006            )
1007        elif _is_jinja_query_begin(tokens, pos):
1008            chunks.append(([token], ChunkType.JINJA_QUERY))
1009            pos += 2
1010        elif _is_jinja_statement_begin(tokens, pos):
1011            chunks.append(([token], ChunkType.JINJA_STATEMENT))
1012            pos += 2
1013        elif _is_virtual_statement_begin(tokens, pos):
1014            chunks.append(([token], ChunkType.VIRTUAL_STATEMENT))
1015            pos += 2
1016            virtual = True
1017        else:
1018            chunks[-1][0].append(token)
1019            pos += 1
1020
1021    parser = dialect.parser()
1022    expressions: t.List[exp.Expr] = []
1023
1024    def parse_sql_chunk(chunk: t.List[Token], meta_sql: bool = True) -> t.List[exp.Expr]:
1025        parsed_expressions: t.List[t.Optional[exp.Expr]] = (
1026            parser.parse(chunk, sql) if into is None else parser.parse_into(into, chunk, sql)
1027        )
1028        expressions = []
1029        for expression in parsed_expressions:
1030            if expression:
1031                if meta_sql:
1032                    expression.meta["sql"] = parser._find_sql(chunk[0], chunk[-1])
1033                expressions.append(expression)
1034        return expressions
1035
1036    def parse_jinja_chunk(chunk: t.List[Token], meta_sql: bool = True) -> exp.Expr:
1037        start, *_, end = chunk
1038        segment = sql[start.end + 2 : end.start - 1]
1039        factory = jinja_query if chunk_type == ChunkType.JINJA_QUERY else jinja_statement
1040        expression = factory(segment.strip())
1041        if meta_sql:
1042            expression.meta["sql"] = sql[start.start : end.end + 1]
1043        return expression
1044
1045    def parse_virtual_statement(
1046        chunks: t.List[t.Tuple[t.List[Token], ChunkType]], pos: int
1047    ) -> t.Tuple[t.List[exp.Expr], int]:
1048        # For virtual statements we need to handle both SQL and Jinja nested blocks within the chunk
1049        virtual_update_statements: t.List[exp.Expr] = []
1050        start = chunks[pos][0][0].start
1051
1052        while (
1053            chunks[pos - 1][0] == [] or chunks[pos - 1][0][-1].text.upper() != ON_VIRTUAL_UPDATE_END
1054        ):
1055            chunk, chunk_type = chunks[pos]
1056            if chunk_type == ChunkType.JINJA_STATEMENT:
1057                virtual_update_statements.append(parse_jinja_chunk(chunk, False))
1058            else:
1059                virtual_update_statements.extend(
1060                    parse_sql_chunk(
1061                        chunk[int(chunk[0].text.upper() == ON_VIRTUAL_UPDATE_BEGIN) : -1], False
1062                    ),
1063                )
1064            pos += 1
1065
1066        if virtual_update_statements:
1067            statements = virtual_statement(virtual_update_statements)
1068            end = chunk[-1].end + 1
1069            statements.meta["sql"] = sql[start:end]
1070            return [statements], pos
1071
1072        return [], pos
1073
1074    pos = 0
1075    total_chunks = len(chunks)
1076    while pos < total_chunks:
1077        chunk, chunk_type = chunks[pos]
1078        if chunk_type == ChunkType.VIRTUAL_STATEMENT:
1079            virtual_expression, pos = parse_virtual_statement(chunks, pos)
1080            expressions.extend(virtual_expression)
1081        elif chunk_type == ChunkType.SQL:
1082            expressions.extend(parse_sql_chunk(chunk))
1083        else:
1084            expressions.append(parse_jinja_chunk(chunk))
1085        pos += 1
1086
1087    return expressions
1088
1089
1090def extend_sqlglot() -> None:
1091    """Extend SQLGlot with SQLMesh's custom macro aware dialect."""
1092    tokenizers = {Tokenizer}
1093    parsers = {Parser}
1094    generators = {Generator}
1095
1096    for dialect in Dialect.classes.values():
1097        # Athena picks a different Tokenizer / Parser / Generator depending on the query
1098        # so this ensures that the extra ones it defines are also extended
1099        if dialect == athena.Athena:
1100            tokenizers.add(athena._TrinoTokenizer)
1101            parsers.add(AthenaTrinoParser)
1102            generators.add(athena_generators.AthenaTrinoGenerator)
1103            generators.add(athena_generators._HiveGenerator)
1104
1105        if hasattr(dialect, "Tokenizer"):
1106            tokenizers.add(dialect.Tokenizer)
1107        if hasattr(dialect, "Parser"):
1108            parsers.add(dialect.Parser)
1109        if hasattr(dialect, "Generator"):
1110            generators.add(dialect.Generator)
1111
1112    for tokenizer in tokenizers:
1113        tokenizer.VAR_SINGLE_TOKENS.update(SQLMESH_MACRO_PREFIX)
1114
1115    for parser in parsers:
1116        parser.FUNCTIONS.update({"JINJA": Jinja.from_arg_list, "METRIC": MetricAgg.from_arg_list})
1117        parser.PLACEHOLDER_PARSERS.update({TokenType.PARAMETER: _parse_macro})
1118        parser.QUERY_MODIFIER_PARSERS.update(
1119            {TokenType.PARAMETER: lambda self: _parse_body_macro(self)}
1120        )
1121
1122    for generator in generators:
1123        if MacroFunc not in generator.TRANSFORMS:
1124            generator.TRANSFORMS.update(
1125                {
1126                    Audit: lambda self, e: _sqlmesh_ddl_sql(self, e, "AUDIT"),
1127                    DColonCast: lambda self, e: f"{self.sql(e, 'this')}::{self.sql(e, 'to')}",
1128                    Jinja: lambda self, e: e.name,
1129                    JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
1130                    JinjaStatement: lambda self, e: (
1131                        f"{JINJA_STATEMENT_BEGIN};\n{e.name}\n{JINJA_END};"
1132                    ),
1133                    VirtualUpdateStatement: lambda self, e: _on_virtual_update_sql(self, e),
1134                    MacroDef: lambda self, e: f"@DEF({self.sql(e.this)}, {self.sql(e.expression)})",
1135                    MacroFunc: _macro_func_sql,
1136                    MacroStrReplace: lambda self, e: f"@{self.sql(e.this)}",
1137                    MacroSQL: lambda self, e: f"@SQL({self.sql(e.this)})",
1138                    MacroVar: lambda self, e: f"@{e.name}",
1139                    Metric: lambda self, e: _sqlmesh_ddl_sql(self, e, "METRIC"),
1140                    Model: lambda self, e: _sqlmesh_ddl_sql(self, e, "MODEL"),
1141                    ModelKind: _model_kind_sql,
1142                    PythonCode: lambda self, e: self.expressions(e, sep="\n", indent=False),
1143                    StagedFilePath: lambda self, e: self.table_sql(e),
1144                    exp.Whens: _whens_sql,
1145                }
1146            )
1147        if MacroDef not in generator.WITH_SEPARATED_COMMENTS:
1148            generator.WITH_SEPARATED_COMMENTS = (
1149                *generator.WITH_SEPARATED_COMMENTS,
1150                Model,
1151                MacroDef,
1152            )
1153
1154        generator.UNWRAPPED_INTERVAL_VALUES = (
1155            *generator.UNWRAPPED_INTERVAL_VALUES,
1156            MacroStrReplace,
1157            MacroVar,
1158        )
1159
1160    _override(Parser, _parse_select)
1161    _override(Parser, _parse_statement)
1162    _override(Parser, _parse_join)
1163    _override(Parser, _parse_order)
1164    _override(Parser, _parse_where)
1165    _override(Parser, _parse_group)
1166    _override(Parser, _parse_with)
1167    _override(Parser, _parse_having)
1168    _override(Parser, _parse_limit)
1169    _override(Parser, _parse_value)
1170    _override(Parser, _parse_lambda)
1171    _override(Parser, _parse_types)
1172    _override(Parser, _parse_if)
1173    _override(TSQL.Parser, Parser._parse_if)
1174    _override(Parser, _parse_id_var)
1175    _override(Parser, _parse_interval_span)
1176    _override(Parser, _warn_unsupported)
1177    _override(Snowflake.Parser, _parse_table_parts)
1178
1179    # DuckDB's prefix absolute power operator `@` clashes with the macro syntax
1180    DuckDB.Parser.NO_PAREN_FUNCTION_PARSERS.pop("@", None)
1181
1182
1183def select_from_values(
1184    values: t.List[t.Tuple[t.Any, ...]],
1185    columns_to_types: t.Dict[str, exp.DataType],
1186    batch_size: int = 0,
1187    alias: str = "t",
1188) -> t.Iterator[exp.Select]:
1189    """Generate a VALUES expression that has a select wrapped around it to cast the values to their correct types.
1190
1191    Args:
1192        values: List of values to use for the VALUES expression.
1193        columns_to_types: Mapping of column names to types to assign to the values.
1194        batch_size: The maximum number of tuples per batches. Defaults to sys.maxsize if <= 0.
1195        alias: The alias to assign to the values expression. If not provided then will default to "t"
1196
1197    Returns:
1198        This method operates as a generator and yields a VALUES expression.
1199    """
1200    if batch_size <= 0:
1201        batch_size = sys.maxsize
1202    num_rows = len(values)
1203    for i in range(0, num_rows, batch_size):
1204        yield select_from_values_for_batch_range(
1205            values=values,
1206            target_columns_to_types=columns_to_types,
1207            batch_start=i,
1208            batch_end=min(i + batch_size, num_rows),
1209            alias=alias,
1210        )
1211
1212
1213def select_from_values_for_batch_range(
1214    values: t.List[t.Tuple[t.Any, ...]],
1215    target_columns_to_types: t.Dict[str, exp.DataType],
1216    batch_start: int,
1217    batch_end: int,
1218    alias: str = "t",
1219    source_columns: t.Optional[t.List[str]] = None,
1220) -> exp.Select:
1221    source_columns = source_columns or list(target_columns_to_types)
1222    source_columns_to_types = get_source_columns_to_types(target_columns_to_types, source_columns)
1223
1224    if not values:
1225        # Ensures we don't generate an empty VALUES clause & forces a zero-row output
1226        where = exp.false()
1227        expressions = [
1228            tuple(exp.cast(exp.null(), to=kind) for kind in source_columns_to_types.values())
1229        ]
1230    else:
1231        where = None
1232        expressions = [
1233            tuple(transform_values(v, source_columns_to_types))
1234            for v in values[batch_start:batch_end]
1235        ]
1236
1237    values_exp = exp.values(expressions, alias=alias, columns=source_columns_to_types)
1238    if values:
1239        # BigQuery crashes on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([NULL]) AS x`, but not
1240        # on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([CAST(NULL AS TIMESTAMP)]) AS x`. This
1241        # ensures nulls under the `Values` expression are cast to avoid similar issues.
1242        for value, kind in zip(
1243            values_exp.expressions[0].expressions, source_columns_to_types.values()
1244        ):
1245            if isinstance(value, exp.Null):
1246                value.replace(exp.cast(value, to=kind))
1247
1248    casted_columns = [
1249        exp.alias_(
1250            exp.cast(
1251                exp.column(column) if column in source_columns_to_types else exp.Null(), to=kind
1252            ),
1253            column,
1254            copy=False,
1255        )
1256        for column, kind in target_columns_to_types.items()
1257    ]
1258    return exp.select(*casted_columns).from_(values_exp, copy=False).where(where, copy=False)
1259
1260
1261def pandas_to_sql(
1262    df: pd.DataFrame,
1263    columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
1264    batch_size: int = 0,
1265    alias: str = "t",
1266) -> t.Iterator[exp.Select]:
1267    """Convert a pandas dataframe into a VALUES sql statement.
1268
1269    Args:
1270        df: A pandas dataframe to convert.
1271        columns_to_types: Mapping of column names to types to assign to the values.
1272        batch_size: The maximum number of tuples per batches. Defaults to sys.maxsize if <= 0.
1273        alias: The alias to assign to the values expression. If not provided then will default to "t"
1274
1275    Returns:
1276        This method operates as a generator and yields a VALUES expression.
1277    """
1278    yield from select_from_values(
1279        values=list(df.itertuples(index=False, name=None)),
1280        columns_to_types=columns_to_types or columns_to_types_from_df(df),
1281        batch_size=batch_size,
1282        alias=alias,
1283    )
1284
1285
1286def set_default_catalog(
1287    table: str | exp.Table,
1288    default_catalog: t.Optional[str],
1289) -> exp.Table:
1290    table = exp.to_table(table)
1291
1292    if default_catalog and not table.catalog and table.db:
1293        table.set("catalog", exp.parse_identifier(default_catalog))
1294
1295    return table
1296
1297
1298@lru_cache(maxsize=16384)
1299def normalize_model_name(
1300    table: str | exp.Table | exp.Column,
1301    default_catalog: t.Optional[str],
1302    dialect: DialectType = None,
1303) -> str:
1304    if isinstance(table, exp.Column):
1305        table = exp.table_(table.this, db=table.args.get("table"), catalog=table.args.get("db"))
1306    else:
1307        # We are relying on sqlglot's flexible parsing here to accept quotes from other dialects.
1308        # Ex: I have a a normalized name of '"my_table"' but the dialect is spark and therefore we should
1309        # expect spark quotes to be backticks ('`') instead of double quotes ('"'). sqlglot today is flexible
1310        # and will still parse this correctly and we rely on that.
1311        table = exp.to_table(table, dialect=dialect)
1312
1313    table = set_default_catalog(table, default_catalog)
1314    # An alternative way to do this is the following: exp.table_name(table, dialect=dialect, identify=True)
1315    # This though would result in the names being normalized to the target dialect AND the quotes while the below
1316    # approach just normalizes the names.
1317    # By just normalizing names and using sqlglot dialect for quotes this makes it easier for dialects that have
1318    # compatible normalization strategies but incompatible quoting to still work together without user hassle
1319    return exp.table_name(normalize_identifiers(table, dialect=dialect), identify=True)
1320
1321
1322def find_tables(
1323    expression: exp.Expr, default_catalog: t.Optional[str], dialect: DialectType = None
1324) -> t.Set[str]:
1325    """Find all tables referenced in a query.
1326
1327    Caches the result in the meta field 'tables'.
1328
1329    Args:
1330        expressions: The query to find the tables in.
1331        dialect: The dialect to use for normalization of table names.
1332
1333    Returns:
1334        A Set of all the table names.
1335    """
1336    if TABLES_META not in expression.meta:
1337        expression.meta[TABLES_META] = {
1338            normalize_model_name(table, default_catalog=default_catalog, dialect=dialect)
1339            for scope in traverse_scope(expression)
1340            for table in scope.tables
1341            if table.name and table.name not in scope.cte_sources
1342        }
1343    return expression.meta[TABLES_META]
1344
1345
1346def add_table(node: exp.Expr, table: str) -> exp.Expr:
1347    """Add a table to all columns in an expression."""
1348
1349    def _transform(node: exp.Expr) -> exp.Expr:
1350        if isinstance(node, exp.Column) and not node.table:
1351            return exp.column(node.this, table=table)
1352        if isinstance(node, exp.Identifier):
1353            return exp.column(node, table=table)
1354        return node
1355
1356    return node.transform(_transform)
1357
1358
1359def transform_values(
1360    values: t.Tuple[t.Any, ...], columns_to_types: t.Dict[str, exp.DataType]
1361) -> t.Iterator[t.Any]:
1362    """Perform transformations on values given columns_to_types."""
1363
1364    def _transform_value(value: t.Any, dtype: exp.DataType) -> t.Any:
1365        if (
1366            isinstance(value, list)
1367            and dtype.is_type(*exp.DataType.ARRAY_TYPES)
1368            and len(dtype.expressions) == 1
1369        ):
1370            element_type = dtype.expressions[0]
1371            return exp.convert([_transform_value(v, element_type) for v in value])
1372
1373        if (
1374            isinstance(value, dict)
1375            and dtype.is_type(*exp.DataType.STRUCT_TYPES)
1376            and len(value) == len(dtype.expressions)
1377        ):
1378            expressions = []
1379            for (field_name, field_value), field_type in zip(value.items(), dtype.expressions):
1380                if isinstance(field_type, exp.ColumnDef):
1381                    field_type = field_type.kind
1382                else:
1383                    field_type = exp.DataType.build(exp.DataType.Type.UNKNOWN)
1384
1385                expressions.append(
1386                    exp.PropertyEQ(
1387                        this=exp.to_identifier(field_name),
1388                        expression=_transform_value(field_value, field_type),
1389                    )
1390                )
1391
1392            return exp.Struct(expressions=expressions)
1393
1394        if dtype.is_type(exp.DataType.Type.JSON):
1395            return exp.func("PARSE_JSON", f"'{value}'")
1396
1397        return exp.convert(value)
1398
1399    for col_value, col_type in zip(values, columns_to_types.values()):
1400        yield _transform_value(col_value, col_type)
1401
1402
1403def to_schema(sql_path: str | exp.Table, dialect: DialectType = None) -> exp.Table:
1404    if isinstance(sql_path, exp.Table) and sql_path.this is None:
1405        return sql_path
1406    table = exp.to_table(
1407        sql_path.copy() if isinstance(sql_path, exp.Table) else sql_path, dialect=dialect
1408    )
1409    table.set("catalog", table.args.get("db"))
1410    table.set("db", table.args.get("this"))
1411    table.set("this", None)
1412    return table
1413
1414
1415def schema_(
1416    db: exp.Identifier | str,
1417    catalog: t.Optional[exp.Identifier | str] = None,
1418    quoted: t.Optional[bool] = None,
1419) -> exp.Table:
1420    """Build a Schema.
1421
1422    Args:
1423        db: Database name.
1424        catalog: Catalog name.
1425        quoted: Whether to force quotes on the schema's identifiers.
1426
1427    Returns:
1428        The new Schema instance.
1429    """
1430    return exp.Table(
1431        this=None,
1432        db=exp.to_identifier(db, quoted=quoted) if db else None,
1433        catalog=exp.to_identifier(catalog, quoted=quoted) if catalog else None,
1434    )
1435
1436
1437def normalize_mapping_schema(schema: t.Dict, dialect: DialectType) -> MappingSchema:
1438    return MappingSchema(_unquote_schema(schema), dialect=dialect, normalize=False)
1439
1440
1441def _unquote_schema(schema: t.Dict) -> t.Dict:
1442    """SQLGlot schema expects unquoted normalized keys."""
1443    return {
1444        k.strip('"'): _unquote_schema(v) if isinstance(v, dict) else v for k, v in schema.items()
1445    }
1446
1447
1448@contextmanager
1449def normalize_and_quote(
1450    query: E, dialect: DialectType, default_catalog: t.Optional[str], quote: bool = True
1451) -> t.Iterator[E]:
1452    qualify_tables(query, catalog=default_catalog, dialect=dialect)
1453    normalize_identifiers(query, dialect=dialect)
1454    yield query
1455    if quote:
1456        quote_identifiers(query, dialect=dialect)
1457
1458
1459def interpret_expression(e: exp.Expr) -> exp.Expr | str | int | float | bool:
1460    if e.is_int:
1461        return int(e.this)
1462    if e.is_number:
1463        return float(e.this)
1464    if isinstance(e, (exp.Literal, exp.Boolean)):
1465        return e.this
1466    return e
1467
1468
1469def interpret_key_value_pairs(
1470    e: exp.Tuple,
1471) -> t.Dict[str, exp.Expr | str | int | float | bool]:
1472    return {i.this.name: interpret_expression(i.expression) for i in e.expressions}
1473
1474
1475def extract_func_call(
1476    v: exp.Expr, allow_tuples: bool = False
1477) -> t.Tuple[str, t.Dict[str, exp.Expr]]:
1478    kwargs = {}
1479
1480    if isinstance(v, exp.Anonymous):
1481        func = v.name
1482        args = v.expressions
1483    elif isinstance(v, exp.Func):
1484        func = v.sql_name()
1485        args = list(v.args.values())
1486    elif isinstance(v, exp.Paren):
1487        func = ""
1488        args = [v.this]
1489    elif isinstance(v, exp.Tuple):  # airflow only
1490        if not allow_tuples:
1491            raise ConfigError("Audit name is missing (eg. MY_AUDIT())")
1492
1493        func = ""
1494        args = v.expressions
1495    else:
1496        return v.name.lower(), {}
1497
1498    for arg in args:
1499        if not isinstance(arg, (exp.PropertyEQ, exp.EQ)):
1500            raise ConfigError(
1501                f"Function '{func}' must be called with key-value arguments like {func}(arg := value)."
1502            )
1503        kwargs[arg.left.name.lower()] = arg.right
1504    return func.lower(), kwargs
1505
1506
1507def extract_function_calls(func_calls: t.Any, allow_tuples: bool = False) -> t.Any:
1508    """Used for extracting function calls for signals or audits."""
1509
1510    if isinstance(func_calls, (exp.Tuple, exp.Array)):
1511        return [extract_func_call(i, allow_tuples=allow_tuples) for i in func_calls.expressions]
1512    if isinstance(func_calls, exp.Paren):
1513        return [extract_func_call(func_calls.this, allow_tuples=allow_tuples)]
1514    if isinstance(func_calls, exp.Expr):
1515        return [extract_func_call(func_calls, allow_tuples=allow_tuples)]
1516    if isinstance(func_calls, list):
1517        function_calls = []
1518        for entry in func_calls:
1519            if isinstance(entry, dict):
1520                args = entry
1521                name = "" if allow_tuples else entry.pop("name")
1522            elif isinstance(entry, (tuple, list)):
1523                name, args = entry
1524            else:
1525                raise ConfigError(f"Audit must be a dictionary or named tuple. Got {entry}.")
1526
1527            function_calls.append(
1528                (
1529                    name.lower(),
1530                    {
1531                        key: parse_one(value) if isinstance(value, str) else value
1532                        for key, value in args.items()
1533                    },
1534                )
1535            )
1536
1537        return function_calls
1538
1539    return func_calls or []
1540
1541
1542def is_meta_expression(v: t.Any) -> bool:
1543    return isinstance(v, (Audit, Metric, Model))
1544
1545
1546def replace_merge_table_aliases(expression: exp.Expr, dialect: t.Optional[str] = None) -> exp.Expr:
1547    """
1548    Resolves references from the "source" and "target" tables (or their DBT equivalents)
1549    with the corresponding SQLMesh merge aliases (MERGE_SOURCE_ALIAS and MERGE_TARGET_ALIAS)
1550    """
1551    from sqlmesh.core.engine_adapter.base import MERGE_SOURCE_ALIAS, MERGE_TARGET_ALIAS
1552
1553    if isinstance(expression, exp.Column) and (first_part := expression.parts[0]):
1554        if first_part.this.lower() in ("target", "dbt_internal_dest", "__merge_target__"):
1555            first_part.replace(exp.to_identifier(MERGE_TARGET_ALIAS, quoted=True))
1556        elif first_part.this.lower() in ("source", "dbt_internal_source", "__merge_source__"):
1557            first_part.replace(exp.to_identifier(MERGE_SOURCE_ALIAS, quoted=True))
1558
1559    return expression