sqlmesh.core.macros

   1from __future__ import annotations
   2
   3import inspect
   4import sys
   5import types
   6import typing as t
   7from enum import Enum
   8from functools import lru_cache, reduce
   9from itertools import chain
  10from pathlib import Path
  11from string import Template
  12from datetime import datetime, date
  13
  14import sqlglot
  15from sqlglot import Generator, exp, parse_one
  16from sqlglot.executor.env import ENV
  17from sqlglot.executor.python import Python
  18from sqlglot.generators.python import PythonGenerator
  19from sqlglot.helper import csv, ensure_collection
  20from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
  21from sqlglot.schema import MappingSchema
  22
  23from sqlmesh.core import constants as c
  24from sqlmesh.core.dialect import (
  25    SQLMESH_MACRO_PREFIX,
  26    Dialect,
  27    MacroDef,
  28    MacroFunc,
  29    MacroSQL,
  30    MacroStrReplace,
  31    MacroVar,
  32    StagedFilePath,
  33    normalize_model_name,
  34)
  35from sqlmesh.utils import (
  36    DECORATOR_RETURN_TYPE,
  37    UniqueKeyDict,
  38    columns_to_types_all_known,
  39    registry_decorator,
  40)
  41from sqlmesh.utils.date import DatetimeRanges, to_datetime, to_date
  42from sqlmesh.utils.errors import MacroEvalError, SQLMeshError
  43from sqlmesh.utils.metaprogramming import (
  44    Executable,
  45    SqlValue,
  46    format_evaluated_code_exception,
  47    prepare_env,
  48)
  49
  50if t.TYPE_CHECKING:
  51    from sqlglot.dialects.dialect import DialectType
  52    from sqlmesh.core._typing import TableName
  53    from sqlmesh.core.engine_adapter import EngineAdapter
  54    from sqlmesh.core.snapshot import Snapshot
  55    from sqlmesh.core.environment import EnvironmentNamingInfo
  56
  57
  58if sys.version_info >= (3, 10):
  59    UNION_TYPES = (t.Union, types.UnionType)
  60else:
  61    UNION_TYPES = (t.Union,)
  62
  63
  64class RuntimeStage(Enum):
  65    LOADING = "loading"
  66    CREATING = "creating"
  67    EVALUATING = "evaluating"
  68    PROMOTING = "promoting"
  69    DEMOTING = "demoting"
  70    AUDITING = "auditing"
  71    TESTING = "testing"
  72    BEFORE_ALL = "before_all"
  73    AFTER_ALL = "after_all"
  74
  75
  76class MacroStrTemplate(Template):
  77    delimiter = SQLMESH_MACRO_PREFIX
  78
  79
  80EXPRESSIONS_NAME_MAP = {}
  81SQL = t.NewType("SQL", str)
  82
  83
  84@lru_cache()
  85def get_supported_types() -> t.Dict[str, t.Any]:
  86    from sqlmesh.core.context import ExecutionContext
  87
  88    return {
  89        "t": t,
  90        "typing": t,
  91        "List": t.List,
  92        "Tuple": t.Tuple,
  93        "Union": t.Union,
  94        "DatetimeRanges": DatetimeRanges,
  95        "exp": exp,
  96        "SQL": SQL,
  97        "MacroEvaluator": MacroEvaluator,
  98        "ExecutionContext": ExecutionContext,
  99    }
 100
 101
 102for klass in sqlglot.Parser.EXPRESSION_PARSERS:
 103    name = klass if isinstance(klass, str) else klass.__name__  # type: ignore
 104    EXPRESSIONS_NAME_MAP[name.lower()] = name
 105
 106
 107def _macro_sql(sql: str, into: t.Optional[str] = None) -> str:
 108    args = [_macro_str_replace(sql)]
 109    if into in EXPRESSIONS_NAME_MAP:
 110        args.append(f"into=exp.{EXPRESSIONS_NAME_MAP[into]}")
 111    return f"self.parse_one({', '.join(args)})"
 112
 113
 114def _macro_func_sql(self: Generator, e: exp.Expr) -> str:
 115    func = e.this
 116
 117    if isinstance(func, exp.Anonymous):
 118        return f"""self.send({csv("'" + func.name + "'", self.expressions(func))})"""
 119    return self.sql(func)
 120
 121
 122def _macro_str_replace(text: str) -> str:
 123    """Stringifies python code for variable replacement
 124    Args:
 125        text: text string
 126    Returns:
 127        Stringified python code to execute variable replacement
 128    """
 129    return f"self.template({text}, locals())"
 130
 131
 132class CaseInsensitiveMapping(t.Dict[str, t.Any]):
 133    def __init__(self, data: t.Dict[str, t.Any]) -> None:
 134        super().__init__(data)
 135
 136    def __getitem__(self, key: str) -> t.Any:
 137        return super().__getitem__(key.lower())
 138
 139    def get(self, key: str, default: t.Any = None, /) -> t.Any:
 140        return super().get(key.lower(), default)
 141
 142
 143class MacroDialect(Python):
 144    class Generator(PythonGenerator):
 145        TRANSFORMS = {
 146            **PythonGenerator.TRANSFORMS,
 147            exp.Column: lambda self, e: f"exp.to_column('{self.sql(e, 'this')}')",
 148            exp.Lambda: lambda self, e: f"lambda {self.expressions(e)}: {self.sql(e, 'this')}",
 149            MacroFunc: _macro_func_sql,
 150            MacroSQL: lambda self, e: _macro_sql(self.sql(e, "this"), e.args.get("into")),
 151            MacroStrReplace: lambda self, e: _macro_str_replace(self.sql(e, "this")),
 152        }
 153
 154
 155class MacroEvaluator:
 156    """The class responsible for evaluating SQLMesh Macros/SQL.
 157
 158    SQLMesh supports special preprocessed SQL prefixed with `@`. Although it provides similar power to
 159    traditional methods like string templating, there is semantic understanding of SQL which prevents
 160    common errors like leading/trailing commas, syntax errors, etc.
 161
 162    SQLMesh SQL allows for macro variables and macro functions. Macro variables take the form of @variable. These are used for variable substitution.
 163
 164    SELECT * FROM foo WHERE ds BETWEEN @start_date AND @end_date
 165
 166    Macro variables can be defined with a special macro function.
 167
 168    @DEF(start_date, '2021-01-01')
 169
 170    Args:
 171        dialect: Dialect of the SQL to evaluate.
 172        python_env: Serialized Python environment.
 173    """
 174
 175    def __init__(
 176        self,
 177        dialect: DialectType = "",
 178        python_env: t.Optional[t.Dict[str, Executable]] = None,
 179        schema: t.Optional[MappingSchema] = None,
 180        runtime_stage: RuntimeStage = RuntimeStage.LOADING,
 181        resolve_table: t.Optional[t.Callable[[str | exp.Table], str]] = None,
 182        resolve_tables: t.Optional[t.Callable[[exp.Expr], exp.Expr]] = None,
 183        snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
 184        default_catalog: t.Optional[str] = None,
 185        path: t.Optional[Path] = None,
 186        environment_naming_info: t.Optional[EnvironmentNamingInfo] = None,
 187        model_fqn: t.Optional[str] = None,
 188    ):
 189        self.dialect = dialect
 190        self.generator = MacroDialect().generator()
 191        self.locals: t.Dict[str, t.Any] = {
 192            "runtime_stage": runtime_stage.value,
 193            "default_catalog": default_catalog,
 194        }
 195        self.env = {
 196            **ENV,
 197            "self": self,
 198            "SQL": SQL,
 199            "MacroEvaluator": MacroEvaluator,
 200        }
 201        self.python_env = python_env or {}
 202        self.macros = {normalize_macro_name(k): v.func for k, v in macro.get_registry().items()}
 203        self.columns_to_types_called = False
 204        self.default_catalog = default_catalog
 205
 206        self._schema = schema
 207        self._resolve_table = resolve_table
 208        self._resolve_tables = resolve_tables
 209        self._snapshots = snapshots if snapshots is not None else {}
 210        self._path = path
 211        self._environment_naming_info = environment_naming_info
 212        self._model_fqn = model_fqn
 213
 214        prepare_env(self.python_env, self.env)
 215        for k, v in self.python_env.items():
 216            if v.is_definition:
 217                self.macros[normalize_macro_name(k)] = self.env[v.name or k]
 218            elif v.is_import and getattr(self.env.get(k), c.SQLMESH_MACRO, None):
 219                self.macros[normalize_macro_name(k)] = self.env[k]
 220            elif v.is_value:
 221                value = self.env[k]
 222                if k in (
 223                    c.SQLMESH_VARS,
 224                    c.SQLMESH_VARS_METADATA,
 225                    c.SQLMESH_BLUEPRINT_VARS,
 226                    c.SQLMESH_BLUEPRINT_VARS_METADATA,
 227                ):
 228                    value = {
 229                        var_name: (
 230                            self.parse_one(var_value.sql)
 231                            if isinstance(var_value, SqlValue)
 232                            else var_value
 233                        )
 234                        for var_name, var_value in value.items()
 235                    }
 236
 237                self.locals[k] = value
 238
 239    def send(
 240        self, name: str, *args: t.Any, **kwargs: t.Any
 241    ) -> t.Union[None, exp.Expr, t.List[exp.Expr]]:
 242        func = self.macros.get(normalize_macro_name(name))
 243
 244        if not callable(func):
 245            raise MacroEvalError(f"Macro '{name}' does not exist.")
 246
 247        try:
 248            return call_macro(
 249                func, self.dialect, self._path, provided_args=(self, *args), provided_kwargs=kwargs
 250            )  # type: ignore
 251        except Exception as e:
 252            raise MacroEvalError(
 253                f"An error occurred during evaluation of '{name}'\n\n"
 254                + format_evaluated_code_exception(e, self.python_env)
 255            )
 256
 257    def transform(self, expression: exp.Expr) -> exp.Expr | t.List[exp.Expr] | None:
 258        changed = False
 259
 260        def evaluate_macros(
 261            node: exp.Expr,
 262        ) -> exp.Expr | t.List[exp.Expr] | None:
 263            nonlocal changed
 264
 265            if isinstance(node, MacroVar):
 266                changed = True
 267                variables = self.variables
 268
 269                # This makes all variables case-insensitive, e.g. @X is the same as @x. We do this
 270                # for consistency, since `variables` and `blueprint_variables` are normalized.
 271                var_name = node.name.lower()
 272
 273                if var_name not in self.locals and var_name not in variables:
 274                    if not isinstance(node.parent, StagedFilePath):
 275                        raise SQLMeshError(f"Macro variable '{node.name}' is undefined.")
 276
 277                    return node
 278
 279                # Precedence order is locals (e.g. @DEF) > blueprint variables > config variables
 280                value = self.locals.get(var_name, variables.get(var_name))
 281                if isinstance(value, list):
 282                    return exp.convert(
 283                        tuple(self.transform(v) if isinstance(v, exp.Expr) else v for v in value)
 284                    )
 285
 286                return exp.convert(self.transform(value) if isinstance(value, exp.Expr) else value)
 287            if isinstance(node, exp.Identifier) and "@" in node.this:
 288                text = self.template(node.this, {})
 289                if node.this != text:
 290                    changed = True
 291                    return exp.to_identifier(text, quoted=node.quoted or None)
 292            if isinstance(node, MacroFunc):
 293                changed = True
 294                return self.evaluate(node)
 295            return node
 296
 297        transformed = exp.replace_tree(
 298            expression.copy(),
 299            evaluate_macros,  # type: ignore[arg-type]
 300            prune=lambda n: isinstance(n, exp.Lambda),
 301        )
 302
 303        if changed:
 304            # the transformations could have corrupted the ast, turning this into sql and reparsing ensures
 305            # that the ast is correct
 306            if isinstance(transformed, list):
 307                return [
 308                    self.parse_one(node.sql(dialect=self.dialect, copy=False))
 309                    for node in transformed
 310                ]
 311            if isinstance(transformed, exp.Expr):
 312                return self.parse_one(transformed.sql(dialect=self.dialect, copy=False))
 313
 314        return transformed
 315
 316    def template(self, text: t.Any, local_variables: t.Dict[str, t.Any]) -> str:
 317        """Substitute @vars with locals.
 318
 319        Args:
 320            text: The string to do substitition on.
 321            local_variables: Local variables in the context so that lambdas can be used.
 322
 323        Returns:
 324           The rendered string.
 325        """
 326        # We try to convert all variables into sqlglot expressions because they're going to be converted
 327        # into strings; in sql we don't convert strings because that would result in adding quotes
 328        base_mapping = {
 329            k.lower(): convert_sql(v, self.dialect)
 330            for k, v in chain(self.variables.items(), self.locals.items(), local_variables.items())
 331            if k.lower()
 332            not in (
 333                "engine_adapter",
 334                "snapshot",
 335            )
 336        }
 337        return MacroStrTemplate(str(text)).safe_substitute(CaseInsensitiveMapping(base_mapping))
 338
 339    def evaluate(self, node: MacroFunc) -> exp.Expr | t.List[exp.Expr] | None:
 340        if isinstance(node, MacroDef):
 341            if isinstance(node.expression, exp.Lambda):
 342                _, fn = _norm_var_arg_lambda(self, node.expression)
 343                self.macros[normalize_macro_name(node.name)] = lambda _, *args: fn(
 344                    args[0] if len(args) == 1 else exp.Tuple(expressions=list(args))
 345                )
 346            else:
 347                # Make variables defined through `@DEF` case-insensitive
 348                self.locals[node.name.lower()] = self.transform(node.expression)
 349
 350            return node
 351
 352        if isinstance(node, (MacroSQL, MacroStrReplace)):
 353            result: t.Optional[exp.Expr | t.List[exp.Expr]] = exp.convert(
 354                self.eval_expression(node)
 355            )
 356        else:
 357            func = t.cast(exp.Anonymous, node.this)
 358
 359            args = []
 360            kwargs = {}
 361            for e in func.expressions:
 362                if isinstance(e, exp.PropertyEQ):
 363                    kwargs[e.this.name] = e.expression
 364                else:
 365                    if kwargs:
 366                        raise MacroEvalError(
 367                            "Positional argument cannot follow keyword argument.\n  "
 368                            f"{func.sql(dialect=self.dialect)} at '{self._path}'"
 369                        )
 370
 371                    args.append(e)
 372
 373            result = self.send(func.name, *args, **kwargs)
 374
 375        if result is None:
 376            return None
 377
 378        if isinstance(result, (tuple, list)):
 379            result = [self.parse_one(item) for item in result if item is not None]
 380
 381            if (
 382                len(result) == 1
 383                and isinstance(result[0], (exp.Array, exp.Tuple))
 384                and node.find_ancestor(MacroFunc)
 385            ):
 386                """
 387                if:
 388                 - the output of evaluating this node is being passed as an argument to another macro function
 389                 - and that output is something that _norm_var_arg_lambda() will unpack into varargs
 390                   > (a list containing a single item of type exp.Tuple/exp.Array)
 391                then we will get inconsistent behaviour depending on if this node emits a list with a single item vs multiple items.
 392                
 393                In the first case, emitting a list containing a single array item will cause that array to get unpacked and its *members* passed to the calling macro
 394                In the second case, emitting a list containing multiple array items will cause each item to get passed as-is to the calling macro
 395                
 396                To prevent this inconsistency, we wrap this node output in an exp.Array so that _norm_var_arg_lambda() can "unpack" that into the
 397                actual argument we want to pass to the parent macro function
 398                
 399                Note we only do this for evaluation results that get passed as an argument to another macro, because when the final
 400                result is given to something like SELECT, we still want that to be unpacked into a list of items like:
 401                 - SELECT ARRAY(1), ARRAY(2)
 402                rather than a single item like:
 403                 - SELECT ARRAY(ARRAY(1), ARRAY(2))                
 404                """
 405                result = [exp.Array(expressions=result)]
 406        else:
 407            result = self.parse_one(result)
 408
 409        return result
 410
 411    def eval_expression(self, node: t.Any) -> t.Any:
 412        """Converts a SQLGlot expression into executable Python code and evals it.
 413
 414        If the node is not an expression, it will simply be returned.
 415
 416        Args:
 417            node: expression
 418        Returns:
 419            The return value of the evaled Python Code.
 420        """
 421        if not isinstance(node, exp.Expr):
 422            return node
 423        code = node.sql()
 424        try:
 425            code = self.generator.generate(node)
 426            return eval(code, self.env, self.locals)
 427        except Exception as e:
 428            raise MacroEvalError(
 429                f"Error trying to eval macro.\n\nGenerated code: {code}\n\nOriginal sql: {node}\n\n"
 430                + format_evaluated_code_exception(e, self.python_env)
 431            )
 432
 433    def parse_one(
 434        self, sql: str | exp.Expr, into: t.Optional[exp.IntoType] = None, **opts: t.Any
 435    ) -> exp.Expr:
 436        """Parses the given SQL string and returns a syntax tree for the first
 437        parsed SQL statement.
 438
 439        Args:
 440            sql: the SQL code or expression to parse.
 441            into: the Expression to parse into
 442            **opts: other options
 443
 444        Returns:
 445            Expression: the syntax tree for the first parsed statement
 446        """
 447        return sqlglot.maybe_parse(sql, dialect=self.dialect, into=into, **opts)
 448
 449    def columns_to_types(self, model_name: TableName | exp.Column) -> t.Dict[str, exp.DataType]:
 450        """Returns the columns-to-types mapping corresponding to the specified model."""
 451
 452        # We only return this dummy schema at load time, because if we don't actually know the
 453        # target model's schema at creation/evaluation time, returning a dummy schema could lead
 454        # to unintelligible errors when the query is executed
 455        if (self._schema is None or self._schema.empty) and self.runtime_stage == "loading":
 456            self.columns_to_types_called = True
 457            return {"__schema_unavailable_at_load__": exp.DataType.build("unknown")}
 458
 459        normalized_model_name = normalize_model_name(
 460            model_name,
 461            default_catalog=self.default_catalog,
 462            dialect=self.dialect,
 463        )
 464        model_name = exp.to_table(normalized_model_name)
 465
 466        columns_to_types = (
 467            self._schema.find(model_name, ensure_data_types=True) if self._schema else None
 468        )
 469        if columns_to_types is None:
 470            snapshot = self.get_snapshot(model_name)
 471            if snapshot and snapshot.node.is_model:
 472                columns_to_types = snapshot.node.columns_to_types  # type: ignore
 473
 474        if columns_to_types is None:
 475            raise SQLMeshError(f"Schema for model '{model_name}' can't be statically determined.")
 476
 477        return columns_to_types
 478
 479    def get_snapshot(self, model_name: TableName | exp.Column) -> t.Optional[Snapshot]:
 480        """Returns the snapshot that corresponds to the given model name."""
 481        return self._snapshots.get(
 482            normalize_model_name(
 483                model_name,
 484                default_catalog=self.default_catalog,
 485                dialect=self.dialect,
 486            )
 487        )
 488
 489    def resolve_table(self, table: str | exp.Table) -> str:
 490        """Gets the physical table name for a given model."""
 491        if not self._resolve_table:
 492            raise SQLMeshError(
 493                "Macro evaluator not properly initialized with resolve_table lambda."
 494            )
 495        return self._resolve_table(table)
 496
 497    def resolve_tables(self, query: exp.Expr) -> exp.Expr:
 498        """Resolves queries with references to SQLMesh model names to their physical tables."""
 499        if not self._resolve_tables:
 500            raise SQLMeshError(
 501                "Macro evaluator not properly initialized with resolve_tables lambda."
 502            )
 503        return self._resolve_tables(query)
 504
 505    @property
 506    def runtime_stage(self) -> RuntimeStage:
 507        """Returns the current runtime stage of the macro evaluation."""
 508        return self.locals["runtime_stage"]
 509
 510    @property
 511    def this_model(self) -> str:
 512        """Returns the resolved name of the surrounding model."""
 513        this_model = self.locals.get("this_model")
 514        if not this_model:
 515            raise SQLMeshError("Model name is not available in the macro evaluator.")
 516        return this_model.sql(dialect=self.dialect, identify=True, comments=False)
 517
 518    @property
 519    def this_model_fqn(self) -> str:
 520        if self._model_fqn is None:
 521            raise SQLMeshError("Model name is not available in the macro evaluator.")
 522        return self._model_fqn
 523
 524    @property
 525    def engine_adapter(self) -> EngineAdapter:
 526        engine_adapter = self.locals.get("engine_adapter")
 527        if not engine_adapter:
 528            raise SQLMeshError(
 529                "The engine adapter is not available while models are loading."
 530                " You can gate these calls by checking in Python: evaluator.runtime_stage != 'loading' or SQL: @runtime_stage <> 'loading'."
 531            )
 532        return self.locals["engine_adapter"]
 533
 534    @property
 535    def gateway(self) -> t.Optional[str]:
 536        """Returns the gateway name."""
 537        return self.var(c.GATEWAY)
 538
 539    @property
 540    def snapshots(self) -> t.Dict[str, Snapshot]:
 541        """Returns the snapshots if available."""
 542        return self._snapshots
 543
 544    @property
 545    def this_env(self) -> str:
 546        """Returns the name of the current environment in before after all."""
 547        if "this_env" not in self.locals:
 548            raise SQLMeshError("Environment name is only available in before_all and after_all")
 549        return self.locals["this_env"]
 550
 551    @property
 552    def schemas(self) -> t.List[str]:
 553        """Returns the schemas of the current environment in before after all macros."""
 554        if "schemas" not in self.locals:
 555            raise SQLMeshError("Schemas are only available in before_all and after_all")
 556        return self.locals["schemas"]
 557
 558    @property
 559    def views(self) -> t.List[str]:
 560        """Returns the views of the current environment in before after all macros."""
 561        if "views" not in self.locals:
 562            raise SQLMeshError("Views are only available in before_all and after_all")
 563        return self.locals["views"]
 564
 565    def var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
 566        """Returns the value of the specified variable, or the default value if it doesn't exist."""
 567        return {
 568            **(self.locals.get(c.SQLMESH_VARS) or {}),
 569            **(self.locals.get(c.SQLMESH_VARS_METADATA) or {}),
 570        }.get(var_name.lower(), default)
 571
 572    def blueprint_var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
 573        """Returns the value of the specified blueprint variable, or the default value if it doesn't exist."""
 574        return {
 575            **(self.locals.get(c.SQLMESH_BLUEPRINT_VARS) or {}),
 576            **(self.locals.get(c.SQLMESH_BLUEPRINT_VARS_METADATA) or {}),
 577        }.get(var_name.lower(), default)
 578
 579    @property
 580    def variables(self) -> t.Dict[str, t.Any]:
 581        return {
 582            **self.locals.get(c.SQLMESH_VARS, {}),
 583            **self.locals.get(c.SQLMESH_VARS_METADATA, {}),
 584            **self.locals.get(c.SQLMESH_BLUEPRINT_VARS, {}),
 585            **self.locals.get(c.SQLMESH_BLUEPRINT_VARS_METADATA, {}),
 586        }
 587
 588    def _coerce(self, expr: exp.Expr, typ: t.Any, strict: bool = False) -> t.Any:
 589        """Coerces the given expression to the specified type on a best-effort basis."""
 590        return _coerce(expr, typ, self.dialect, self._path, strict)
 591
 592
 593class macro(registry_decorator):
 594    """Specifies a function is a macro and registers it the global MACROS registry.
 595
 596    Registered macros can be referenced in SQL statements to make queries more dynamic/cleaner.
 597
 598    Example:
 599        from sqlglot import exp
 600        from sqlmesh.core.macros import MacroEvaluator, macro
 601
 602        @macro()
 603        def add_one(evaluator: MacroEvaluator, column: exp.Literal) -> exp.Add:
 604            return evaluator.parse_one(f"{column} + 1")
 605
 606    Args:
 607        name: A custom name for the macro, the default is the name of the function.
 608    """
 609
 610    registry_name = "macros"
 611
 612    def __init__(self, *args: t.Any, metadata_only: bool = False, **kwargs: t.Any) -> None:
 613        super().__init__(*args, **kwargs)
 614        self.metadata_only = metadata_only
 615
 616    def __call__(
 617        self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
 618    ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
 619        if self.metadata_only:
 620            setattr(func, c.SQLMESH_METADATA, self.metadata_only)
 621        wrapper = super().__call__(func)
 622
 623        # This is used to identify macros at runtime to unwrap during serialization.
 624        setattr(wrapper, c.SQLMESH_MACRO, True)
 625        return wrapper
 626
 627
 628ExecutableOrMacro = t.Union[Executable, macro]
 629MacroRegistry = UniqueKeyDict[str, ExecutableOrMacro]
 630
 631
 632def _norm_var_arg_lambda(
 633    evaluator: MacroEvaluator, func: exp.Lambda, *items: t.Any
 634) -> t.Tuple[t.Iterable, t.Callable]:
 635    """
 636    Converts sql literal array and lambda into actual python iterable + callable.
 637
 638    In order to support expressions like @EACH([a, b, c], x -> @SQL('@x')), the lambda var x
 639    needs be passed to the local state.
 640
 641    Args:
 642        evaluator: MacroEvaluator that invoked the macro
 643        func: Lambda SQLGlot expression.
 644        items: Array or items of SQLGlot expressions.
 645    """
 646
 647    def substitute(
 648        node: exp.Expr, args: t.Dict[str, exp.Expr]
 649    ) -> exp.Expr | t.List[exp.Expr] | None:
 650        if isinstance(node, (exp.Identifier, exp.Var)):
 651            name = node.name.lower()
 652            if name in args:
 653                return args[name].copy()
 654            if not isinstance(node.parent, exp.Column):
 655                if name in evaluator.locals:
 656                    return exp.convert(evaluator.locals[name])
 657            if SQLMESH_MACRO_PREFIX in node.name:
 658                return node.__class__(
 659                    this=evaluator.template(node.name, {k: v.name for k, v in args.items()})
 660                )
 661        elif isinstance(node, MacroFunc):
 662            local_copy = evaluator.locals.copy()
 663            evaluator.locals.update(args)
 664            result = evaluator.transform(node)
 665            evaluator.locals = local_copy
 666            return result
 667        return node
 668
 669    if len(items) == 1:
 670        item = items[0]
 671        expressions = (
 672            item.expressions
 673            if isinstance(item, (exp.Array, exp.Tuple))
 674            else [item.this]
 675            if isinstance(item, exp.Paren)
 676            else item
 677        )
 678    else:
 679        expressions = items
 680
 681    if not callable(func):
 682        return expressions, lambda args: func.this.transform(
 683            substitute,
 684            {
 685                expression.name.lower(): arg
 686                for expression, arg in zip(
 687                    func.expressions, args.expressions if isinstance(args, exp.Tuple) else [args]
 688                )
 689            },
 690        )
 691
 692    return expressions, func
 693
 694
 695@macro()
 696def each(
 697    evaluator: MacroEvaluator,
 698    *args: t.Any,
 699) -> t.List[t.Any]:
 700    """Iterates through items calling func on each.
 701
 702    If a func call on item returns None, it will be excluded from the list.
 703
 704    Args:
 705        evaluator: MacroEvaluator that invoked the macro
 706        args: The last argument should be a lambda of the form x -> x +1. The first argument can be
 707            an Array or var args can be used.
 708
 709    Returns:
 710        A list of items that is the result of func
 711    """
 712    *items, func = args
 713    items, func = _norm_var_arg_lambda(evaluator, func, *items)  # type: ignore
 714    return [item for item in map(func, ensure_collection(items)) if item is not None]
 715
 716
 717@macro("IF")
 718def if_(
 719    evaluator: MacroEvaluator,
 720    condition: t.Any,
 721    true: t.Any,
 722    false: t.Any = None,
 723) -> t.Any:
 724    """Evaluates a given condition and returns the second argument if true or else the third argument.
 725
 726    If false is not passed in, the default return value will be None.
 727
 728    Example:
 729        >>> from sqlglot import parse_one
 730        >>> from sqlmesh.core.macros import MacroEvaluator
 731        >>> MacroEvaluator().transform(parse_one("@IF('a' = 1, a, b)")).sql()
 732        'b'
 733
 734        >>> MacroEvaluator().transform(parse_one("@IF('a' = 1, a)"))
 735    """
 736
 737    if evaluator.eval_expression(condition):
 738        return true
 739    return false
 740
 741
 742@macro("REDUCE")
 743def reduce_(evaluator: MacroEvaluator, *args: t.Any) -> t.Any:
 744    """Iterates through items applying provided function that takes two arguments
 745    cumulatively to the items of iterable items, from left to right, so as to reduce
 746    the iterable to a single item.
 747
 748    Example:
 749        >>> from sqlglot import parse_one
 750        >>> from sqlmesh.core.macros import MacroEvaluator
 751        >>> sql = "@SQL(@REDUCE([100, 200, 300, 400], (x, y) -> x + y))"
 752        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 753        '1000'
 754
 755    Args:
 756        evaluator: MacroEvaluator that invoked the macro
 757        args: The last argument should be a lambda of the form (x, y) -> x + y. The first argument can be
 758            an Array or var args can be used.
 759    Returns:
 760        A single item that is the result of applying func cumulatively to items
 761    """
 762    *items, func = args
 763    items, func = _norm_var_arg_lambda(evaluator, func, *items)  # type: ignore
 764    return reduce(lambda a, b: func(exp.Tuple(expressions=[a, b])), ensure_collection(items))
 765
 766
 767@macro("FILTER")
 768def filter_(evaluator: MacroEvaluator, *args: t.Any) -> t.List[t.Any]:
 769    """Iterates through items, applying provided function to each item and removing
 770    all items where the function returns False
 771
 772    Example:
 773        >>> from sqlglot import parse_one
 774        >>> from sqlmesh.core.macros import MacroEvaluator
 775        >>> sql = "@REDUCE(@FILTER([1, 2, 3], x -> x > 1), (x, y) -> x + y)"
 776        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 777        '2 + 3'
 778
 779        >>> sql = "@EVAL(@REDUCE(@FILTER([1, 2, 3], x -> x > 1), (x, y) -> x + y))"
 780        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 781        '5'
 782
 783    Args:
 784        evaluator: MacroEvaluator that invoked the macro
 785        args: The last argument should be a lambda of the form x -> x > 1. The first argument can be
 786            an Array or var args can be used.
 787    Returns:
 788        The items for which the func returned True
 789    """
 790    *items, func = args
 791    items, func = _norm_var_arg_lambda(evaluator, func, *items)  # type: ignore
 792    return list(filter(lambda arg: evaluator.eval_expression(func(arg)), items))
 793
 794
 795def _optional_expression(
 796    evaluator: MacroEvaluator,
 797    condition: exp.Condition,
 798    expression: exp.Expr,
 799) -> t.Optional[exp.Expr]:
 800    """Inserts expression when the condition is True
 801
 802    The following examples express the usage of this function in the context of the macros which wrap it.
 803
 804    Examples:
 805        >>> from sqlglot import parse_one
 806        >>> from sqlmesh.core.macros import MacroEvaluator
 807        >>> sql = "@WITH(True) all_cities as (select * from city) select all_cities"
 808        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 809        'WITH all_cities AS (SELECT * FROM city) SELECT all_cities'
 810        >>> sql = "select * from city left outer @JOIN(True) country on city.country = country.name"
 811        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 812        'SELECT * FROM city LEFT OUTER JOIN country ON city.country = country.name'
 813        >>> sql = "select * from city @GROUP_BY(True) country, population"
 814        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 815        'SELECT * FROM city GROUP BY country, population'
 816        >>> sql = "select * from city group by country @HAVING(True) population > 100 and country = 'Mexico'"
 817        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 818        "SELECT * FROM city GROUP BY country HAVING population > 100 AND country = 'Mexico'"
 819
 820    Args:
 821        evaluator: MacroEvaluator that invoked the macro
 822        condition: Condition expression
 823        expression: SQL expression
 824    Returns:
 825        Expression if the conditional is True; otherwise None
 826    """
 827    return expression if evaluator.eval_expression(condition) else None
 828
 829
 830with_ = macro("WITH")(_optional_expression)
 831join = macro("JOIN")(_optional_expression)
 832where = macro("WHERE")(_optional_expression)
 833group_by = macro("GROUP_BY")(_optional_expression)
 834having = macro("HAVING")(_optional_expression)
 835order_by = macro("ORDER_BY")(_optional_expression)
 836limit = macro("LIMIT")(_optional_expression)
 837
 838
 839@macro("eval")
 840def eval_(evaluator: MacroEvaluator, condition: exp.Condition) -> t.Any:
 841    """Evaluate the given condition in a Python/SQL interpretor.
 842
 843    Example:
 844        >>> from sqlglot import parse_one
 845        >>> from sqlmesh.core.macros import MacroEvaluator
 846        >>> sql = "@EVAL(1 + 1)"
 847        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 848        '2'
 849    """
 850    return evaluator.eval_expression(condition)
 851
 852
 853# macros with union types need to use t.Union since | isn't available until 3.9
 854@macro()
 855def star(
 856    evaluator: MacroEvaluator,
 857    relation: exp.Table,
 858    alias: exp.Column = t.cast(exp.Column, exp.column("")),
 859    exclude: t.Union[exp.Array, exp.Tuple] = exp.Tuple(expressions=[]),
 860    prefix: exp.Literal = exp.Literal.string(""),
 861    suffix: exp.Literal = exp.Literal.string(""),
 862    quote_identifiers: exp.Boolean = exp.true(),
 863    except_: t.Union[exp.Array, exp.Tuple] = exp.Tuple(expressions=[]),
 864) -> t.List[exp.Expr]:
 865    """Returns a list of projections for the given relation.
 866
 867    Args:
 868        evaluator: MacroEvaluator that invoked the macro
 869        relation: The relation to select star from
 870        alias: The alias of the relation
 871        exclude: Columns to exclude
 872        prefix: A prefix to use for all selections
 873        suffix: A suffix to use for all selections
 874        quote_identifiers: Whether or not quote the resulting aliases, defaults to true
 875        except_: Alias for exclude (TODO: deprecate this, update docs)
 876
 877    Returns:
 878        An array of columns.
 879
 880    Example:
 881        >>> from sqlglot import parse_one, exp
 882        >>> from sqlglot.schema import MappingSchema
 883        >>> from sqlmesh.core.macros import MacroEvaluator
 884        >>> sql = "SELECT @STAR(foo, bar, exclude := [c], prefix := 'baz_') FROM foo AS bar"
 885        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": exp.DataType.build("string"), "b": exp.DataType.build("string"), "c": exp.DataType.build("string"), "d": exp.DataType.build("int")}})).transform(parse_one(sql)).sql()
 886        'SELECT CAST("bar"."a" AS TEXT) AS "baz_a", CAST("bar"."b" AS TEXT) AS "baz_b", CAST("bar"."d" AS INT) AS "baz_d" FROM foo AS bar'
 887    """
 888    if alias and not isinstance(alias, (exp.Identifier, exp.Column)):
 889        raise SQLMeshError(f"Invalid alias '{alias}'. Expected an identifier.")
 890    if exclude and not isinstance(exclude, (exp.Array, exp.Tuple)):
 891        raise SQLMeshError(f"Invalid exclude '{exclude}'. Expected an array.")
 892    if except_ != exp.tuple_():
 893        from sqlmesh.core.console import get_console
 894
 895        get_console().log_warning(
 896            "The 'except_' argument in @STAR will soon be deprecated. Use 'exclude' instead."
 897        )
 898        if not isinstance(exclude, (exp.Array, exp.Tuple)):
 899            raise SQLMeshError(f"Invalid exclude_ '{exclude}'. Expected an array.")
 900    if prefix and not isinstance(prefix, exp.Literal):
 901        raise SQLMeshError(f"Invalid prefix '{prefix}'. Expected a literal.")
 902    if suffix and not isinstance(suffix, exp.Literal):
 903        raise SQLMeshError(f"Invalid suffix '{suffix}'. Expected a literal.")
 904    if not isinstance(quote_identifiers, exp.Boolean):
 905        raise SQLMeshError(f"Invalid quote_identifiers '{quote_identifiers}'. Expected a boolean.")
 906
 907    excluded_names = {
 908        normalize_identifiers(excluded, dialect=evaluator.dialect).name
 909        for excluded in exclude.expressions or except_.expressions
 910    }
 911    quoted = quote_identifiers.this
 912    table_identifier = normalize_identifiers(
 913        alias if alias.name else relation, dialect=evaluator.dialect
 914    ).name
 915
 916    columns_to_types = {
 917        k: v for k, v in evaluator.columns_to_types(relation).items() if k not in excluded_names
 918    }
 919    if columns_to_types_all_known(columns_to_types):
 920        return [
 921            exp.cast(
 922                exp.column(column, table=table_identifier, quoted=quoted),
 923                dtype,
 924                dialect=evaluator.dialect,
 925            ).as_(f"{prefix.this}{column}{suffix.this}", quoted=quoted)
 926            for column, dtype in columns_to_types.items()
 927        ]
 928    return [
 929        exp.column(column, table=table_identifier, quoted=quoted).as_(
 930            f"{prefix.this}{column}{suffix.this}", quoted=quoted
 931        )
 932        for column, type_ in columns_to_types.items()
 933    ]
 934
 935
 936@macro()
 937def generate_surrogate_key(
 938    evaluator: MacroEvaluator,
 939    *fields: exp.Expr,
 940    hash_function: exp.Literal = exp.Literal.string("MD5"),
 941) -> exp.Func:
 942    """Generates a surrogate key (string) for the given fields.
 943
 944    Example:
 945        >>> from sqlglot import parse_one
 946        >>> from sqlmesh.core.macros import MacroEvaluator
 947        >>>
 948        >>> sql = "SELECT @GENERATE_SURROGATE_KEY(a, b, c) FROM foo"
 949        >>> MacroEvaluator(dialect="bigquery").transform(parse_one(sql, dialect="bigquery")).sql("bigquery")
 950        "SELECT TO_HEX(MD5(CONCAT(COALESCE(CAST(a AS STRING), '_sqlmesh_surrogate_key_null_'), '|', COALESCE(CAST(b AS STRING), '_sqlmesh_surrogate_key_null_'), '|', COALESCE(CAST(c AS STRING), '_sqlmesh_surrogate_key_null_')))) FROM foo"
 951        >>>
 952        >>> sql = "SELECT @GENERATE_SURROGATE_KEY(a, b, c, hash_function := 'SHA256') FROM foo"
 953        >>> MacroEvaluator(dialect="bigquery").transform(parse_one(sql, dialect="bigquery")).sql("bigquery")
 954        "SELECT SHA256(CONCAT(COALESCE(CAST(a AS STRING), '_sqlmesh_surrogate_key_null_'), '|', COALESCE(CAST(b AS STRING), '_sqlmesh_surrogate_key_null_'), '|', COALESCE(CAST(c AS STRING), '_sqlmesh_surrogate_key_null_'))) FROM foo"
 955    """
 956    string_fields: t.List[exp.Expr] = []
 957    for i, field in enumerate(fields):
 958        if i > 0:
 959            string_fields.append(exp.Literal.string("|"))
 960        string_fields.append(
 961            exp.func(
 962                "COALESCE",
 963                exp.cast(field, exp.DataType.build("text")),
 964                exp.Literal.string("_sqlmesh_surrogate_key_null_"),
 965            )
 966        )
 967
 968    func = exp.func(
 969        hash_function.name,
 970        exp.func("CONCAT", *string_fields),
 971        dialect=evaluator.dialect,
 972    )
 973    if isinstance(func, exp.MD5Digest):
 974        func = exp.MD5(this=func.this)
 975
 976    return func
 977
 978
 979@macro()
 980def safe_add(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case:
 981    """Adds numbers together, substitutes nulls for 0s and only returns null if all fields are null.
 982
 983    Example:
 984        >>> from sqlglot import parse_one
 985        >>> from sqlmesh.core.macros import MacroEvaluator
 986        >>> sql = "SELECT @SAFE_ADD(a, b) FROM foo"
 987        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 988        'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) + COALESCE(b, 0) END FROM foo'
 989    """
 990    return (
 991        exp.Case()
 992        .when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null())
 993        .else_(reduce(lambda a, b: a + b, [exp.func("COALESCE", field, 0) for field in fields]))  # type: ignore
 994    )
 995
 996
 997@macro()
 998def safe_sub(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case:
 999    """Subtract numbers, substitutes nulls for 0s and only returns null if all fields are null.
1000
1001    Example:
1002        >>> from sqlglot import parse_one
1003        >>> from sqlmesh.core.macros import MacroEvaluator
1004        >>> sql = "SELECT @SAFE_SUB(a, b) FROM foo"
1005        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1006        'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) - COALESCE(b, 0) END FROM foo'
1007    """
1008    return (
1009        exp.Case()
1010        .when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null())
1011        .else_(reduce(lambda a, b: a - b, [exp.func("COALESCE", field, 0) for field in fields]))  # type: ignore
1012    )
1013
1014
1015@macro()
1016def safe_div(_: MacroEvaluator, numerator: exp.Expr, denominator: exp.Expr) -> exp.Div:
1017    """Divides numbers, returns null if the denominator is 0.
1018
1019    Example:
1020        >>> from sqlglot import parse_one
1021        >>> from sqlmesh.core.macros import MacroEvaluator
1022        >>> sql = "SELECT @SAFE_DIV(a, b) FROM foo"
1023        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1024        'SELECT a / NULLIF(b, 0) FROM foo'
1025    """
1026    return numerator / exp.func("NULLIF", denominator, 0)
1027
1028
1029@macro()
1030def union(
1031    evaluator: MacroEvaluator,
1032    *args: exp.Expr,
1033) -> exp.Query:
1034    """Returns a UNION of the given tables. Only choosing columns that have the same name and type.
1035
1036    Args:
1037        evaluator: MacroEvaluator that invoked the macro
1038        args: Variable arguments that can be:
1039            - First argument can be a condition (exp.Condition)
1040            - A union type ('ALL' or 'DISTINCT') as exp.Literal
1041            - Tables (exp.Table)
1042
1043    Example:
1044        >>> from sqlglot import parse_one
1045        >>> from sqlglot.schema import MappingSchema
1046        >>> from sqlmesh.core.macros import MacroEvaluator
1047        >>> sql = "@UNION('distinct', foo, bar)"
1048        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": "int", "b": "string", "c": "string"}, "bar": {"c": "string", "a": "int", "b": "int"}})).transform(parse_one(sql)).sql()
1049        'SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM foo UNION SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM bar'
1050        >>> sql = "@UNION(True, 'distinct', foo, bar)"
1051        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": "int", "b": "string", "c": "string"}, "bar": {"c": "string", "a": "int", "b": "int"}})).transform(parse_one(sql)).sql()
1052        'SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM foo UNION SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM bar'
1053    """
1054
1055    if not args:
1056        raise SQLMeshError("At least one table is required for the @UNION macro.")
1057
1058    arg_idx = 0
1059    # Check for condition
1060    condition = evaluator.eval_expression(args[arg_idx])
1061    if isinstance(condition, bool):
1062        arg_idx += 1
1063        if arg_idx >= len(args):
1064            raise SQLMeshError("Expected more arguments after the condition of the `@UNION` macro.")
1065
1066    # Check for union type
1067    type_ = exp.Literal.string("ALL")
1068    if isinstance(args[arg_idx], exp.Literal):
1069        type_ = args[arg_idx]  # type: ignore
1070        arg_idx += 1
1071    kind = type_.name.upper()
1072    if kind not in ("ALL", "DISTINCT"):
1073        raise SQLMeshError(f"Invalid type '{type_}'. Expected 'ALL' or 'DISTINCT'.")
1074
1075    # Remaining args should be tables
1076    tables = [
1077        exp.to_table(e.sql(evaluator.dialect), dialect=evaluator.dialect) for e in args[arg_idx:]
1078    ]
1079
1080    columns = {
1081        column
1082        for column, _ in reduce(
1083            lambda a, b: a & b,  # type: ignore
1084            (evaluator.columns_to_types(table).items() for table in tables),
1085        )
1086    }
1087
1088    projections = [
1089        exp.cast(column, type_, dialect=evaluator.dialect).as_(column)
1090        for column, type_ in evaluator.columns_to_types(tables[0]).items()
1091        if column in columns
1092    ]
1093
1094    # Skip the union if condition is False
1095    if condition == False:
1096        return exp.select(*projections).from_(tables[0])
1097
1098    return reduce(
1099        lambda a, b: a.union(b, distinct=kind == "DISTINCT"),  # type: ignore
1100        [exp.select(*projections).from_(t) for t in tables],
1101    )
1102
1103
1104@macro()
1105def haversine_distance(
1106    _: MacroEvaluator,
1107    lat1: exp.Expr,
1108    lon1: exp.Expr,
1109    lat2: exp.Expr,
1110    lon2: exp.Expr,
1111    unit: exp.Literal = exp.Literal.string("mi"),
1112) -> exp.Mul:
1113    """Returns the haversine distance between two points.
1114
1115    Example:
1116        >>> from sqlglot import parse_one
1117        >>> from sqlmesh.core.macros import MacroEvaluator
1118        >>> sql = "SELECT @HAVERSINE_DISTANCE(driver_y, driver_x, passenger_y, passenger_x, 'mi') FROM rides"
1119        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1120        'SELECT 7922 * ASIN(SQRT((POWER(SIN(RADIANS((passenger_y - driver_y) / 2)), 2)) + (COS(RADIANS(driver_y)) * COS(RADIANS(passenger_y)) * POWER(SIN(RADIANS((passenger_x - driver_x) / 2)), 2)))) * 1.0 FROM rides'
1121    """
1122    if unit.this == "mi":
1123        conversion_rate = 1.0
1124    elif unit.this == "km":
1125        conversion_rate = 1.60934
1126    else:
1127        raise SQLMeshError(f"Invalid unit '{unit}'. Expected 'mi' or 'km'.")
1128
1129    return (
1130        2
1131        * 3961
1132        * exp.func(
1133            "ASIN",
1134            exp.func(
1135                "SQRT",
1136                exp.func("POWER", exp.func("SIN", exp.func("RADIANS", (lat2 - lat1) / 2)), 2)
1137                + exp.func("COS", exp.func("RADIANS", lat1))
1138                * exp.func("COS", exp.func("RADIANS", lat2))
1139                * exp.func("POWER", exp.func("SIN", exp.func("RADIANS", (lon2 - lon1) / 2)), 2),
1140            ),
1141        )
1142        * conversion_rate
1143    )
1144
1145
1146@macro()
1147def pivot(
1148    evaluator: MacroEvaluator,
1149    column: SQL,
1150    values: t.List[exp.Expr],
1151    alias: bool = True,
1152    agg: exp.Expr = exp.Literal.string("SUM"),
1153    cmp: exp.Expr = exp.Literal.string("="),
1154    prefix: exp.Expr = exp.Literal.string(""),
1155    suffix: exp.Expr = exp.Literal.string(""),
1156    then_value: SQL = SQL("1"),
1157    else_value: SQL = SQL("0"),
1158    quote: bool = True,
1159    distinct: bool = False,
1160) -> t.List[exp.Expr]:
1161    """Returns a list of projections as a result of pivoting the given column on the given values.
1162
1163    Example:
1164        >>> from sqlglot import parse_one
1165        >>> from sqlmesh.core.macros import MacroEvaluator
1166        >>> sql = "SELECT date_day, @PIVOT(status, ['cancelled', 'completed']) FROM rides GROUP BY 1"
1167        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1168        'SELECT date_day, SUM(CASE WHEN status = \\'cancelled\\' THEN 1 ELSE 0 END) AS "cancelled", SUM(CASE WHEN status = \\'completed\\' THEN 1 ELSE 0 END) AS "completed" FROM rides GROUP BY 1'
1169        >>> sql = "SELECT @PIVOT(a, ['v'], then_value := tv, suffix := '_sfx', quote := FALSE)"
1170        >>> MacroEvaluator(dialect="bigquery").transform(parse_one(sql)).sql("bigquery")
1171        "SELECT SUM(CASE WHEN a = 'v' THEN tv ELSE 0 END) AS v_sfx"
1172    """
1173    aggregates: t.List[exp.Expr] = []
1174    for value in values:
1175        proj = f"{agg.name}("
1176        if distinct:
1177            proj += "DISTINCT "
1178
1179        proj += f"CASE WHEN {column} {cmp.name} {value.sql(evaluator.dialect)} THEN {then_value} ELSE {else_value} END) "
1180        node: exp.Expr = evaluator.parse_one(proj)
1181
1182        if alias:
1183            node = node.as_(
1184                f"{prefix.name}{value.name}{suffix.name}",
1185                quoted=quote,
1186                copy=False,
1187                dialect=evaluator.dialect,
1188            )
1189
1190        aggregates.append(node)
1191
1192    return aggregates
1193
1194
1195@macro("AND")
1196def and_(evaluator: MacroEvaluator, *expressions: t.Optional[exp.Expr]) -> exp.Condition:
1197    """Returns an AND statement filtering out any NULL expressions."""
1198    conditions = [e for e in expressions if not isinstance(e, exp.Null)]
1199
1200    if not conditions:
1201        return exp.true()
1202
1203    return exp.and_(*conditions, dialect=evaluator.dialect)
1204
1205
1206@macro("OR")
1207def or_(evaluator: MacroEvaluator, *expressions: t.Optional[exp.Expr]) -> exp.Condition:
1208    """Returns an OR statement filtering out any NULL expressions."""
1209    conditions = [e for e in expressions if not isinstance(e, exp.Null)]
1210
1211    if not conditions:
1212        return exp.true()
1213
1214    return exp.or_(*conditions, dialect=evaluator.dialect)
1215
1216
1217@macro("VAR")
1218def var(
1219    evaluator: MacroEvaluator, var_name: exp.Expr, default: t.Optional[exp.Expr] = None
1220) -> exp.Expr:
1221    """Returns the value of a variable or the default value if the variable is not set."""
1222    if not var_name.is_string:
1223        raise SQLMeshError(f"Invalid variable name '{var_name.sql()}'. Expected a string literal.")
1224
1225    return exp.convert(evaluator.var(var_name.this, default))
1226
1227
1228@macro("BLUEPRINT_VAR")
1229def blueprint_var(
1230    evaluator: MacroEvaluator, var_name: exp.Expr, default: t.Optional[exp.Expr] = None
1231) -> exp.Expr:
1232    """Returns the value of a blueprint variable or the default value if the variable is not set."""
1233    if not var_name.is_string:
1234        raise SQLMeshError(
1235            f"Invalid blueprint variable name '{var_name.sql()}'. Expected a string literal."
1236        )
1237
1238    return exp.convert(evaluator.blueprint_var(var_name.this, default))
1239
1240
1241@macro()
1242def deduplicate(
1243    evaluator: MacroEvaluator,
1244    relation: exp.Expr,
1245    partition_by: t.List[exp.Expr],
1246    order_by: t.List[str],
1247) -> exp.Query:
1248    """Returns a QUERY to deduplicate rows within a table
1249
1250    Args:
1251        relation: table or CTE name to deduplicate
1252        partition_by: column names, or expressions to use to identify a window of rows out of which to select one as the deduplicated row
1253        order_by: A list of strings representing the ORDER BY clause
1254
1255    Example:
1256        >>> from sqlglot import parse_one
1257        >>> from sqlglot.schema import MappingSchema
1258        >>> from sqlmesh.core.macros import MacroEvaluator
1259        >>> sql = "@deduplicate(demo.table, [user_id, cast(timestamp as date)], ['timestamp desc', 'status asc'])"
1260        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1261        'SELECT * FROM demo.table QUALIFY ROW_NUMBER() OVER (PARTITION BY user_id, CAST(timestamp AS DATE) ORDER BY timestamp DESC, status ASC) = 1'
1262    """
1263    if not isinstance(partition_by, list):
1264        raise SQLMeshError(
1265            "partition_by must be a list of columns: [<column>, cast(<column> as <type>)]"
1266        )
1267
1268    if not isinstance(order_by, list):
1269        raise SQLMeshError(
1270            "order_by must be a list of strings, optional - nulls ordering: ['<column> <asc|desc> nulls <first|last>']"
1271        )
1272
1273    partition_clause = exp.tuple_(*partition_by)
1274
1275    order_expressions = [
1276        evaluator.transform(parse_one(order_item, into=exp.Ordered, dialect=evaluator.dialect))
1277        for order_item in order_by
1278    ]
1279
1280    if not order_expressions:
1281        raise SQLMeshError(
1282            "order_by must be a list of strings, optional - nulls ordering: ['<column> <asc|desc> nulls <first|last>']"
1283        )
1284
1285    order_clause = exp.Order(expressions=order_expressions)
1286
1287    window_function = exp.Window(
1288        this=exp.RowNumber(), partition_by=partition_clause, order=order_clause
1289    )
1290
1291    first_unique_row = window_function.eq(1)
1292
1293    query = exp.select("*").from_(relation).qualify(first_unique_row)
1294
1295    return query
1296
1297
1298@macro()
1299def date_spine(
1300    evaluator: MacroEvaluator,
1301    datepart: exp.Expr,
1302    start_date: exp.Expr,
1303    end_date: exp.Expr,
1304) -> exp.Select:
1305    """Returns a query that produces a date spine with the given datepart, and range of start_date and end_date. Useful for joining as a date lookup table.
1306
1307    Args:
1308        datepart: The datepart to use for the date spine - day, week, month, quarter, year
1309        start_date: The start date for the date spine in format YYYY-MM-DD
1310        end_date: The end date for the date spine in format YYYY-MM-DD
1311
1312    Example:
1313        >>> from sqlglot import parse_one
1314        >>> from sqlglot.schema import MappingSchema
1315        >>> from sqlmesh.core.macros import MacroEvaluator
1316        >>> sql = "@date_spine('week', '2022-01-20', '2024-12-16')"
1317        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1318        "SELECT date_week FROM UNNEST(GENERATE_DATE_ARRAY(CAST(\'2022-01-20\' AS DATE), CAST(\'2024-12-16\' AS DATE), INTERVAL \'1\' WEEK)) AS _exploded(date_week)"
1319    """
1320    datepart_name = datepart.name.lower()
1321    if datepart_name not in ("day", "week", "month", "quarter", "year"):
1322        raise SQLMeshError(
1323            f"Invalid datepart '{datepart_name}'. Expected: 'day', 'week', 'month', 'quarter', or 'year'"
1324        )
1325
1326    start_date_name = start_date.name
1327    end_date_name = end_date.name
1328
1329    try:
1330        if start_date.is_string and end_date.is_string:
1331            start_date_obj = datetime.strptime(start_date_name, "%Y-%m-%d").date()
1332            end_date_obj = datetime.strptime(end_date_name, "%Y-%m-%d").date()
1333        else:
1334            start_date_obj = None
1335            end_date_obj = None
1336    except Exception as e:
1337        raise SQLMeshError(
1338            f"Invalid date format - start_date and end_date must be in format: YYYY-MM-DD. Error: {e}"
1339        )
1340
1341    if start_date_obj and end_date_obj:
1342        if start_date_obj > end_date_obj:
1343            raise SQLMeshError(
1344                f"Invalid date range - start_date '{start_date_name}' is after end_date '{end_date_name}'."
1345            )
1346
1347        start_date = exp.cast(start_date, "DATE")
1348        end_date = exp.cast(end_date, "DATE")
1349
1350    if datepart_name == "quarter" and evaluator.dialect in (
1351        "spark",
1352        "spark2",
1353        "databricks",
1354        "postgres",
1355    ):
1356        date_interval = exp.Interval(this=exp.Literal.number(3), unit=exp.var("month"))
1357    else:
1358        date_interval = exp.Interval(this=exp.Literal.number(1), unit=exp.var(datepart_name))
1359
1360    generate_date_array = exp.func(
1361        "GENERATE_DATE_ARRAY",
1362        start_date,
1363        end_date,
1364        date_interval,
1365    )
1366
1367    alias_name = f"date_{datepart_name}"
1368    exploded = exp.alias_(exp.func("unnest", generate_date_array), "_exploded", table=[alias_name])
1369
1370    return exp.select(alias_name).from_(exploded)
1371
1372
1373@macro()
1374def resolve_template(
1375    evaluator: MacroEvaluator,
1376    template: exp.Literal,
1377    mode: str = "literal",
1378) -> t.Union[exp.Literal, exp.Table]:
1379    """
1380    Generates either a String literal or an exp.Table representing a physical table location, based on rendering the provided template String literal.
1381
1382    Note: It relies on the @this_model variable being available in the evaluation context (@this_model resolves to an exp.Table object
1383    representing the current physical table).
1384    Therefore, the @resolve_template macro must be used at creation or evaluation time and not at load time.
1385
1386    Args:
1387        template: Template string literal. Can contain the following placeholders:
1388            @{catalog_name} -> replaced with the catalog of the exp.Table returned from @this_model
1389            @{schema_name} -> replaced with the schema of the exp.Table returned from @this_model
1390            @{table_name} -> replaced with the name of the exp.Table returned from @this_model
1391        mode: What to return.
1392            'literal' -> return an exp.Literal string
1393            'table' -> return an exp.Table
1394
1395    Example:
1396        >>> from sqlglot import parse_one, exp
1397        >>> from sqlmesh.core.macros import MacroEvaluator, RuntimeStage
1398        >>> sql = "@resolve_template('s3://data-bucket/prod/@{catalog_name}/@{schema_name}/@{table_name}')"
1399        >>> evaluator = MacroEvaluator(runtime_stage=RuntimeStage.CREATING)
1400        >>> evaluator.locals.update({"this_model": exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505")})
1401        >>> evaluator.transform(parse_one(sql)).sql()
1402        "'s3://data-bucket/prod/test_catalog/sqlmesh__test/test__test_model__2517971505'"
1403    """
1404    if "this_model" in evaluator.locals:
1405        this_model = exp.to_table(evaluator.locals["this_model"], dialect=evaluator.dialect)
1406        template_str: str = template.this
1407        result = (
1408            template_str.replace("@{catalog_name}", this_model.catalog)
1409            .replace("@{schema_name}", this_model.db)
1410            .replace("@{table_name}", this_model.name)
1411        )
1412
1413        if mode.lower() == "table":
1414            return exp.to_table(result, dialect=evaluator.dialect)
1415        return exp.Literal.string(result)
1416    if evaluator.runtime_stage != RuntimeStage.LOADING.value:
1417        # only error if we are CREATING, EVALUATING or TESTING and @this_model is not present; this could indicate a bug
1418        # otherwise, for LOADING, it's a no-op
1419        raise SQLMeshError(
1420            "@this_model must be present in the macro evaluation context in order to use @resolve_template"
1421        )
1422
1423    return template
1424
1425
1426def normalize_macro_name(name: str) -> str:
1427    """Prefix macro name with @ and upcase"""
1428    return f"@{name.upper()}"
1429
1430
1431for m in macro.get_registry().values():
1432    setattr(m, c.SQLMESH_BUILTIN, True)
1433
1434
1435def call_macro(
1436    func: t.Callable,
1437    dialect: DialectType,
1438    path: t.Optional[Path],
1439    provided_args: t.Tuple[t.Any, ...],
1440    provided_kwargs: t.Dict[str, t.Any],
1441    **optional_kwargs: t.Any,
1442) -> t.Any:
1443    # Bind the macro's actual parameters to its formal parameters
1444    sig = inspect.signature(func)
1445
1446    if optional_kwargs:
1447        provided_kwargs = provided_kwargs.copy()
1448
1449    for k, v in optional_kwargs.items():
1450        if k in sig.parameters:
1451            provided_kwargs[k] = v
1452
1453    bound = sig.bind(*provided_args, **provided_kwargs)
1454    bound.apply_defaults()
1455
1456    try:
1457        annotations = t.get_type_hints(func, localns=get_supported_types())
1458    except (NameError, TypeError):  # forward references aren't handled
1459        annotations = {}
1460
1461    # If the macro is annotated, we try coerce the actual parameters to the corresponding types
1462    if annotations:
1463        for arg, value in bound.arguments.items():
1464            typ = annotations.get(arg)
1465            if not typ:
1466                continue
1467
1468            # Changes to bound.arguments will reflect in bound.args and bound.kwargs
1469            # https://docs.python.org/3/library/inspect.html#inspect.BoundArguments.arguments
1470            param = sig.parameters[arg]
1471            if param.kind is inspect.Parameter.VAR_POSITIONAL:
1472                bound.arguments[arg] = tuple(_coerce(v, typ, dialect, path) for v in value)
1473            elif param.kind is inspect.Parameter.VAR_KEYWORD:
1474                bound.arguments[arg] = {k: _coerce(v, typ, dialect, path) for k, v in value.items()}
1475            else:
1476                bound.arguments[arg] = _coerce(value, typ, dialect, path)
1477
1478    return func(*bound.args, **bound.kwargs)
1479
1480
1481def _coerce(
1482    expr: t.Any,
1483    typ: t.Any,
1484    dialect: DialectType,
1485    path: t.Optional[Path] = None,
1486    strict: bool = False,
1487) -> t.Any:
1488    """Coerces the given expression to the specified type on a best-effort basis."""
1489    base_err_msg = f"Failed to coerce expression '{expr}' to type '{typ}'."
1490    try:
1491        if typ is None or typ is t.Any or not isinstance(expr, exp.Expr):
1492            return expr
1493        base = t.get_origin(typ) or typ
1494
1495        # We need to handle Union and TypeVars first since we cannot use isinstance with it
1496        if base in UNION_TYPES:
1497            for branch in t.get_args(typ):
1498                try:
1499                    return _coerce(expr, branch, dialect, path, strict=True)
1500                except Exception:
1501                    pass
1502            raise SQLMeshError(base_err_msg)
1503        if base is SQL and isinstance(expr, exp.Expr):
1504            return expr.sql(dialect)
1505
1506        if base is t.Literal:
1507            if not isinstance(expr, (exp.Literal, exp.Boolean)):
1508                raise SQLMeshError(
1509                    f"{base_err_msg} Coercion to {base} requires a literal expression."
1510                )
1511            literal_type_args = t.get_args(typ)
1512            try:
1513                for literal_type_arg in literal_type_args:
1514                    expr_is_bool = isinstance(expr.this, bool)
1515                    literal_is_bool = isinstance(literal_type_arg, bool)
1516                    if (expr_is_bool and literal_is_bool and literal_type_arg == expr.this) or (
1517                        not expr_is_bool
1518                        and not literal_is_bool
1519                        and str(literal_type_arg) == str(expr.this)
1520                    ):
1521                        return type(literal_type_arg)(expr.this)
1522            except Exception:
1523                raise SQLMeshError(base_err_msg)
1524            raise SQLMeshError(base_err_msg)
1525
1526        if isinstance(expr, base):
1527            return expr
1528        if issubclass(base, exp.Expr):
1529            d = Dialect.get_or_raise(dialect)
1530            into = base if base in d.parser_class.EXPRESSION_PARSERS else None
1531            if into is None:
1532                if isinstance(expr, exp.Literal):
1533                    coerced = parse_one(expr.this)
1534                else:
1535                    raise SQLMeshError(
1536                        f"{base_err_msg} Coercion to {base} requires a literal expression."
1537                    )
1538            else:
1539                coerced = parse_one(
1540                    expr.this if isinstance(expr, exp.Literal) else expr.sql(), into=into
1541                )
1542            if isinstance(coerced, base):
1543                return coerced
1544            raise SQLMeshError(base_err_msg)
1545
1546        if base in (int, float, str) and isinstance(expr, exp.Literal):
1547            return base(expr.this)
1548        if base is str and isinstance(expr, exp.Column) and not expr.table:
1549            return expr.name
1550        if base is bool and isinstance(expr, exp.Boolean):
1551            return expr.this
1552        if base is datetime and isinstance(expr, exp.Literal):
1553            return to_datetime(expr.this)
1554        if base is date and isinstance(expr, exp.Literal):
1555            return to_date(expr.this)
1556        if base is tuple and isinstance(expr, (exp.Tuple, exp.Array)):
1557            generic = t.get_args(typ)
1558            if not generic:
1559                return tuple(expr.expressions)
1560            if generic[-1] is ...:
1561                return tuple(_coerce(expr, generic[0], dialect, path) for expr in expr.expressions)
1562            if len(generic) == len(expr.expressions):
1563                return tuple(
1564                    _coerce(expr, generic[i], dialect, path)
1565                    for i, expr in enumerate(expr.expressions)
1566                )
1567            raise SQLMeshError(f"{base_err_msg} Expected {len(generic)} items.")
1568        if base is list and isinstance(expr, (exp.Array, exp.Tuple)):
1569            generic = t.get_args(typ)
1570            if not generic:
1571                return expr.expressions
1572            return [_coerce(expr, generic[0], dialect, path) for expr in expr.expressions]
1573        raise SQLMeshError(base_err_msg)
1574    except Exception:
1575        if strict:
1576            raise
1577
1578        from sqlmesh.core.console import get_console
1579
1580        get_console().log_error(
1581            f"Coercion of expression '{expr}' to type '{typ}' failed. Using non coerced expression at '{path}'",
1582        )
1583        return expr
1584
1585
1586def convert_sql(v: t.Any, dialect: DialectType) -> t.Any:
1587    try:
1588        return _cache_convert_sql(v, dialect, v.__class__)
1589    # dicts aren't hashable but are convertable
1590    except TypeError:
1591        return _convert_sql(v, dialect)
1592
1593
1594def _convert_sql(v: t.Any, dialect: DialectType) -> t.Any:
1595    if not isinstance(v, str):
1596        try:
1597            v = exp.convert(v)
1598        # we use bare Exception instead of ValueError because there's
1599        # a recursive error with MagicMock.
1600        except Exception:
1601            pass
1602
1603    if isinstance(v, exp.Expr):
1604        if (isinstance(v, exp.Column) and not v.table) or (
1605            isinstance(v, exp.Identifier) or v.is_string
1606        ):
1607            return v.name
1608        v = v.sql(dialect=dialect)
1609    return v
1610
1611
1612@lru_cache(maxsize=16384)
1613def _cache_convert_sql(v: t.Any, dialect: DialectType, t: type) -> t.Any:
1614    return _convert_sql(v, dialect)