sqlmesh.core.macros

   1from __future__ import annotations
   2
   3import inspect
   4import logging
   5import sys
   6import types
   7import typing as t
   8from enum import Enum
   9from functools import reduce
  10from itertools import chain
  11from pathlib import Path
  12from string import Template
  13
  14import sqlglot
  15from jinja2 import Environment
  16from sqlglot import Generator, exp, parse_one
  17from sqlglot.executor.env import ENV
  18from sqlglot.executor.python import Python
  19from sqlglot.helper import csv, ensure_collection
  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.errors import MacroEvalError, SQLMeshError
  41from sqlmesh.utils.jinja import JinjaMacroRegistry, has_jinja
  42from sqlmesh.utils.metaprogramming import Executable, prepare_env, print_exception
  43
  44if t.TYPE_CHECKING:
  45    from sqlmesh.core._typing import TableName
  46    from sqlmesh.core.engine_adapter import EngineAdapter
  47    from sqlmesh.core.snapshot import Snapshot
  48
  49
  50if sys.version_info >= (3, 10):
  51    UNION_TYPES = (t.Union, types.UnionType)
  52else:
  53    UNION_TYPES = (t.Union,)
  54
  55
  56logger = logging.getLogger(__name__)
  57
  58
  59class RuntimeStage(Enum):
  60    LOADING = "loading"
  61    CREATING = "creating"
  62    EVALUATING = "evaluating"
  63    TESTING = "testing"
  64
  65
  66class MacroStrTemplate(Template):
  67    delimiter = SQLMESH_MACRO_PREFIX
  68
  69
  70EXPRESSIONS_NAME_MAP = {}
  71SQL = t.NewType("SQL", str)
  72
  73for klass in sqlglot.Parser.EXPRESSION_PARSERS:
  74    name = klass if isinstance(klass, str) else klass.__name__  # type: ignore
  75    EXPRESSIONS_NAME_MAP[name.lower()] = name
  76
  77
  78def _macro_sql(sql: str, into: t.Optional[str] = None) -> str:
  79    args = [_macro_str_replace(sql)]
  80    if into in EXPRESSIONS_NAME_MAP:
  81        args.append(f"into=exp.{EXPRESSIONS_NAME_MAP[into]}")
  82    return f"self.parse_one({', '.join(args)})"
  83
  84
  85def _macro_func_sql(self: Generator, e: exp.Expression) -> str:
  86    func = e.this
  87
  88    if isinstance(func, exp.Anonymous):
  89        return f"""self.send({csv("'" + func.name + "'", self.expressions(func))})"""
  90    return self.sql(func)
  91
  92
  93def _macro_str_replace(text: str) -> str:
  94    """Stringifies python code for variable replacement
  95    Args:
  96        text: text string
  97    Returns:
  98        Stringified python code to execute variable replacement
  99    """
 100    return f"self.template({text}, locals())"
 101
 102
 103class MacroDialect(Python):
 104    class Generator(Python.Generator):
 105        TRANSFORMS = {
 106            **Python.Generator.TRANSFORMS,  # type: ignore
 107            exp.Column: lambda self, e: f"exp.to_column('{self.sql(e, 'this')}')",
 108            exp.Lambda: lambda self, e: f"lambda {self.expressions(e)}: {self.sql(e, 'this')}",
 109            MacroFunc: _macro_func_sql,
 110            MacroSQL: lambda self, e: _macro_sql(self.sql(e, "this"), e.args.get("into")),
 111            MacroStrReplace: lambda self, e: _macro_str_replace(self.sql(e, "this")),
 112        }
 113
 114
 115class MacroEvaluator:
 116    """The class responsible for evaluating SQLMesh Macros/SQL.
 117
 118    SQLMesh supports special preprocessed SQL prefixed with `@`. Although it provides similar power to
 119    traditional methods like string templating, there is semantic understanding of SQL which prevents
 120    common errors like leading/trailing commas, syntax errors, etc.
 121
 122    SQLMesh SQL allows for macro variables and macro functions. Macro variables take the form of @variable. These are used for variable substitution.
 123
 124    SELECT * FROM foo WHERE ds BETWEEN @start_date AND @end_date
 125
 126    Macro variables can be defined with a special macro function.
 127
 128    @DEF(start_date, '2021-01-01')
 129
 130    Args:
 131        dialect: Dialect of the SQL to evaluate.
 132        python_env: Serialized Python environment.
 133    """
 134
 135    def __init__(
 136        self,
 137        dialect: str = "",
 138        python_env: t.Optional[t.Dict[str, Executable]] = None,
 139        jinja_env: t.Optional[Environment] = None,
 140        schema: t.Optional[MappingSchema] = None,
 141        runtime_stage: RuntimeStage = RuntimeStage.LOADING,
 142        resolve_tables: t.Optional[t.Callable[[exp.Expression], exp.Expression]] = None,
 143        snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
 144        default_catalog: t.Optional[str] = None,
 145        path: Path = Path(),
 146    ):
 147        self.dialect = dialect
 148        self.generator = MacroDialect().generator()
 149        self.locals: t.Dict[str, t.Any] = {
 150            "runtime_stage": runtime_stage.value,
 151            "default_catalog": default_catalog,
 152        }
 153        self.env = {
 154            **ENV,
 155            "self": self,
 156            "SQL": SQL,
 157            "MacroEvaluator": MacroEvaluator,
 158        }
 159        self.python_env = python_env or {}
 160        self._jinja_env: t.Optional[Environment] = jinja_env
 161        self.macros = {normalize_macro_name(k): v.func for k, v in macro.get_registry().items()}
 162        self._schema = schema
 163        self._resolve_tables = resolve_tables
 164        self.columns_to_types_called = False
 165        self._snapshots = snapshots if snapshots is not None else {}
 166        self.default_catalog = default_catalog
 167        self._path = path
 168
 169        prepare_env(self.python_env, self.env)
 170        for k, v in self.python_env.items():
 171            if v.is_definition:
 172                self.macros[normalize_macro_name(k)] = self.env[v.name or k]
 173            elif v.is_import and getattr(self.env.get(k), c.SQLMESH_MACRO, None):
 174                self.macros[normalize_macro_name(k)] = self.env[k]
 175            elif v.is_value:
 176                self.locals[k] = self.env[k]
 177
 178    def send(
 179        self, name: str, *args: t.Any, **kwargs: t.Any
 180    ) -> t.Union[None, exp.Expression, t.List[exp.Expression]]:
 181        func = self.macros.get(normalize_macro_name(name))
 182
 183        if not callable(func):
 184            raise SQLMeshError(f"Macro '{name}' does not exist.")
 185
 186        try:
 187            annotations = t.get_type_hints(func)
 188        except NameError:  # forward references aren't handled
 189            annotations = {}
 190
 191        if annotations:
 192            spec = inspect.getfullargspec(func)
 193            callargs = inspect.getcallargs(func, self, *args, **kwargs)
 194            new_args: t.List[t.Any] = []
 195
 196            for arg, value in callargs.items():
 197                typ = annotations.get(arg)
 198
 199                if value is self:
 200                    continue
 201                if arg == spec.varargs:
 202                    new_args.extend(self._coerce(v, typ) for v in value)
 203                elif arg == spec.varkw:
 204                    for k, v in value.items():
 205                        kwargs[k] = self._coerce(v, typ)
 206                elif arg in kwargs:
 207                    kwargs[arg] = self._coerce(value, typ)
 208                else:
 209                    new_args.append(self._coerce(value, typ))
 210
 211            args = new_args  # type: ignore
 212
 213        try:
 214            return func(self, *args, **kwargs)
 215        except Exception as e:
 216            print_exception(e, self.python_env)
 217            raise MacroEvalError("Error trying to eval macro.") from e
 218
 219    def transform(
 220        self, expression: exp.Expression
 221    ) -> exp.Expression | t.List[exp.Expression] | None:
 222        changed = False
 223
 224        def evaluate_macros(
 225            node: exp.Expression,
 226        ) -> exp.Expression | t.List[exp.Expression] | None:
 227            nonlocal changed
 228
 229            if isinstance(node, MacroVar):
 230                changed = True
 231                variables = self.locals.get(c.SQLMESH_VARS, {})
 232                if node.name not in self.locals and node.name.lower() not in variables:
 233                    if not isinstance(node.parent, StagedFilePath):
 234                        raise SQLMeshError(f"Macro variable '{node.name}' is undefined.")
 235
 236                    return node
 237
 238                value = self.locals.get(node.name, variables.get(node.name.lower()))
 239                if isinstance(value, list):
 240                    return exp.convert(
 241                        tuple(
 242                            self.transform(v) if isinstance(v, exp.Expression) else v for v in value
 243                        )
 244                    )
 245                return exp.convert(
 246                    self.transform(value) if isinstance(value, exp.Expression) else value
 247                )
 248            if isinstance(node, exp.Identifier) and "@" in node.this:
 249                text = self.template(node.this, self.locals)
 250                if node.this != text:
 251                    changed = True
 252                    node.args["this"] = text
 253                    return node
 254            if node.is_string:
 255                text = node.this
 256                if has_jinja(text):
 257                    changed = True
 258                    node.set("this", self.jinja_env.from_string(node.this).render())
 259                return node
 260            if isinstance(node, MacroFunc):
 261                changed = True
 262                return self.evaluate(node)
 263            return node
 264
 265        transformed = exp.replace_tree(
 266            expression.copy(), evaluate_macros, prune=lambda n: isinstance(n, exp.Lambda)
 267        )
 268
 269        if changed:
 270            # the transformations could have corrupted the ast, turning this into sql and reparsing ensures
 271            # that the ast is correct
 272            if isinstance(transformed, list):
 273                return [
 274                    self.parse_one(node.sql(dialect=self.dialect, copy=False))
 275                    for node in transformed
 276                ]
 277            elif isinstance(transformed, exp.Expression):
 278                return self.parse_one(transformed.sql(dialect=self.dialect, copy=False))
 279
 280        return transformed
 281
 282    def template(self, text: t.Any, local_variables: t.Dict[str, t.Any]) -> str:
 283        """Substitute @vars with locals.
 284
 285        Args:
 286            text: The string to do substitition on.
 287            local_variables: Local variables in the context so that lambdas can be used.
 288
 289        Returns:
 290           The rendered string.
 291        """
 292        mapping = {}
 293
 294        variables = self.locals.get(c.SQLMESH_VARS, {})
 295
 296        for k, v in chain(variables.items(), self.locals.items(), local_variables.items()):
 297            # try to convert all variables into sqlglot expressions
 298            # because they're going to be converted into strings in sql
 299            # we use bare Exception instead of ValueError because there's
 300            # a recursive error with MagicMock.
 301            # we don't convert strings because that would result in adding quotes
 302            if not isinstance(v, str):
 303                try:
 304                    v = exp.convert(v)
 305                except Exception:
 306                    pass
 307
 308            if isinstance(v, exp.Expression):
 309                v = v.sql(dialect=self.dialect)
 310            mapping[k] = v
 311
 312        return MacroStrTemplate(str(text)).safe_substitute(mapping)
 313
 314    def evaluate(self, node: MacroFunc) -> exp.Expression | t.List[exp.Expression] | None:
 315        if isinstance(node, MacroDef):
 316            if isinstance(node.expression, exp.Lambda):
 317                _, fn = _norm_var_arg_lambda(self, node.expression)
 318                self.macros[normalize_macro_name(node.name)] = lambda _, *args: fn(
 319                    args[0] if len(args) == 1 else exp.Tuple(expressions=list(args))
 320                )
 321            else:
 322                self.locals[node.name] = self.transform(node.expression)
 323            return node
 324
 325        if isinstance(node, (MacroSQL, MacroStrReplace)):
 326            result: t.Optional[exp.Expression | t.List[exp.Expression]] = exp.convert(
 327                self.eval_expression(node)
 328            )
 329        else:
 330            func = t.cast(exp.Anonymous, node.this)
 331
 332            args = []
 333            kwargs = {}
 334            for e in func.expressions:
 335                if isinstance(e, exp.PropertyEQ):
 336                    kwargs[e.this.name] = e.expression
 337                else:
 338                    if kwargs:
 339                        raise MacroEvalError(
 340                            f"Positional argument cannot follow keyword argument.\n  {func.sql(dialect=self.dialect)} at '{self._path}'"
 341                        )
 342
 343                    args.append(e)
 344
 345            result = self.send(func.name, *args, **kwargs)
 346
 347        if result is None:
 348            return None
 349
 350        if isinstance(result, (tuple, list)):
 351            return [self.parse_one(item) for item in result if item is not None]
 352        return self.parse_one(result)
 353
 354    def eval_expression(self, node: t.Any) -> t.Any:
 355        """Converts a SQLGlot expression into executable Python code and evals it.
 356
 357        If the node is not an expression, it will simply be returned.
 358
 359        Args:
 360            node: expression
 361        Returns:
 362            The return value of the evaled Python Code.
 363        """
 364        if not isinstance(node, exp.Expression):
 365            return node
 366        code = node.sql()
 367        try:
 368            code = self.generator.generate(node)
 369            return eval(code, self.env, self.locals)
 370        except Exception as e:
 371            print_exception(e, self.python_env)
 372            raise MacroEvalError(
 373                f"Error trying to eval macro.\n\nGenerated code: {code}\n\nOriginal sql: {node}"
 374            ) from e
 375
 376    def parse_one(
 377        self, sql: str | exp.Expression, into: t.Optional[exp.IntoType] = None, **opts: t.Any
 378    ) -> exp.Expression:
 379        """Parses the given SQL string and returns a syntax tree for the first
 380        parsed SQL statement.
 381
 382        Args:
 383            sql: the SQL code or expression to parse.
 384            into: the Expression to parse into
 385            **opts: other options
 386
 387        Returns:
 388            Expression: the syntax tree for the first parsed statement
 389        """
 390        return sqlglot.maybe_parse(sql, dialect=self.dialect, into=into, **opts)
 391
 392    @property
 393    def jinja_env(self) -> Environment:
 394        if not self._jinja_env:
 395            jinja_env_methods = {**self.locals, **self.env}
 396            del jinja_env_methods["self"]
 397            self._jinja_env = JinjaMacroRegistry().build_environment(**jinja_env_methods)
 398        return self._jinja_env
 399
 400    def columns_to_types(self, model_name: TableName | exp.Column) -> t.Dict[str, exp.DataType]:
 401        """Returns the columns-to-types mapping corresponding to the specified model."""
 402        if self._schema is None or self._schema.empty:
 403            self.columns_to_types_called = True
 404            return {"__schema_unavailable_at_load__": exp.DataType.build("unknown")}
 405
 406        normalized_model_name = normalize_model_name(
 407            model_name,
 408            default_catalog=self.default_catalog,
 409            dialect=self.dialect,
 410        )
 411        columns_to_types = self._schema.find(exp.to_table(normalized_model_name))
 412        if columns_to_types is None:
 413            raise SQLMeshError(f"Schema for model '{model_name}' can't be statically determined.")
 414
 415        return columns_to_types
 416
 417    def get_snapshot(self, model_name: TableName | exp.Column) -> t.Optional[Snapshot]:
 418        """Returns the snapshot that corresponds to the given model name."""
 419        return self._snapshots.get(
 420            normalize_model_name(
 421                model_name,
 422                default_catalog=self.default_catalog,
 423                dialect=self.dialect,
 424            )
 425        )
 426
 427    def resolve_tables(self, query: exp.Expression) -> exp.Expression:
 428        """Resolves queries with references to SQLMesh model names to their physical tables."""
 429        if not self._resolve_tables:
 430            raise SQLMeshError(
 431                "Macro evaluator not properly initialized with resolve_tables lambda."
 432            )
 433        return self._resolve_tables(query)
 434
 435    @property
 436    def runtime_stage(self) -> RuntimeStage:
 437        """Returns the current runtime stage of the macro evaluation."""
 438        return self.locals["runtime_stage"]
 439
 440    @property
 441    def engine_adapter(self) -> EngineAdapter:
 442        engine_adapter = self.locals.get("engine_adapter")
 443        if not engine_adapter:
 444            raise SQLMeshError(
 445                "The engine adapter is not available while models are loading."
 446                " You can gate these calls by checking in Python: evaluator.runtime_stage != 'loading' or SQL: @runtime_stage <> 'loading'."
 447            )
 448        return self.locals["engine_adapter"]
 449
 450    @property
 451    def gateway(self) -> t.Optional[str]:
 452        """Returns the gateway name."""
 453        return self.var(c.GATEWAY)
 454
 455    def var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
 456        """Returns the value of the specified variable, or the default value if it doesn't exist."""
 457        return (self.locals.get(c.SQLMESH_VARS) or {}).get(var_name.lower(), default)
 458
 459    def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:
 460        """Coerces the given expression to the specified type on a best-effort basis."""
 461        base_err_msg = f"Failed to coerce expression '{expr}' to type '{typ}'."
 462        try:
 463            if typ is None or typ is t.Any:
 464                return expr
 465            base = t.get_origin(typ) or typ
 466
 467            # We need to handle Union and TypeVars first since we cannot use isinstance with it
 468            if base in UNION_TYPES:
 469                for branch in t.get_args(typ):
 470                    try:
 471                        return self._coerce(expr, branch, True)
 472                    except Exception:
 473                        pass
 474                raise SQLMeshError(base_err_msg)
 475            if base is SQL and isinstance(expr, exp.Expression):
 476                return expr.sql(self.dialect)
 477
 478            if isinstance(expr, base):
 479                return expr
 480            if issubclass(base, exp.Expression):
 481                d = Dialect.get_or_raise(self.dialect)
 482                into = base if base in d.parser().EXPRESSION_PARSERS else None
 483                if into is None:
 484                    if isinstance(expr, exp.Literal):
 485                        coerced = parse_one(expr.this)
 486                    else:
 487                        raise SQLMeshError(
 488                            f"{base_err_msg} Coercion to {base} requires a literal expression."
 489                        )
 490                else:
 491                    coerced = parse_one(
 492                        expr.this if isinstance(expr, exp.Literal) else expr.sql(), into=into
 493                    )
 494                if isinstance(coerced, base):
 495                    return coerced
 496                raise SQLMeshError(base_err_msg)
 497
 498            if base in (int, float, str) and isinstance(expr, exp.Literal):
 499                return base(expr.this)
 500            if base is str and isinstance(expr, exp.Column) and not expr.table:
 501                return expr.name
 502            if base is bool and isinstance(expr, exp.Boolean):
 503                return expr.this
 504            # if base is str and isinstance(expr, exp.Expression):
 505            #    return expr.sql(self.dialect)
 506            if base is tuple and isinstance(expr, (exp.Tuple, exp.Array)):
 507                generic = t.get_args(typ)
 508                if not generic:
 509                    return tuple(expr.expressions)
 510                if generic[-1] is ...:
 511                    return tuple(self._coerce(expr, generic[0]) for expr in expr.expressions)
 512                elif len(generic) == len(expr.expressions):
 513                    return tuple(
 514                        self._coerce(expr, generic[i]) for i, expr in enumerate(expr.expressions)
 515                    )
 516                raise SQLMeshError(f"{base_err_msg} Expected {len(generic)} items.")
 517            if base is list and isinstance(expr, (exp.Array, exp.Tuple)):
 518                generic = t.get_args(typ)
 519                if not generic:
 520                    return expr.expressions
 521                return [self._coerce(expr, generic[0]) for expr in expr.expressions]
 522            raise SQLMeshError(base_err_msg)
 523        except Exception:
 524            if strict:
 525                raise
 526            logger.error(
 527                "Coercion of expression '%s' to type '%s' failed. Using non coerced expression at '%s'",
 528                expr,
 529                typ,
 530                self._path,
 531            )
 532            return expr
 533
 534
 535class macro(registry_decorator):
 536    """Specifies a function is a macro and registers it the global MACROS registry.
 537
 538    Registered macros can be referenced in SQL statements to make queries more dynamic/cleaner.
 539
 540    Example:
 541        from sqlglot import exp
 542        from sqlmesh.core.macros import MacroEvaluator, macro
 543
 544        @macro()
 545        def add_one(evaluator: MacroEvaluator, column: exp.Literal) -> exp.Add:
 546            return evaluator.parse_one(f"{column} + 1")
 547
 548    Args:
 549        name: A custom name for the macro, the default is the name of the function.
 550    """
 551
 552    registry_name = "macros"
 553
 554    def __call__(
 555        self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
 556    ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
 557        wrapper = super().__call__(func)
 558
 559        # This is used to identify macros at runtime to unwrap during serialization.
 560        setattr(wrapper, c.SQLMESH_MACRO, True)
 561        return wrapper
 562
 563
 564ExecutableOrMacro = t.Union[Executable, macro]
 565MacroRegistry = UniqueKeyDict[str, ExecutableOrMacro]
 566
 567
 568def _norm_var_arg_lambda(
 569    evaluator: MacroEvaluator, func: exp.Lambda, *items: t.Any
 570) -> t.Tuple[t.Iterable, t.Callable]:
 571    """
 572    Converts sql literal array and lambda into actual python iterable + callable.
 573
 574    In order to support expressions like @EACH([a, b, c], x -> @SQL('@x')), the lambda var x
 575    needs be passed to the local state.
 576
 577    Args:
 578        evaluator: MacroEvaluator that invoked the macro
 579        func: Lambda SQLGlot expression.
 580        items: Array or items of SQLGlot expressions.
 581    """
 582
 583    def substitute(
 584        node: exp.Expression, args: t.Dict[str, exp.Expression]
 585    ) -> exp.Expression | t.List[exp.Expression] | None:
 586        if isinstance(node, (exp.Identifier, exp.Var)):
 587            if not isinstance(node.parent, exp.Column):
 588                name = node.name
 589                if name in args:
 590                    return args[name].copy()
 591                if name in evaluator.locals:
 592                    return exp.convert(evaluator.locals[name])
 593            if SQLMESH_MACRO_PREFIX in node.name:
 594                return node.__class__(
 595                    this=evaluator.template(node.name, {k: v.name for k, v in args.items()})
 596                )
 597        elif isinstance(node, MacroFunc):
 598            local_copy = evaluator.locals.copy()
 599            evaluator.locals.update(args)
 600            result = evaluator.transform(node)
 601            evaluator.locals = local_copy
 602            return result
 603        return node
 604
 605    if len(items) == 1:
 606        item = items[0]
 607        expressions = item.expressions if isinstance(item, (exp.Array, exp.Tuple)) else item
 608    else:
 609        expressions = items
 610
 611    if not callable(func):
 612        return expressions, lambda args: func.this.transform(
 613            substitute,
 614            {
 615                expression.name: arg
 616                for expression, arg in zip(
 617                    func.expressions, args.expressions if isinstance(args, exp.Tuple) else [args]
 618                )
 619            },
 620        )
 621
 622    return expressions, func
 623
 624
 625@macro()
 626def each(
 627    evaluator: MacroEvaluator,
 628    *args: t.Any,
 629) -> t.List[t.Any]:
 630    """Iterates through items calling func on each.
 631
 632    If a func call on item returns None, it will be excluded from the list.
 633
 634    Args:
 635        evaluator: MacroEvaluator that invoked the macro
 636        args: The last argument should be a lambda of the form x -> x +1. The first argument can be
 637            an Array or var args can be used.
 638
 639    Returns:
 640        A list of items that is the result of func
 641    """
 642    *items, func = args
 643    items, func = _norm_var_arg_lambda(evaluator, func, *items)  # type: ignore
 644    return [item for item in map(func, ensure_collection(items)) if item is not None]
 645
 646
 647@macro("IF")
 648def if_(
 649    evaluator: MacroEvaluator,
 650    condition: t.Any,
 651    true: t.Any,
 652    false: t.Any = None,
 653) -> t.Any:
 654    """Evaluates a given condition and returns the second argument if true or else the third argument.
 655
 656    If false is not passed in, the default return value will be None.
 657
 658    Example:
 659        >>> from sqlglot import parse_one
 660        >>> from sqlmesh.core.macros import MacroEvaluator
 661        >>> MacroEvaluator().transform(parse_one("@IF('a' = 1, a, b)")).sql()
 662        'b'
 663
 664        >>> MacroEvaluator().transform(parse_one("@IF('a' = 1, a)"))
 665    """
 666
 667    if evaluator.eval_expression(condition):
 668        return true
 669    return false
 670
 671
 672@macro("REDUCE")
 673def reduce_(evaluator: MacroEvaluator, *args: t.Any) -> t.Any:
 674    """Iterates through items applying provided function that takes two arguments
 675    cumulatively to the items of iterable items, from left to right, so as to reduce
 676    the iterable to a single item.
 677
 678    Example:
 679        >>> from sqlglot import parse_one
 680        >>> from sqlmesh.core.macros import MacroEvaluator
 681        >>> sql = "@SQL(@REDUCE([100, 200, 300, 400], (x, y) -> x + y))"
 682        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 683        '1000'
 684
 685    Args:
 686        evaluator: MacroEvaluator that invoked the macro
 687        args: The last argument should be a lambda of the form (x, y) -> x + y. The first argument can be
 688            an Array or var args can be used.
 689    Returns:
 690        A single item that is the result of applying func cumulatively to items
 691    """
 692    *items, func = args
 693    items, func = _norm_var_arg_lambda(evaluator, func, *items)  # type: ignore
 694    return reduce(lambda a, b: func(exp.Tuple(expressions=[a, b])), ensure_collection(items))
 695
 696
 697@macro("FILTER")
 698def filter_(evaluator: MacroEvaluator, *args: t.Any) -> t.List[t.Any]:
 699    """Iterates through items, applying provided function to each item and removing
 700    all items where the function returns False
 701
 702    Example:
 703        >>> from sqlglot import parse_one
 704        >>> from sqlmesh.core.macros import MacroEvaluator
 705        >>> sql = "@REDUCE(@FILTER([1, 2, 3], x -> x > 1), (x, y) -> x + y)"
 706        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 707        '2 + 3'
 708
 709        >>> sql = "@EVAL(@REDUCE(@FILTER([1, 2, 3], x -> x > 1), (x, y) -> x + y))"
 710        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 711        '5'
 712
 713    Args:
 714        evaluator: MacroEvaluator that invoked the macro
 715        args: The last argument should be a lambda of the form x -> x > 1. The first argument can be
 716            an Array or var args can be used.
 717    Returns:
 718        The items for which the func returned True
 719    """
 720    *items, func = args
 721    items, func = _norm_var_arg_lambda(evaluator, func, *items)  # type: ignore
 722    return list(filter(lambda arg: evaluator.eval_expression(func(arg)), items))
 723
 724
 725def _optional_expression(
 726    evaluator: MacroEvaluator,
 727    condition: exp.Condition,
 728    expression: exp.Expression,
 729) -> t.Optional[exp.Expression]:
 730    """Inserts expression when the condition is True
 731
 732    The following examples express the usage of this function in the context of the macros which wrap it.
 733
 734    Examples:
 735        >>> from sqlglot import parse_one
 736        >>> from sqlmesh.core.macros import MacroEvaluator
 737        >>> sql = "@WITH(True) all_cities as (select * from city) select all_cities"
 738        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 739        'WITH all_cities AS (SELECT * FROM city) SELECT all_cities'
 740        >>> sql = "select * from city left outer @JOIN(True) country on city.country = country.name"
 741        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 742        'SELECT * FROM city LEFT OUTER JOIN country ON city.country = country.name'
 743        >>> sql = "select * from city @GROUP_BY(True) country, population"
 744        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 745        'SELECT * FROM city GROUP BY country, population'
 746        >>> sql = "select * from city group by country @HAVING(True) population > 100 and country = 'Mexico'"
 747        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 748        "SELECT * FROM city GROUP BY country HAVING population > 100 AND country = 'Mexico'"
 749
 750    Args:
 751        evaluator: MacroEvaluator that invoked the macro
 752        condition: Condition expression
 753        expression: SQL expression
 754    Returns:
 755        Expression if the conditional is True; otherwise None
 756    """
 757    return expression if evaluator.eval_expression(condition) else None
 758
 759
 760with_ = macro("WITH")(_optional_expression)
 761join = macro("JOIN")(_optional_expression)
 762where = macro("WHERE")(_optional_expression)
 763group_by = macro("GROUP_BY")(_optional_expression)
 764having = macro("HAVING")(_optional_expression)
 765order_by = macro("ORDER_BY")(_optional_expression)
 766limit = macro("LIMIT")(_optional_expression)
 767
 768
 769@macro("eval")
 770def eval_(evaluator: MacroEvaluator, condition: exp.Condition) -> t.Any:
 771    """Evaluate the given condition in a Python/SQL interpretor.
 772
 773    Example:
 774        >>> from sqlglot import parse_one
 775        >>> from sqlmesh.core.macros import MacroEvaluator
 776        >>> sql = "@EVAL(1 + 1)"
 777        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 778        '2'
 779    """
 780    return evaluator.eval_expression(condition)
 781
 782
 783# macros with union types need to use t.Union since | isn't available until 3.9
 784@macro()
 785def star(
 786    evaluator: MacroEvaluator,
 787    relation: exp.Table,
 788    alias: exp.Column = t.cast(exp.Column, exp.column("")),
 789    except_: t.Union[exp.Array, exp.Tuple] = exp.Tuple(this=[]),
 790    prefix: exp.Literal = exp.Literal.string(""),
 791    suffix: exp.Literal = exp.Literal.string(""),
 792    quote_identifiers: exp.Boolean = exp.true(),
 793) -> t.List[exp.Alias]:
 794    """Returns a list of projections for the given relation.
 795
 796    Args:
 797        evaluator: MacroEvaluator that invoked the macro
 798        relation: The relation to select star from
 799        alias: The alias of the relation
 800        except_: Columns to exclude
 801        prefix: A prefix to use for all selections
 802        suffix: A suffix to use for all selections
 803        quote_identifiers: Whether or not quote the resulting aliases, defaults to true
 804
 805    Returns:
 806        An array of columns.
 807
 808    Example:
 809        >>> from sqlglot import parse_one, exp
 810        >>> from sqlglot.schema import MappingSchema
 811        >>> from sqlmesh.core.macros import MacroEvaluator
 812        >>> sql = "SELECT @STAR(foo, bar, [c], 'baz_') FROM foo AS bar"
 813        >>> 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()
 814        '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'
 815    """
 816    if alias and not isinstance(alias, (exp.Identifier, exp.Column)):
 817        raise SQLMeshError(f"Invalid alias '{alias}'. Expected an identifier.")
 818    if except_ and not isinstance(except_, (exp.Array, exp.Tuple)):
 819        raise SQLMeshError(f"Invalid except '{except_}'. Expected an array.")
 820    if prefix and not isinstance(prefix, exp.Literal):
 821        raise SQLMeshError(f"Invalid prefix '{prefix}'. Expected a literal.")
 822    if suffix and not isinstance(suffix, exp.Literal):
 823        raise SQLMeshError(f"Invalid suffix '{suffix}'. Expected a literal.")
 824    if not isinstance(quote_identifiers, exp.Boolean):
 825        raise SQLMeshError(f"Invalid quote_identifiers '{quote_identifiers}'. Expected a boolean.")
 826
 827    exclude = {e.name for e in except_.expressions}
 828    quoted = quote_identifiers.this
 829    table_identifier = alias.name or relation.name
 830
 831    columns_to_types = {
 832        k: v for k, v in evaluator.columns_to_types(relation).items() if k not in exclude
 833    }
 834    if columns_to_types_all_known(columns_to_types):
 835        return [
 836            exp.cast(
 837                exp.column(column, table=table_identifier, quoted=quoted),
 838                dtype,
 839                dialect=evaluator.dialect,
 840            ).as_(f"{prefix.this}{column}{suffix.this}", quoted=quoted)
 841            for column, dtype in columns_to_types.items()
 842        ]
 843    return [
 844        exp.column(column, table=table_identifier, quoted=quoted).as_(
 845            f"{prefix.this}{column}{suffix.this}", quoted=quoted
 846        )
 847        for column, type_ in evaluator.columns_to_types(relation).items()
 848    ]
 849
 850
 851@macro()
 852def generate_surrogate_key(_: MacroEvaluator, *fields: exp.Expression) -> exp.Func:
 853    """Generates a surrogate key for the given fields.
 854
 855    Example:
 856        >>> from sqlglot import parse_one
 857        >>> from sqlmesh.core.macros import MacroEvaluator
 858        >>> sql = "SELECT @GENERATE_SURROGATE_KEY(a, b, c) FROM foo"
 859        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 860        "SELECT MD5(CONCAT(COALESCE(CAST(a AS TEXT), '_sqlmesh_surrogate_key_null_'), '|', COALESCE(CAST(b AS TEXT), '_sqlmesh_surrogate_key_null_'), '|', COALESCE(CAST(c AS TEXT), '_sqlmesh_surrogate_key_null_'))) FROM foo"
 861    """
 862    string_fields: t.List[exp.Expression] = []
 863    for i, field in enumerate(fields):
 864        if i > 0:
 865            string_fields.append(exp.Literal.string("|"))
 866        string_fields.append(
 867            exp.func(
 868                "COALESCE",
 869                exp.cast(field, exp.DataType.build("text")),
 870                exp.Literal.string("_sqlmesh_surrogate_key_null_"),
 871            )
 872        )
 873    return exp.func("MD5", exp.func("CONCAT", *string_fields))
 874
 875
 876@macro()
 877def safe_add(_: MacroEvaluator, *fields: exp.Expression) -> exp.Case:
 878    """Adds numbers together, substitutes nulls for 0s and only returns null if all fields are null.
 879
 880    Example:
 881        >>> from sqlglot import parse_one
 882        >>> from sqlmesh.core.macros import MacroEvaluator
 883        >>> sql = "SELECT @SAFE_ADD(a, b) FROM foo"
 884        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 885        'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) + COALESCE(b, 0) END FROM foo'
 886    """
 887    return (
 888        exp.Case()
 889        .when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null())
 890        .else_(reduce(lambda a, b: a + b, [exp.func("COALESCE", field, 0) for field in fields]))  # type: ignore
 891    )
 892
 893
 894@macro()
 895def safe_sub(_: MacroEvaluator, *fields: exp.Expression) -> exp.Case:
 896    """Subtract numbers, substitutes nulls for 0s and only returns null if all fields are null.
 897
 898    Example:
 899        >>> from sqlglot import parse_one
 900        >>> from sqlmesh.core.macros import MacroEvaluator
 901        >>> sql = "SELECT @SAFE_SUB(a, b) FROM foo"
 902        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 903        'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) - COALESCE(b, 0) END FROM foo'
 904    """
 905    return (
 906        exp.Case()
 907        .when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null())
 908        .else_(reduce(lambda a, b: a - b, [exp.func("COALESCE", field, 0) for field in fields]))  # type: ignore
 909    )
 910
 911
 912@macro()
 913def safe_div(_: MacroEvaluator, numerator: exp.Expression, denominator: exp.Expression) -> exp.Div:
 914    """Divides numbers, returns null if the denominator is 0.
 915
 916    Example:
 917        >>> from sqlglot import parse_one
 918        >>> from sqlmesh.core.macros import MacroEvaluator
 919        >>> sql = "SELECT @SAFE_DIV(a, b) FROM foo"
 920        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 921        'SELECT a / NULLIF(b, 0) FROM foo'
 922    """
 923    return numerator / exp.func("NULLIF", denominator, 0)
 924
 925
 926@macro()
 927def union(
 928    evaluator: MacroEvaluator,
 929    type_: exp.Literal = exp.Literal.string("ALL"),
 930    *tables: exp.Table,
 931) -> exp.Query:
 932    """Returns a UNION of the given tables. Only choosing columns that have the same name and type.
 933
 934    Example:
 935        >>> from sqlglot import parse_one
 936        >>> from sqlglot.schema import MappingSchema
 937        >>> from sqlmesh.core.macros import MacroEvaluator
 938        >>> sql = "@UNION('distinct', foo, bar)"
 939        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": "int", "b": "string", "c": "string"}, "bar": {"c": "string", "a": "int", "b": "int"}})).transform(parse_one(sql)).sql()
 940        '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'
 941    """
 942    kind = type_.name.upper()
 943    if kind not in ("ALL", "DISTINCT"):
 944        raise SQLMeshError(f"Invalid type '{type_}'. Expected 'ALL' or 'DISTINCT'.")
 945
 946    columns = {
 947        column
 948        for column, _ in reduce(
 949            lambda a, b: a & b,  # type: ignore
 950            (evaluator.columns_to_types(table).items() for table in tables),
 951        )
 952    }
 953
 954    projections = [
 955        exp.cast(column, type_, dialect=evaluator.dialect).as_(column)
 956        for column, type_ in evaluator.columns_to_types(tables[0]).items()
 957        if column in columns
 958    ]
 959
 960    return reduce(
 961        lambda a, b: a.union(b, distinct=kind == "DISTINCT"),  # type: ignore
 962        [exp.select(*projections).from_(t) for t in tables],
 963    )
 964
 965
 966@macro()
 967def haversine_distance(
 968    _: MacroEvaluator,
 969    lat1: exp.Expression,
 970    lon1: exp.Expression,
 971    lat2: exp.Expression,
 972    lon2: exp.Expression,
 973    unit: exp.Literal = exp.Literal.string("mi"),
 974) -> exp.Mul:
 975    """Returns the haversine distance between two points.
 976
 977    Example:
 978        >>> from sqlglot import parse_one
 979        >>> from sqlmesh.core.macros import MacroEvaluator
 980        >>> sql = "SELECT @HAVERSINE_DISTANCE(driver_y, driver_x, passenger_y, passenger_x, 'mi') FROM rides"
 981        >>> MacroEvaluator().transform(parse_one(sql)).sql()
 982        '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'
 983    """
 984    if unit.this == "mi":
 985        conversion_rate = 1.0
 986    elif unit.this == "km":
 987        conversion_rate = 1.60934
 988    else:
 989        raise SQLMeshError(f"Invalid unit '{unit}'. Expected 'mi' or 'km'.")
 990
 991    return (
 992        2
 993        * 3961
 994        * exp.func(
 995            "ASIN",
 996            exp.func(
 997                "SQRT",
 998                exp.func("POWER", exp.func("SIN", exp.func("RADIANS", (lat2 - lat1) / 2)), 2)
 999                + exp.func("COS", exp.func("RADIANS", lat1))
1000                * exp.func("COS", exp.func("RADIANS", lat2))
1001                * exp.func("POWER", exp.func("SIN", exp.func("RADIANS", (lon2 - lon1) / 2)), 2),
1002            ),
1003        )
1004        * conversion_rate
1005    )
1006
1007
1008@macro()
1009def pivot(
1010    evaluator: MacroEvaluator,
1011    column: exp.Column,
1012    values: t.Union[exp.Array, exp.Tuple],
1013    alias: exp.Boolean = exp.true(),
1014    agg: exp.Literal = exp.Literal.string("SUM"),
1015    cmp: exp.Literal = exp.Literal.string("="),
1016    prefix: exp.Literal = exp.Literal.string(""),
1017    suffix: exp.Literal = exp.Literal.string(""),
1018    then_value: exp.Literal = exp.Literal.number(1),
1019    else_value: exp.Literal = exp.Literal.number(0),
1020    quote: exp.Boolean = exp.true(),
1021    distinct: exp.Boolean = exp.false(),
1022) -> t.List[exp.Expression]:
1023    """Returns a list of projections as a result of pivoting the given column on the given values.
1024
1025    Example:
1026        >>> from sqlglot import parse_one
1027        >>> from sqlmesh.core.macros import MacroEvaluator
1028        >>> sql = "SELECT date_day, @PIVOT(status, ['cancelled', 'completed']) FROM rides GROUP BY 1"
1029        >>> MacroEvaluator().transform(parse_one(sql)).sql()
1030        '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'
1031    """
1032    aggregates: t.List[exp.Expression] = []
1033    for value in values.expressions:
1034        proj = f"{agg.this}("
1035        if distinct.this:
1036            proj += "DISTINCT "
1037        proj += f"CASE WHEN {column} {cmp.this} {value} THEN {then_value} ELSE {else_value} END) "
1038        node = evaluator.parse_one(proj)
1039        if alias.this:
1040            node = node.as_(f"{prefix.this}{value}{suffix.this}", quoted=quote.this, copy=False)
1041        aggregates.append(node)
1042    return aggregates
1043
1044
1045@macro("AND")
1046def and_(evaluator: MacroEvaluator, *expressions: t.Optional[exp.Expression]) -> exp.Condition:
1047    """Returns an AND statement filtering out any NULL expressions."""
1048    conditions = [e for e in expressions if not isinstance(e, exp.Null)]
1049
1050    if not conditions:
1051        return exp.true()
1052
1053    return exp.and_(*conditions, dialect=evaluator.dialect)
1054
1055
1056@macro("OR")
1057def or_(evaluator: MacroEvaluator, *expressions: t.Optional[exp.Expression]) -> exp.Condition:
1058    """Returns an OR statement filtering out any NULL expressions."""
1059    conditions = [e for e in expressions if not isinstance(e, exp.Null)]
1060
1061    if not conditions:
1062        return exp.true()
1063
1064    return exp.or_(*conditions, dialect=evaluator.dialect)
1065
1066
1067@macro("VAR")
1068def var(
1069    evaluator: MacroEvaluator, var_name: exp.Expression, default: t.Optional[exp.Expression] = None
1070) -> exp.Expression:
1071    """Returns the value of a variable or the default value if the variable is not set."""
1072    if not var_name.is_string:
1073        raise SQLMeshError(f"Invalid variable name '{var_name.sql()}'. Expected a string literal.")
1074
1075    return exp.convert(evaluator.var(var_name.this, default))
1076
1077
1078def normalize_macro_name(name: str) -> str:
1079    """Prefix macro name with @ and upcase"""
1080    return f"@{name.upper()}"
1081
1082
1083for m in macro.get_registry().values():
1084    setattr(m, c.SQLMESH_BUILTIN, True)