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    normalize_functions: t.Union[str, bool, None] = False,
 794    **kwargs: t.Any,
 795) -> str:
 796    """Format a model's expressions into a standardized format.
 797
 798    Args:
 799        expressions: The model's expressions, must be at least model def + query.
 800        dialect: The dialect to render the expressions as.
 801        rewrite_casts: Whether to rewrite all casts to use the :: syntax.
 802        normalize_functions: How to normalize function name casing.
 803
 804            * ``False`` (default) — preserves the original spelling of custom and audit
 805              function names.  SQLGlot built-in functions may still canonicalize because
 806              the parser discards the original token.
 807            * ``"upper"`` — uppercases all function names including custom audit
 808              references.
 809            * ``"lower"`` — lowercases all function names including built-ins.
 810            * ``True`` — defers to SQLGlot's generator default (uppercase).
 811            * ``None`` — passes ``None`` directly to the SQLGlot generator, which
 812              defers to SQLGlot's own default (typically uppercase, but may vary by
 813              dialect).  Note: this is the **direct generator API** behaviour.  When
 814              called via ``FormatConfig``, ``None`` is excluded by Pydantic's
 815              ``exclude_none`` serialization and this function receives its own ``False``
 816              default instead — so the two paths are not equivalent.
 817        **kwargs: Additional keyword arguments to pass to the sql generator.
 818
 819    Returns:
 820        A string representing the formatted model.
 821    """
 822    if len(expressions) == 1 and is_meta_expression(expressions[0]):
 823        # Meta expressions (MODEL/AUDIT/METRIC) are SQLMesh DDL, not standard SQL,
 824        # so they must never be transpiled to the target dialect (e.g. tsql would
 825        # rewrite a boolean property like `allow_partials TRUE` to `(1 = 1)`).
 826        return expressions[0].sql(
 827            pretty=True, dialect=None, normalize_functions=normalize_functions
 828        )
 829
 830    if rewrite_casts:
 831
 832        def cast_to_colon(node: exp.Expr) -> exp.Expr:
 833            # Directly check type instead of isinstance to avoid rewriting subclasses of CAST, e.g. JSONCast
 834            if type(node) is exp.Cast and not any(
 835                # Only convert CAST into :: if it doesn't have additional args set, otherwise this
 836                # conversion could alter the semantics (eg. changing SAFE_CAST in BigQuery to CAST)
 837                arg
 838                for name, arg in node.args.items()
 839                if name not in ("this", "to")
 840            ):
 841                this = node.this
 842
 843                if not isinstance(this, (exp.Binary, exp.Unary)) or isinstance(this, exp.Paren):
 844                    cast = DColonCast(this=this, to=node.to)
 845                    cast.comments = node.comments
 846                    node = cast
 847
 848            exp.replace_children(node, cast_to_colon)
 849            return node
 850
 851        new_expressions = []
 852        for expression in expressions:
 853            expression = expression.copy()
 854            exp.replace_children(expression, cast_to_colon)
 855            new_expressions.append(expression)
 856
 857        expressions = new_expressions
 858
 859    return ";\n\n".join(
 860        # Meta expressions (MODEL/AUDIT/METRIC) are SQLMesh DDL and must stay
 861        # dialect-agnostic; only the actual query/statement expressions transpile.
 862        expression.sql(
 863            pretty=True,
 864            dialect=None if is_meta_expression(expression) else dialect,
 865            normalize_functions=normalize_functions,
 866            **kwargs,
 867        )
 868        for expression in expressions
 869    ).strip()
 870
 871
 872def text_diff(
 873    a: t.List[exp.Expr],
 874    b: t.List[exp.Expr],
 875    a_dialect: t.Optional[str] = None,
 876    b_dialect: t.Optional[str] = None,
 877) -> str:
 878    """Find the unified text diff between two expressions."""
 879    a_sql = [
 880        line
 881        for expr in a
 882        for line in expr.sql(pretty=True, comments=False, dialect=a_dialect).split("\n")
 883    ]
 884    b_sql = [
 885        line
 886        for expr in b
 887        for line in expr.sql(pretty=True, comments=False, dialect=b_dialect).split("\n")
 888    ]
 889    return "\n".join(unified_diff(a_sql, b_sql))
 890
 891
 892WS_OR_COMMENT = r"(?:\s|--[^\n]*\n|/\*.*?\*/)"
 893HEADER = r"\b(?:model|audit)\b(?=\s*\()"
 894KEY_BOUNDARY = r"(?:\(|,)"  # key is preceded by either '(' or ','
 895DIALECT_VALUE = r"['\"]?(?P<dialect>[a-z][a-z0-9]*)['\"]?"
 896VALUE_BOUNDARY = r"(?=,|\))"  # value is followed by comma or closing paren
 897
 898DIALECT_PATTERN = re.compile(
 899    rf"{HEADER}.*?{KEY_BOUNDARY}{WS_OR_COMMENT}*dialect{WS_OR_COMMENT}+{DIALECT_VALUE}{WS_OR_COMMENT}*{VALUE_BOUNDARY}",
 900    re.IGNORECASE | re.DOTALL,
 901)
 902
 903
 904def _is_command_statement(command: str, tokens: t.List[Token], pos: int) -> bool:
 905    try:
 906        return (
 907            tokens[pos].text.upper() == command.upper()
 908            and tokens[pos + 1].token_type == TokenType.SEMICOLON
 909        )
 910    except IndexError:
 911        return False
 912
 913
 914JINJA_QUERY_BEGIN = "JINJA_QUERY_BEGIN"
 915JINJA_STATEMENT_BEGIN = "JINJA_STATEMENT_BEGIN"
 916JINJA_END = "JINJA_END"
 917ON_VIRTUAL_UPDATE_BEGIN = "ON_VIRTUAL_UPDATE_BEGIN"
 918ON_VIRTUAL_UPDATE_END = "ON_VIRTUAL_UPDATE_END"
 919
 920
 921def _is_jinja_statement_begin(tokens: t.List[Token], pos: int) -> bool:
 922    return _is_command_statement(JINJA_STATEMENT_BEGIN, tokens, pos)
 923
 924
 925def _is_jinja_query_begin(tokens: t.List[Token], pos: int) -> bool:
 926    return _is_command_statement(JINJA_QUERY_BEGIN, tokens, pos)
 927
 928
 929def _is_jinja_end(tokens: t.List[Token], pos: int) -> bool:
 930    return _is_command_statement(JINJA_END, tokens, pos)
 931
 932
 933def jinja_query(query: str) -> JinjaQuery:
 934    return JinjaQuery(this=exp.Literal.string(query.strip()))
 935
 936
 937def jinja_statement(statement: str) -> JinjaStatement:
 938    return JinjaStatement(this=exp.Literal.string(statement.strip()))
 939
 940
 941def _is_virtual_statement_begin(tokens: t.List[Token], pos: int) -> bool:
 942    return _is_command_statement(ON_VIRTUAL_UPDATE_BEGIN, tokens, pos)
 943
 944
 945def _is_virtual_statement_end(tokens: t.List[Token], pos: int) -> bool:
 946    return _is_command_statement(ON_VIRTUAL_UPDATE_END, tokens, pos)
 947
 948
 949def virtual_statement(statements: t.List[exp.Expr]) -> VirtualUpdateStatement:
 950    return VirtualUpdateStatement(expressions=statements)
 951
 952
 953class ChunkType(Enum):
 954    JINJA_QUERY = auto()
 955    JINJA_STATEMENT = auto()
 956    SQL = auto()
 957    VIRTUAL_STATEMENT = auto()
 958    VIRTUAL_JINJA_STATEMENT = auto()
 959
 960
 961def parse_one(
 962    sql: str, dialect: t.Optional[str] = None, into: t.Optional[exp.IntoType] = None
 963) -> exp.Expr:
 964    expressions = parse(sql, default_dialect=dialect, match_dialect=False, into=into)
 965    if not expressions:
 966        raise SQLMeshError(f"No expressions found in '{sql}'")
 967    elif len(expressions) > 1:
 968        raise SQLMeshError(f"Multiple expressions found in '{sql}'")
 969    return expressions[0]
 970
 971
 972def parse(
 973    sql: str,
 974    default_dialect: t.Optional[str] = None,
 975    match_dialect: bool = True,
 976    into: t.Optional[exp.IntoType] = None,
 977) -> t.List[exp.Expr]:
 978    """Parse a sql string.
 979
 980    Supports parsing model definition.
 981    If a jinja block is detected, the query is stored as raw string in a Jinja node.
 982
 983    Args:
 984        sql: The sql based definition.
 985        default_dialect: The dialect to use if the model does not specify one.
 986
 987    Returns:
 988        A list of the parsed expressions: [Model, *Statements, Query, *Statements]
 989    """
 990    match = match_dialect and DIALECT_PATTERN.search(sql[:MAX_MODEL_DEFINITION_SIZE])
 991    dialect_str = match.group("dialect") if match else None
 992    dialect = Dialect.get_or_raise(dialect_str or default_dialect)
 993
 994    tokens = dialect.tokenize(sql)
 995    chunks: t.List[t.Tuple[t.List[Token], ChunkType]] = [([], ChunkType.SQL)]
 996    total = len(tokens)
 997
 998    pos = 0
 999    virtual = False
