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