1000    while pos < total:
1001        token = tokens[pos]
1002        if _is_virtual_statement_end(tokens, pos):
1003            chunks[-1][0].append(token)
1004            virtual = False
1005            chunks.append(([], ChunkType.SQL))
1006            pos += 2
1007        elif _is_jinja_end(tokens, pos) or (
1008            chunks[-1][1] == ChunkType.SQL
1009            and token.token_type == TokenType.SEMICOLON
1010            and pos < total - 1
1011        ):
1012            if token.token_type == TokenType.SEMICOLON:
1013                pos += 1
1014            else:
1015                # Jinja end statement
1016                chunks[-1][0].append(token)
1017                pos += 2
1018            chunks.append(
1019                (
1020                    [],
1021                    ChunkType.VIRTUAL_STATEMENT
1022                    if virtual and tokens[pos] != ON_VIRTUAL_UPDATE_END
1023                    else ChunkType.SQL,
1024                )
1025            )
1026        elif _is_jinja_query_begin(tokens, pos):
1027            chunks.append(([token], ChunkType.JINJA_QUERY))
1028            pos += 2
1029        elif _is_jinja_statement_begin(tokens, pos):
1030            chunks.append(([token], ChunkType.JINJA_STATEMENT))
1031            pos += 2
1032        elif _is_virtual_statement_begin(tokens, pos):
1033            chunks.append(([token], ChunkType.VIRTUAL_STATEMENT))
1034            pos += 2
1035            virtual = True
1036        else:
1037            chunks[-1][0].append(token)
1038            pos += 1
1039
1040    parser = dialect.parser()
1041    expressions: t.List[exp.Expr] = []
1042
1043    def parse_sql_chunk(chunk: t.List[Token], meta_sql: bool = True) -> t.List[exp.Expr]:
1044        parsed_expressions: t.List[t.Optional[exp.Expr]] = (
1045            parser.parse(chunk, sql) if into is None else parser.parse_into(into, chunk, sql)
1046        )
1047        expressions = []
1048        for expression in parsed_expressions:
1049            if expression:
1050                if meta_sql:
1051                    expression.meta["sql"] = parser._find_sql(chunk[0], chunk[-1])
1052                expressions.append(expression)
1053        return expressions
1054
1055    def parse_jinja_chunk(chunk: t.List[Token], meta_sql: bool = True) -> exp.Expr:
1056        start, *_, end = chunk
1057        segment = sql[start.end + 2 : end.start - 1]
1058        factory = jinja_query if chunk_type == ChunkType.JINJA_QUERY else jinja_statement
1059        expression = factory(segment.strip())
1060        if meta_sql:
1061            expression.meta["sql"] = sql[start.start : end.end + 1]
1062        return expression
1063
1064    def parse_virtual_statement(
1065        chunks: t.List[t.Tuple[t.List[Token], ChunkType]], pos: int
1066    ) -> t.Tuple[t.List[exp.Expr], int]:
1067        # For virtual statements we need to handle both SQL and Jinja nested blocks within the chunk
1068        virtual_update_statements: t.List[exp.Expr] = []
1069        start = chunks[pos][0][0].start
1070
1071        while (
1072            chunks[pos - 1][0] == [] or chunks[pos - 1][0][-1].text.upper() != ON_VIRTUAL_UPDATE_END
1073        ):
1074            chunk, chunk_type = chunks[pos]
1075            if chunk_type == ChunkType.JINJA_STATEMENT:
1076                virtual_update_statements.append(parse_jinja_chunk(chunk, False))
1077            else:
1078                virtual_update_statements.extend(
1079                    parse_sql_chunk(
1080                        chunk[int(chunk[0].text.upper() == ON_VIRTUAL_UPDATE_BEGIN) : -1], False
1081                    ),
1082                )
1083            pos += 1
1084
1085        if virtual_update_statements:
1086            statements = virtual_statement(virtual_update_statements)
1087            end = chunk[-1].end + 1
1088            statements.meta["sql"] = sql[start:end]
1089            return [statements], pos
1090
1091        return [], pos
1092
1093    pos = 0
1094    total_chunks = len(chunks)
1095    while pos < total_chunks:
1096        chunk, chunk_type = chunks[pos]
1097        if chunk_type == ChunkType.VIRTUAL_STATEMENT:
1098            virtual_expression, pos = parse_virtual_statement(chunks, pos)
1099            expressions.extend(virtual_expression)
1100        elif chunk_type == ChunkType.SQL:
1101            expressions.extend(parse_sql_chunk(chunk))
1102        else:
1103            expressions.append(parse_jinja_chunk(chunk))
1104        pos += 1
1105
1106    return expressions
1107
1108
1109def extend_sqlglot() -> None:
1110    """Extend SQLGlot with SQLMesh's custom macro aware dialect."""
1111    tokenizers = {Tokenizer}
1112    parsers = {Parser}
1113    generators = {Generator}
1114
1115    for dialect in Dialect.classes.values():
1116        # Athena picks a different Tokenizer / Parser / Generator depending on the query
1117        # so this ensures that the extra ones it defines are also extended
1118        if dialect == athena.Athena:
1119            tokenizers.add(athena._TrinoTokenizer)
1120            parsers.add(AthenaTrinoParser)
1121            generators.add(athena_generators.AthenaTrinoGenerator)
1122            generators.add(athena_generators._HiveGenerator)
1123
1124        if hasattr(dialect, "Tokenizer"):
1125            tokenizers.add(dialect.Tokenizer)
1126        if hasattr(dialect, "Parser"):
1127            parsers.add(dialect.Parser)
1128        if hasattr(dialect, "Generator"):
1129            generators.add(dialect.Generator)
1130
1131    for tokenizer in tokenizers:
1132        tokenizer.VAR_SINGLE_TOKENS.update(SQLMESH_MACRO_PREFIX)
1133
1134    for parser in parsers:
1135        parser.FUNCTIONS.update({"JINJA": Jinja.from_arg_list, "METRIC": MetricAgg.from_arg_list})
1136        parser.PLACEHOLDER_PARSERS.update({TokenType.PARAMETER: _parse_macro})
1137        parser.QUERY_MODIFIER_PARSERS.update(
1138            {TokenType.PARAMETER: lambda self: _parse_body_macro(self)}
1139        )
1140
1141    for generator in generators:
1142        if MacroFunc not in generator.TRANSFORMS:
1143            generator.TRANSFORMS.update(
1144                {
1145                    Audit: lambda self, e: _sqlmesh_ddl_sql(self, e, "AUDIT"),
1146                    DColonCast: lambda self, e: f"{self.sql(e, 'this')}::{self.sql(e, 'to')}",
1147                    Jinja: lambda self, e: e.name,
1148                    JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
1149                    JinjaStatement: lambda self, e: (
1150                        f"{JINJA_STATEMENT_BEGIN};\n{e.name}\n{JINJA_END};"
1151                    ),
1152                    VirtualUpdateStatement: lambda self, e: _on_virtual_update_sql(self, e),
1153                    MacroDef: lambda self, e: f"@DEF({self.sql(e.this)}, {self.sql(e.expression)})",
1154                    MacroFunc: _macro_func_sql,
1155                    MacroStrReplace: lambda self, e: f"@{self.sql(e.this)}",
1156                    MacroSQL: lambda self, e: f"@SQL({self.sql(e.this)})",
1157                    MacroVar: lambda self, e: f"@{e.name}",
1158                    Metric: lambda self, e: _sqlmesh_ddl_sql(self, e, "METRIC"),
1159                    Model: lambda self, e: _sqlmesh_ddl_sql(self, e, "MODEL"),
1160                    ModelKind: _model_kind_sql,
1161                    PythonCode: lambda self, e: self.expressions(e, sep="\n", indent=False),
1162                    StagedFilePath: lambda self, e: self.table_sql(e),
1163                    exp.Whens: _whens_sql,
1164                }
1165            )
1166        if MacroDef not in generator.WITH_SEPARATED_COMMENTS:
1167            generator.WITH_SEPARATED_COMMENTS = (
1168                *generator.WITH_SEPARATED_COMMENTS,
1169                Model,
1170                MacroDef,
1171            )
1172
1173        generator.UNWRAPPED_INTERVAL_VALUES = (
1174            *generator.UNWRAPPED_INTERVAL_VALUES,
1175            MacroStrReplace,
1176            MacroVar,
1177        )
1178
1179    _override(Parser, _parse_select)
1180    _override(Parser, _parse_statement)
1181    _override(Parser, _parse_join)
1182    _override(Parser, _parse_order)
1183    _override(Parser, _parse_where)
1184    _override(Parser, _parse_group)
1185    _override(Parser, _parse_with)
1186    _override(Parser, _parse_having)
1187    _override(Parser, _parse_limit)
1188    _override(Parser, _parse_value)
1189    _override(Parser, _parse_lambda)
1190    _override(Parser, _parse_types)
1191    _override(Parser, _parse_if)
1192    _override(TSQL.Parser, Parser._parse_if)
1193    _override(Parser, _parse_id_var)
1194    _override(Parser, _parse_interval_span)
1195    _override(Parser, _warn_unsupported)
1196    _override(Snowflake.Parser, _parse_table_parts)
1197
1198    # DuckDB's prefix absolute power operator `@` clashes with the macro syntax
1199    DuckDB.Parser.NO_PAREN_FUNCTION_PARSERS.pop("@", None)
1200
1201
1202def select_from_values(
1203    values: t.List[t.Tuple[t.Any, ...]],
1204    columns_to_types: t.Dict[str, exp.DataType],
1205    batch_size: int = 0,
1206    alias: str = "t",
1207) -> t.Iterator[exp.Select]:
1208    """Generate a VALUES expression that has a select wrapped around it to cast the values to their correct types.
1209
1210    Args:
1211        values: List of values to use for the VALUES expression.
1212        columns_to_types: Mapping of column names to types to assign to the values.
1213        batch_size: The maximum number of tuples per batches. Defaults to sys.maxsize if <= 0.
1214        alias: The alias to assign to the values expression. If not provided then will default to "t"
1215
1216    Returns:
1217        This method operates as a generator and yields a VALUES expression.
1218    """
1219    if batch_size <= 0:
1220        batch_size = sys.maxsize
1221    num_rows = len(values)
1222    for i in range(0, num_rows, batch_size):
1223        yield select_from_values_for_batch_range(
1224            values=values,
1225            target_columns_to_types=columns_to_types,
1226            batch_start=i,
1227            batch_end=min(i + batch_size, num_rows),
1228            alias=alias,
1229        )
1230
1231
1232def select_from_values_for_batch_range(
1233    values: t.List[t.Tuple[t.Any, ...]],
1234    target_columns_to_types: t.Dict[str, exp.DataType],
1235    batch_start: int,
1236    batch_end: int,
1237    alias: str = "t",
1238    source_columns: t.Optional[t.List[str]] = None,
1239) -> exp.Select:
1240    source_columns = source_columns or list(target_columns_to_types)
1241    source_columns_to_types = get_source_columns_to_types(target_columns_to_types, source_columns)
1242
1243    if not values:
1244        # Ensures we don't generate an empty VALUES clause & forces a zero-row output
1245        where = exp.false()
1246        expressions = [
1247            tuple(exp.cast(exp.null(), to=kind) for kind in source_columns_to_types.values())
1248        ]
1249    else:
1250        where = None
1251        expressions = [
1252            tuple(transform_values(v, source_columns_to_types))
1253            for v in values[batch_start:batch_end]
1254        ]
1255
1256    values_exp = exp.values(expressions, alias=alias, columns=source_columns_to_types)
1257    if values:
1258        # BigQuery crashes on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([NULL]) AS x`, but not
1259        # on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([CAST(NULL AS TIMESTAMP)]) AS x`. This
1260        # ensures nulls under the `Values` expression are cast to avoid similar issues.
1261        for value, kind in zip(
1262            values_exp.expressions[0].expressions, source_columns_to_types.values()
1263        ):
1264            if isinstance(value, exp.Null):
1265                value.replace(exp.cast(value, to=kind))
1266
1267    casted_columns = [
1268        exp.alias_(
1269            exp.cast(
1270                exp.column(column) if column in source_columns_to_types else exp.Null(), to=kind
1271            ),
1272            column,
1273            copy=False,
1274        )
1275        for column, kind in target_columns_to_types.items()
1276    ]
1277    return exp.select(*casted_columns).from_(values_exp, copy=False).where(where, copy=False)
1278
1279
1280def pandas_to_sql(
1281    df: pd.DataFrame,
1282    columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
1283    batch_size: int = 0,
1284    alias: str = "t",
1285) -> t.Iterator[exp.Select]:
1286    """Convert a pandas dataframe into a VALUES sql statement.
1287
1288    Args:
1289        df: A pandas dataframe to convert.
1290        columns_to_types: Mapping of column names to types to assign to the values.
1291        batch_size: The maximum number of tuples per batches. Defaults to sys.maxsize if <= 0.
1292        alias: The alias to assign to the values expression. If not provided then will default to "t"
1293
1294    Returns:
1295        This method operates as a generator and yields a VALUES expression.
1296    """
1297    yield from select_from_values(
1298        values=list(df.itertuples(index=False, name=None)),
1299        columns_to_types=columns_to_types or columns_to_types_from_df(df),
1300        batch_size=batch_size,
1301        alias=alias,
1302    )
1303
1304
1305def set_default_catalog(
1306    table: str | exp.Table,
1307    default_catalog: t.Optional[str],
1308) -> exp.Table:
1309    table = exp.to_table(table)
1310
1311    if default_catalog and not table.catalog and table.db:
1312        table.set("catalog", exp.parse_identifier(default_catalog))
1313
1314    return table
1315
1316
1317@lru_cache(maxsize=16384)
1318def normalize_model_name(
1319    table: str | exp.Table | exp.Column,
1320    default_catalog: t.Optional[str],
1321    dialect: DialectType = None,
1322) -> str:
1323    if isinstance(table, exp.Column):
1324        table = exp.table_(table.this, db=table.args.get("table"), catalog=table.args.get("db"))
1325    else:
1326        # We are relying on sqlglot's flexible parsing here to accept quotes from other dialects.
1327        # Ex: I have a a normalized name of '"my_table"' but the dialect is spark and therefore we should
1328        # expect spark quotes to be backticks ('`') instead of double quotes ('"'). sqlglot today is flexible
1329        # and will still parse this correctly and we rely on that.
1330        table = exp.to_table(table, dialect=dialect)
1331
1332    table = set_default_catalog(table, default_catalog)
1333    # An alternative way to do this is the following: exp.table_name(table, dialect=dialect, identify=True)
1334    # This though would result in the names being normalized to the target dialect AND the quotes while the below
1335    # approach just normalizes the names.
1336    # By just normalizing names and using sqlglot dialect for quotes this makes it easier for dialects that have
1337    # compatible normalization strategies but incompatible quoting to still work together without user hassle
1338    return exp.table_name(normalize_identifiers(table, dialect=dialect), identify=True)
1339
1340
1341def find_tables(
1342    expression: exp.Expr, default_catalog: t.Optional[str], dialect: DialectType = None
1343) -> t.Set[str]:
1344    """Find all tables referenced in a query.
1345
1346    Caches the result in the meta field 'tables'.
1347
1348    Args:
1349        expressions: The query to find the tables in.
1350        dialect: The dialect to use for normalization of table names.
1351
1352    Returns:
1353        A Set of all the table names.
1354    """
1355    if TABLES_META not in expression.meta:
1356        expression.meta[TABLES_META] = {
1357            normalize_model_name(table, default_catalog=default_catalog, dialect=dialect)
1358            for scope in traverse_scope(expression)
1359            for table in scope.tables
1360            if table.name and table.name not in scope.cte_sources
1361        }
1362    return expression.meta[TABLES_META]
1363
1364
1365def add_table(node: exp.Expr, table: str) -> exp.Expr:
1366    """Add a table to all columns in an expression."""
1367
1368    def _transform(node: exp.Expr) -> exp.Expr:
1369        if isinstance(node, exp.Column) and not node.table:
1370            return exp.column(node.this, table=table)
1371        if isinstance(node, exp.Identifier):
1372            return exp.column(node, table=table)
1373        return node
1374
1375    return node.transform(_transform)
1376
1377
1378def transform_values(
1379    values: t.Tuple[t.Any, ...], columns_to_types: t.Dict[str, exp.DataType]
1380) -> t.Iterator[t.Any]:
1381    """Perform transformations on values given columns_to_types."""
1382
1383    def _transform_value(value: t.Any, dtype: exp.DataType) -> t.Any:
1384        if (
1385            isinstance(value, list)
1386            and dtype.is_type(*exp.DataType.ARRAY_TYPES)
1387            and len(dtype.expressions) == 1
1388        ):
1389            element_type = dtype.expressions[0]
1390            return exp.convert([_transform_value(v, element_type) for v in value])
1391
1392        if (
1393            isinstance(value, dict)
1394            and dtype.is_type(*exp.DataType.STRUCT_TYPES)
1395            and len(value) == len(dtype.expressions)
1396        ):
1397            expressions = []
1398            for (field_name, field_value), field_type in zip(value.items(), dtype.expressions):
1399                if isinstance(field_type, exp.ColumnDef):
1400                    field_type = field_type.kind
1401                else:
1402                    field_type = exp.DataType.build(exp.DataType.Type.UNKNOWN)
1403
1404                expressions.append(
1405                    exp.PropertyEQ(
1406                        this=exp.to_identifier(field_name),
1407                        expression=_transform_value(field_value, field_type),
1408                    )
1409                )
1410
1411            return exp.Struct(expressions=expressions)
1412
1413        if dtype.is_type(exp.DataType.Type.JSON):
1414            return exp.func("PARSE_JSON", f"'{value}'")
1415
1416        return exp.convert(value)
1417
1418    for col_value, col_type in zip(values, columns_to_types.values()):
1419        yield _transform_value(col_value, col_type)
1420
1421
1422def to_schema(sql_path: str | exp.Table, dialect: DialectType = None) -> exp.Table:
1423    if isinstance(sql_path, exp.Table) and sql_path.this is None:
1424        return sql_path
1425    table = exp.to_table(
1426        sql_path.copy() if isinstance(sql_path, exp.Table) else sql_path, dialect=dialect
1427    )
1428    table.set("catalog", table.args.get("db"))
1429    table.set("db", table.args.get("this"))
1430    table.set("this", None)
1431    return table
1432
1433
1434def schema_(
1435    db: exp.Identifier | str,
1436    catalog: t.Optional[exp.Identifier | str] = None,
1437    quoted: t.Optional[bool] = None,
1438) -> exp.Table:
1439    """Build a Schema.
1440
1441    Args:
1442        db: Database name.
1443        catalog: Catalog name.
1444        quoted: Whether to force quotes on the schema's identifiers.
1445
1446    Returns:
1447        The new Schema instance.
1448    """
1449    return exp.Table(
1450        this=None,
1451        db=exp.to_identifier(db, quoted=quoted) if db else None,
1452        catalog=exp.to_identifier(catalog, quoted=quoted) if catalog else None,
1453    )
1454
1455
1456def normalize_mapping_schema(schema: t.Dict, dialect: DialectType) -> MappingSchema:
1457    return MappingSchema(_unquote_schema(schema), dialect=dialect, normalize=False)
1458
1459
1460def _unquote_schema(schema: t.Dict) -> t.Dict:
1461    """SQLGlot schema expects unquoted normalized keys."""
1462    return {
1463        k.strip('"'): _unquote_schema(v) if isinstance(v, dict) else v for k, v in schema.items()
1464    }
1465
1466
1467@contextmanager
1468def normalize_and_quote(
1469    query: E, dialect: DialectType, default_catalog: t.Optional[str], quote: bool = True
1470) -> t.Iterator[E]:
1471    qualify_tables(query, catalog=default_catalog, dialect=dialect)
1472    normalize_identifiers(query, dialect=dialect)
1473    yield query
1474    if quote:
1475        quote_identifiers(query, dialect=dialect)
1476
1477
1478def interpret_expression(e: exp.Expr) -> exp.Expr | str | int | float | bool:
1479    if e.is_int:
1480        return int(e.this)
1481    if e.is_number:
1482        return float(e.this)
1483    if isinstance(e, (exp.Literal, exp.Boolean)):
1484        return e.this
1485    return e
1486
1487
1488def interpret_key_value_pairs(
1489    e: exp.Tuple,
1490) -> t.Dict[str, exp.Expr | str | int | float | bool]:
1491    return {i.this.name: interpret_expression(i.expression) for i in e.expressions}
1492
1493
1494def extract_func_call(
1495    v: exp.Expr, allow_tuples: bool = False
1496) -> t.Tuple[str, t.Dict[str, exp.Expr]]:
1497    kwargs = {}
1498
1499    if isinstance(v, exp.Anonymous):
1500        func = v.name
1501        args = v.expressions
1502    elif isinstance(v, exp.Func):
1503        func = v.sql_name()
1504        args = list(v.args.values())
1505    elif isinstance(v, exp.Paren):
1506        func = ""
1507        args = [v.this]
1508    elif isinstance(v, exp.Tuple):  # airflow only
1509        if not allow_tuples:
1510            raise ConfigError("Audit name is missing (eg. MY_AUDIT())")
1511
1512        func = ""
1513        args = v.expressions
1514    else:
1515        return v.name.lower(), {}
1516
1517    for arg in args:
1518        if not isinstance(arg, (exp.PropertyEQ, exp.EQ)):
1519            raise ConfigError(
1520                f"Function '{func}' must be called with key-value arguments like {func}(arg := value)."
1521            )
1522        kwargs[arg.left.name.lower()] = arg.right
1523    return func.lower(), kwargs
1524
1525
1526def extract_function_calls(func_calls: t.Any, allow_tuples: bool = False) -> t.Any:
1527    """Used for extracting function calls for signals or audits."""
1528
1529    if isinstance(func_calls, (exp.Tuple, exp.Array)):
1530        return [extract_func_call(i, allow_tuples=allow_tuples) for i in func_calls.expressions]
1531    if isinstance(func_calls, exp.Paren):
1532        return [extract_func_call(func_calls.this, allow_tuples=allow_tuples)]
1533    if isinstance(func_calls, exp.Expr):
1534        return [extract_func_call(func_calls, allow_tuples=allow_tuples)]
1535    if isinstance(func_calls, list):
1536        function_calls = []
1537        for entry in func_calls:
1538            if isinstance(entry, dict):
1539                args = entry
1540                name = "" if allow_tuples else entry.pop("name")
1541            elif isinstance(entry, (tuple, list)):
1542                name, args = entry
1543            else:
1544                raise ConfigError(f"Audit must be a dictionary or named tuple. Got {entry}.")
1545
1546            function_calls.append(
1547                (
1548                    name.lower(),
1549                    {
1550                        key: parse_one(value) if isinstance(value, str) else value
1551                        for key, value in args.items()
1552                    },
1553                )
1554            )
1555
1556        return function_calls
1557
1558    return func_calls or []
1559
1560
1561def is_meta_expression(v: t.Any) -> bool:
1562    return isinstance(v, (Audit, Metric, Model))
1563
1564
1565def replace_merge_table_aliases(expression: exp.Expr, dialect: t.Optional[str] = None) -> exp.Expr:
1566    """
1567    Resolves references from the "source" and "target" tables (or their DBT equivalents)
1568    with the corresponding SQLMesh merge aliases (MERGE_SOURCE_ALIAS and MERGE_TARGET_ALIAS)
1569    """
1570    from sqlmesh.core.engine_adapter.base import MERGE_SOURCE_ALIAS, MERGE_TARGET_ALIAS
1571
1572    if isinstance(expression, exp.Column) and (first_part := expression.parts[0]):
1573        if first_part.this.lower() in ("target", "dbt_internal_dest", "__merge_target__"):
1574            first_part.replace(exp.to_identifier(MERGE_TARGET_ALIAS, quoted=True))
1575        elif first_part.this.lower() in ("source", "dbt_internal_source", "__merge_source__"):
1576            first_part.replace(exp.to_identifier(MERGE_SOURCE_ALIAS, quoted=True))
1577
1578    return expression