Repository: TobikoData/sqlmesh
Pull request #2511 (2024-04-24 21:17:40)

Toby/typed macros


Patch:
diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md
index 2bc18701..f4ec89a7 100644
--- a/docs/concepts/macros/sqlmesh_macros.md
+++ b/docs/concepts/macros/sqlmesh_macros.md
@@ -1451,6 +1451,116 @@ The methods are available because the `column` argument is parsed as a SQLGlot [
 
 Column expressions are sub-classes of the [Condition class](https://sqlglot.com/sqlglot/expressions.html#Condition), so they have builder methods like [`between`](https://sqlglot.com/sqlglot/expressions.html#Condition.between) and [`like`](https://sqlglot.com/sqlglot/expressions.html#Condition.like).
 
+## Typed Macros
+
+Typed macros in SQLMesh bring the power of type hints from Python, enhancing readability, maintainability, and usability of your SQL macros. These macros enable developers to specify expected types for arguments, making the macros more intuitive and less error-prone.
+
+### Benefits of Typed Macros
+
+1. **Improved Readability**: By specifying types, the intent of the macro is clearer to other developers or future you.
+2. **Reduced Boilerplate**: No need for manual type conversion within the macro function, allowing you to focus on the core logic.
+3. **Enhanced Autocompletion**: IDEs can provide better autocompletion and documentation based on the specified types.
+
+### Defining a Typed Macro
+
+Typed macros in SQLMesh use Python's type hints. Here's a simple example of a typed macro that repeats a string a given number of times:
+
+```python linenums="1"
+from sqlmesh import macro
+
+@macro()
+def repeat_string(evaluator, text: str, count: int) -> str:
+    return text * count
+```
+
+Usage in SQLMesh:
+
+```sql linenums="1"
+SELECT
+  @repeat_string('SQLMesh ', 3) as repeated_string
+FROM some_table;
+```
+
+This macro takes two arguments: `text` of type `str` and `count` of type `int`, and it returns a string. Without type hints, the inputs to the macro would have been two `exp.Literal` objects you would have had to convert to strings and integers manually.
+
+### Supported Types
+
+SQLMesh supports common Python types for typed macros including:
+
+- `str` -- This handles string literals and basic identifiers, but won't coerce anything more complicated.
+- `int`
+- `float`
+- `bool`
+- `SQL` -- When you want the SQL string representation of the argument that's passed in
+- `List[T]` - where `T` is any supported type including sqlglot expressions
+- `Tuple[T]` - where `T` is any supported type including sqlglot expressions
+- `Union[T1, T2, ...]` - where `T1`, `T2`, etc. are any supported types including sqlglot expressions
+
+We also support SQLGlot expressions as type hints, allowing you to ensure inputs are coerced to the desired SQL AST node your intending on working with. Some useful examples include:
+
+- `exp.Table`
+- `exp.Column`
+- `exp.Literal`
+- `exp.Identifier`
+
+While these might be obvious examples, you can effectively coerce an input into _any_ SQLGlot expression type, which can be useful for more complex macros. When coercing to more complex types, you will almost certainly need to pass a string literal since expression to expression coercion is limited. When a string literal is passed to a macro that hints at a SQLGlot expression, the string will be parsed using SQLGlot and coerced to the correct type. Failure to coerce to the correct type will result in the original expression being passed to the macro and a warning being logged for the user to address as-needed.
+
+```python linenums="1"
+@macro()
+def stamped(evaluator, query: exp.Select) -> exp.Subquery:
+    return query.select(exp.Literal.string(str(datetime.now())).as_("stamp")).subquery()
+
+# Coercing to a complex node like `exp.Select` works as expected given a string literal input
+# SELECT * FROM @stamped('SELECT a, b, c')
+```
+
+When coercion fails, there will always be a warning logged but we will not crash. We believe the macro should be flexible by default, meaning the default behavior is preserved if we cannot coerce. Give that, the user can express whatever level of additional checks they want. For example, if you would like to raise an error when the coercion fails, you can use an `assert` statement. For example:
+
+```python linenums="1"
+@macro()
+def my_macro(evaluator, table: exp.Table) -> exp.Column:
+    assert isinstance(table, exp.Table)
+    table.set("catalog", "dev")
+    return table
+
+# Works
+# SELECT * FROM @my_macro('some.table')
+# SELECT * FROM @my_macro(some.table)
+
+# Raises an error thanks to the users inclusion of the assert, otherwise would pass through the string literal and log a warning
+# SELECT * FROM @my_macro('SELECT 1 + 1')
+```
+
+In using assert this way, you still get the benefits of reducing/removing the boilerplate needed to coerce types; but you **also** get guarantees about the type of the input. This is a useful pattern and is user-defined, so you can use it as you see fit. It ultimately allows you to keep the macro definition clean and focused on the core business logic.
+
+### Advanced Typed Macros
+
+You can create more complex macros using advanced Python features like generics. For example, a macro that accepts a list of integers and returns their sum:
+
+```python linenums="1"
+from typing import List
+from sqlmesh import macro
+
+@macro()
+def sum_integers(evaluator, numbers: List[int]) -> int:
+    return sum(numbers)
+```
+
+Usage in SQLMesh:
+
+```sql linenums="1"
+SELECT
+  @sum_integers([1, 2, 3, 4, 5]) as total
+FROM some_table;
+```
+
+Generics can be nested and are resolved recursively allowing for fairly robust type hinting.
+
+See examples of the coercion function in action in the test suite [here](../../../tests/core/test_macros.py).
+
+### Conclusion
+
+Typed macros in SQLMesh not only enhance the development experience by making macros more readable and easier to use but also contribute to more robust and maintainable code. By leveraging Python's type hinting system, developers can create powerful and intuitive macros for their SQL queries, further bridging the gap between SQL and Python.
 
 ## Mixing macro systems
 
diff --git a/examples/sushi/macros/macros.py b/examples/sushi/macros/macros.py
index abf0ebb9..f4c2b6a8 100644
--- a/examples/sushi/macros/macros.py
+++ b/examples/sushi/macros/macros.py
@@ -1,8 +1,9 @@
 from macros.utils import between  # type: ignore
+from sqlglot import exp
 
 from sqlmesh import macro
 
 
 @macro()
-def incremental_by_ds(evaluator, column):
+def incremental_by_ds(evaluator, column: exp.Column):
     return between(evaluator, column, evaluator.locals["start_date"], evaluator.locals["end_date"])
diff --git a/examples/sushi/macros/utils.py b/examples/sushi/macros/utils.py
index 2ee74116..5442e196 100644
--- a/examples/sushi/macros/utils.py
+++ b/examples/sushi/macros/utils.py
@@ -1,9 +1,45 @@
-from sqlmesh import macro
+from sqlglot import exp
+
+from sqlmesh import SQL, macro
+
+
+@macro()
+def add_one(evaluator, column: int):
+    # typed column will be cast to an int and return an integer back
+    assert isinstance(column, int)
+    return column + 1
 
 
 @macro()
-def add_one(evaluator, column):
-    return f"{column} + 1"
+def is_positive(evaluator, column):
+    # untyped column will be a sqlglot column and return a sqlglot exp "column > 0"
+    assert isinstance(column, exp.Column)
+    return column > 0
+
+
+@macro()
+def sql_literal(
+    evaluator,
+    column: SQL,
+    str_lit: SQL,
+    string: str,
+    column_str: str,
+    column_quoted: str,
+    column_with_db: str,
+):
+    assert isinstance(column, str)
+    assert isinstance(str_lit, str)
+    assert str_lit == "'x'"
+    assert isinstance(string, str)
+    assert string == "y"
+    assert isinstance(column_str, str)
+    assert column_str == "a"
+    assert isinstance(column_quoted, str)
+    assert column_quoted == "b"
+    assert isinstance(column_with_db, exp.Column)
+    assert column_with_db.sql() == "c.d"
+
+    return column
 
 
 @macro()
diff --git a/examples/sushi/models/customers.sql b/examples/sushi/models/customers.sql
index 7563ff4c..9dd73373 100644
--- a/examples/sushi/models/customers.sql
+++ b/examples/sushi/models/customers.sql
@@ -28,5 +28,5 @@ SELECT DISTINCT
   FROM sushi.orders AS o
 LEFT JOIN current_marketing AS m
   ON o.customer_id = m.customer_id
-  LEFT JOIN raw.demographics AS d
+LEFT JOIN raw.demographics AS d
   ON o.customer_id = d.customer_id
diff --git a/examples/sushi/models/top_waiters.sql b/examples/sushi/models/top_waiters.sql
index ad7cb333..7712686c 100644
--- a/examples/sushi/models/top_waiters.sql
+++ b/examples/sushi/models/top_waiters.sql
@@ -8,6 +8,13 @@ MODEL (
   grain waiter_id
 );
 
+WITH test_macros AS (
+    SELECT
+      @ADD_ONE(1) AS lit_two,
+      @IS_POSITIVE(revenue) AS sql_exp,
+      @SQL_LITERAL(revenue::text, 'x', 'y', a, "b", c.d) AS sql_lit,
+    FROM sushi.waiter_revenue_by_day
+)
 SELECT
   waiter_id::INT AS waiter_id,
   revenue::DOUBLE AS revenue
diff --git a/setup.py b/setup.py
index 3fd5a560..50eaa877 100644
--- a/setup.py
+++ b/setup.py
@@ -46,7 +46,7 @@ setup(
         "requests",
         "rich[jupyter]",
         "ruamel.yaml",
-        "sqlglot[rs]~=23.11.2",
+        "sqlglot[rs]~=23.12.1",
     ],
     extras_require={
         "bigquery": [
diff --git a/sqlmesh/__init__.py b/sqlmesh/__init__.py
index 9324f40d..b65c39d2 100644
--- a/sqlmesh/__init__.py
+++ b/sqlmesh/__init__.py
@@ -21,7 +21,7 @@ from sqlmesh.core import constants as c
 from sqlmesh.core.config import Config
 from sqlmesh.core.context import Context, ExecutionContext
 from sqlmesh.core.engine_adapter import EngineAdapter
-from sqlmesh.core.macros import macro
+from sqlmesh.core.macros import SQL, macro
 from sqlmesh.core.model import Model, model
 from sqlmesh.core.snapshot import Snapshot
 from sqlmesh.utils import debug_mode_enabled, enable_debug_mode
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 2a46bc7f..c1f26434 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -1,14 +1,19 @@
 from __future__ import annotations
 
+import inspect
+import logging
+import sys
+import types
 import typing as t
 from enum import Enum
 from functools import reduce
 from itertools import chain
+from pathlib import Path
 from string import Template
 
 import sqlglot
 from jinja2 import Environment
-from sqlglot import Generator, exp
+from sqlglot import Generator, exp, parse_one
 from sqlglot.executor.env import ENV
 from sqlglot.executor.python import Python
 from sqlglot.helper import csv, ensure_collection
@@ -17,6 +22,7 @@ from sqlglot.schema import MappingSchema
 from sqlmesh.core import constants as c
 from sqlmesh.core.dialect import (
     SQLMESH_MACRO_PREFIX,
+    Dialect,
     MacroDef,
     MacroFunc,
     MacroSQL,
@@ -41,6 +47,15 @@ if t.TYPE_CHECKING:
     from sqlmesh.core.snapshot import Snapshot
 
 
+if sys.version_info >= (3, 10):
+    UNION_TYPES = (t.Union, types.UnionType)
+else:
+    UNION_TYPES = (t.Union,)
+
+
+logger = logging.getLogger(__name__)
+
+
 class RuntimeStage(Enum):
     LOADING = "loading"
     CREATING = "creating"
@@ -52,6 +67,7 @@ class MacroStrTemplate(Template):
 
 
 EXPRESSIONS_NAME_MAP = {}
+SQL = t.NewType("SQL", str)
 
 for klass in sqlglot.Parser.EXPRESSION_PARSERS:
     name = klass if isinstance(klass, str) else klass.__name__  # type: ignore
@@ -125,6 +141,7 @@ class MacroEvaluator:
         resolve_tables: t.Optional[t.Callable[[exp.Expression], exp.Expression]] = None,
         snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
         default_catalog: t.Optional[str] = None,
+        path: Path = Path(),
     ):
         self.dialect = dialect
         self.generator = MacroDialect().generator()
@@ -132,7 +149,12 @@ class MacroEvaluator:
             "runtime_stage": runtime_stage.value,
             "default_catalog": default_catalog,
         }
-        self.env = {**ENV, "self": self}
+        self.env = {
+            **ENV,
+            "self": self,
+            "SQL": SQL,
+            "MacroEvaluator": MacroEvaluator,
+        }
         self.python_env = python_env or {}
         self._jinja_env: t.Optional[Environment] = jinja_env
         self.macros = {normalize_macro_name(k): v.func for k, v in macro.get_registry().items()}
@@ -141,6 +163,7 @@ class MacroEvaluator:
         self.columns_to_types_called = False
         self._snapshots = snapshots if snapshots is not None else {}
         self.default_catalog = default_catalog
+        self._path = path
 
         prepare_env(self.python_env, self.env)
         for k, v in self.python_env.items():
@@ -152,7 +175,7 @@ class MacroEvaluator:
                 self.locals[k] = self.env[k]
 
     def send(
-        self, name: str, *args: t.Any
+        self, name: str, *args: t.Any, **kwargs: t.Any
     ) -> t.Union[None, exp.Expression, t.List[exp.Expression]]:
         func = self.macros.get(normalize_macro_name(name))
 
@@ -160,7 +183,33 @@ class MacroEvaluator:
             raise SQLMeshError(f"Macro '{name}' does not exist.")
 
         try:
-            return func(self, *args)
+            annotations = t.get_type_hints(func)
+        except NameError:  # forward references aren't handled
+            annotations = {}
+
+        if annotations:
+            spec = inspect.getfullargspec(func)
+            callargs = inspect.getcallargs(func, self, *args, **kwargs)
+            new_args = []
+
+            for arg, value in callargs.items():
+                typ = annotations.get(arg)
+
+                if value is self:
+                    continue
+                if arg == spec.varargs:
+                    for v in value:
+                        new_args.append(self._coerce(v, typ))
+                elif arg == spec.varkw:
+                    for k, v in value.items():
+                        kwargs[k] = self._coerce(v, typ)
+                else:
+                    new_args.append(self._coerce(value, typ))
+
+            args = new_args  # type: ignore
+
+        try:
+            return func(self, *args, **kwargs)
         except Exception as e:
             print_exception(e, self.python_env)
             raise MacroEvalError("Error trying to eval macro.") from e
@@ -393,6 +442,81 @@ class MacroEvaluator:
         """Returns the value of the specified variable, or the default value if it doesn't exist."""
         return (self.locals.get(c.SQLMESH_VARS) or {}).get(var_name.lower(), default)
 
+    def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:
+        """Coerces the given expression to the specified type on a best-effort basis."""
+        base_err_msg = f"Failed to coerce expression '{expr}' to type '{typ}'."
+        try:
+            if typ is None or typ is t.Any:
+                return expr
+            base = t.get_origin(typ) or typ
+
+            # We need to handle Union and TypeVars first since we cannot use isinstance with it
+            if base in UNION_TYPES:
+                for branch in t.get_args(typ):
+                    try:
+                        return self._coerce(expr, branch, True)
+                    except Exception:
+                        pass
+                raise SQLMeshError(base_err_msg)
+            if base is SQL and isinstance(expr, exp.Expression):
+                return expr.sql(self.dialect)
+
+            if isinstance(expr, base):
+                return expr
+            if issubclass(base, exp.Expression):
+                d = Dialect.get_or_raise(self.dialect)
+                into = base if base in d.parser().EXPRESSION_PARSERS else None
+                if into is None:
+                    if isinstance(expr, exp.Literal):
+                        coerced = parse_one(expr.this)
+                    else:
+                        raise SQLMeshError(
+                            f"{base_err_msg} Coercion to {base} requires a literal expression."
+                        )
+                else:
+                    coerced = parse_one(
+                        expr.this if isinstance(expr, exp.Literal) else expr.sql(), into=into
+                    )
+                if isinstance(coerced, base):
+                    return coerced
+                raise SQLMeshError(base_err_msg)
+
+            if base in (int, float, str) and isinstance(expr, exp.Literal):
+                return base(expr.this)
+            if base is str and isinstance(expr, exp.Column) and not expr.table:
+                return expr.name
+            if base is bool and isinstance(expr, exp.Boolean):
+                return expr.this
+            # if base is str and isinstance(expr, exp.Expression):
+            #    return expr.sql(self.dialect)
+            if base is tuple and isinstance(expr, (exp.Tuple, exp.Array)):
+                generic = t.get_args(typ)
+                if not generic:
+                    return tuple(expr.expressions)
+                if generic[-1] is ...:
+                    return tuple(self._coerce(expr, generic[0]) for expr in expr.expressions)
+                elif len(generic) == len(expr.expressions):
+                    return tuple(
+                        self._coerce(expr, generic[i]) for i, expr in enumerate(expr.expressions)
+                    )
+                raise SQLMeshError(f"{base_err_msg} Expected {len(generic)} items.")
+            if base is list and isinstance(expr, (exp.Array, exp.Tuple)):
+                generic = t.get_args(typ)
+                if not generic:
+                    return expr.expressions
+                return [self._coerce(expr, generic[0]) for expr in expr.expressions]
+            raise SQLMeshError(base_err_msg)
+        except Exception:
+            if strict:
+                raise
+            logger.error(
+                "Coercion of expression '%s' to type '%s' failed. Using non coerced expression at '%s'",
+                expr,
+                typ,
+                self._path,
+            )
+            return expr
+
 
 class macro(registry_decorator):
     """Specifies a function is a macro and registers it the global MACROS registry.
@@ -418,7 +542,7 @@ class macro(registry_decorator):
     ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
         wrapper = super().__call__(func)
 
-        # This is useful to identify macros at runtime
+        # This is used to identify macros at runtime to unwrap during serialization.
         setattr(wrapper, "__sqlmesh_macro__", True)
         return wrapper
 
@@ -642,12 +766,13 @@ def eval_(evaluator: MacroEvaluator, condition: exp.Condition) -> t.Any:
     return evaluator.eval_expression(condition)
 
 
+# macros with union types need to use t.Union since | isn't available until 3.9
 @macro()
 def star(
     evaluator: MacroEvaluator,
     relation: exp.Table,
     alias: exp.Column = t.cast(exp.Column, exp.column("")),
-    except_: exp.Array | exp.Tuple = exp.Tuple(this=[]),
+    except_: t.Union[exp.Array, exp.Tuple] = exp.Tuple(this=[]),
     prefix: exp.Literal = exp.Literal.string(""),
     suffix: exp.Literal = exp.Literal.string(""),
     quote_identifiers: exp.Boolean = exp.true(),
@@ -710,7 +835,7 @@ def star(
 
 
 @macro()
-def generate_surrogate_key(_: MacroEvaluator, *fields: exp.Column) -> exp.Func:
+def generate_surrogate_key(_: MacroEvaluator, *fields: exp.Expression) -> exp.Func:
     """Generates a surrogate key for the given fields.
 
     Example:
@@ -735,7 +860,7 @@ def generate_surrogate_key(_: MacroEvaluator, *fields: exp.Column) -> exp.Func:
 
 
 @macro()
-def safe_add(_: MacroEvaluator, *fields: exp.Column) -> exp.Case:
+def safe_add(_: MacroEvaluator, *fields: exp.Expression) -> exp.Case:
     """Adds numbers together, substitutes nulls for 0s and only returns null if all fields are null.
 
     Example:
@@ -788,7 +913,7 @@ def safe_div(_: MacroEvaluator, numerator: exp.Expression, denominator: exp.Expr
 def union(
     evaluator: MacroEvaluator,
     type_: exp.Literal = exp.Literal.string("ALL"),
-    *tables: exp.Column,  # These represent tables but the ast node will be columns
+    *tables: exp.Table,
 ) -> exp.Query:
     """Returns a UNION of the given tables. Only choosing columns that have the same name and type.
 
@@ -870,7 +995,7 @@ def haversine_distance(
 def pivot(
     evaluator: MacroEvaluator,
     column: exp.Column,
-    values: exp.Array | exp.Tuple,
+    values: t.Union[exp.Array, exp.Tuple],
     alias: exp.Boolean = exp.true(),
     agg: exp.Literal = exp.Literal.string("SUM"),
     cmp: exp.Literal = exp.Literal.string("="),
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index a0b6457c..77decd1d 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -158,6 +158,7 @@ class BaseExpressionRenderer:
             ),
             snapshots=snapshots,
             default_catalog=self._default_catalog,
+            path=self._path,
         )
 
         for definition in self._macro_definitions:
diff --git a/sqlmesh/utils/metaprogramming.py b/sqlmesh/utils/metaprogramming.py
index e74c47bf..5efa729c 100644
--- a/sqlmesh/utils/metaprogramming.py
+++ b/sqlmesh/utils/metaprogramming.py
@@ -215,8 +215,6 @@ def normalize_source(obj: t.Any) -> str:
             # remove function return type annotation
             if isinstance(node, ast.FunctionDef):
                 node.returns = None
-        elif isinstance(node, ast.arg):
-            node.annotation = None
 
     return to_source(root_node).strip()
 
@@ -278,7 +276,7 @@ def build_env(
     if name not in env:
         # We only need to add the undecorated code of @macro() functions in env, which
         # is accessible through the `__wrapped__` attribute added by functools.wraps
-        env[name] = obj.__wrapped__ if getattr(obj, "__sqlmesh_macro__", None) else obj
+        env[name] = obj.__wrapped__ if hasattr(obj, "__sqlmesh_macro__") else obj
 
         if (
             obj_module


Test patch:
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index df73591f..0bf341e3 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -1,3 +1,4 @@
+import functools
 import sys
 import typing as t
 from unittest import mock
@@ -6,7 +7,7 @@ import pytest
 from sqlglot import MappingSchema, exp, parse_one
 
 from sqlmesh.core.dialect import StagedFilePath
-from sqlmesh.core.macros import MacroEvaluator, macro
+from sqlmesh.core.macros import SQL, MacroEvaluator, macro
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.metaprogramming import Executable, ExecutableKind
 
@@ -27,6 +28,45 @@ def macro_evaluator() -> MacroEvaluator:
     def noop(evaluator: MacroEvaluator):
         return None
 
+    @macro()
+    def bitshift_square(evaluator: MacroEvaluator, x: int, y: int) -> int:
+        return (x >> y) ** 2
+
+    @macro()
+    def prefix_db(evaluator: MacroEvaluator, table: exp.Table, prefix: str) -> exp.Table:
+        table.set("db", prefix + table.db)
+        return table
+
+    @macro()
+    def repeated(evaluator: MacroEvaluator, expr: str, times: int = 2, multi: bool = False):
+        if multi is True:
+            return (expr,) * times
+        return expr * times
+
+    @macro()
+    def split(evaluator: MacroEvaluator, string: str, sep: str = ","):
+        return string.split(sep)
+
+    @macro()
+    def cte_tag_name(evaluator: MacroEvaluator, with_: exp.Select):
+        for cte in with_.find_all(exp.CTE):
+            name = cte.alias_or_name
+            for query in cte.find_all(exp.Select):
+                query.select(exp.Literal.string(name).as_("source"), copy=False)
+        return with_
+
+    @macro()
+    def suffix_idents(evaluator: MacroEvaluator, items: t.List[str], suffix: str):
+        return [item + suffix for item in items]
+
+    @macro()
+    def suffix_idents_2(evaluator: MacroEvaluator, items: t.Tuple[str, ...], suffix: str):
+        return [item + suffix for item in items]
+
+    @macro()
+    def stamped(evaluator, query: exp.Select) -> exp.Subquery:
+        return query.select(exp.Literal.string("2024-01-01").as_("stamp")).subquery()
+
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload=f"def test(_):\n    return 'test'")},
@@ -59,12 +99,15 @@ def test_external_macro() -> None:
     def foo(evaluator: MacroEvaluator) -> str:
         return "foo"
 
+    @functools.wraps(foo)
+    def wrapper(*args, **kwargs):
+        return foo(*args, **kwargs)
+
     # Mimic a SQLMesh macro definition by setting the func's metadata appropriately
-    setattr(foo, "__wrapped__", foo)
-    setattr(foo, "__sqlmesh_macro__", True)
+    setattr(wrapper, "__sqlmesh_macro__", True)
 
     sys.modules["pkg"] = mock.Mock()
-    with mock.patch("pkg.foo", foo):
+    with mock.patch("pkg.foo", wrapper):
         evaluator = MacroEvaluator(
             python_env={
                 "foo": Executable(
@@ -320,6 +363,56 @@ def test_ast_correctness(macro_evaluator):
             "SELECT d",
             {},
         ),
+        (
+            """@BITSHIFT_SQUARE(50, '3')""",
+            "36",
+            {},
+        ),
+        (
+            """@PREFIX_DB(my.schema.table, 'dev_')""",
+            "my.dev_schema.table",
+            {},
+        ),
+        (
+            """select @REPEATED(test, 3)""",
+            "SELECT testtesttest",
+            {},
+        ),
+        (
+            """select @REPEATED(test, 3, true)""",
+            "SELECT test, test, test",
+            {},
+        ),
+        (
+            """select @SPLIT('a,b,c')""",
+            "SELECT a, b, c",
+            {},
+        ),
+        (
+            """@CTE_TAG_NAME(WITH step1 AS (SELECT 1) SELECT * FROM step1)""",
+            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
+            {},
+        ),
+        (
+            """@CTE_TAG_NAME('WITH step1 AS (SELECT 1) SELECT * FROM step1')""",
+            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
+            {},
+        ),
+        (
+            """SELECT @SUFFIX_IDENTS(['a', 'b', 'c'], 'z')""",
+            "SELECT az, bz, cz",
+            {},
+        ),
+        (
+            """SELECT @SUFFIX_IDENTS_2(['a', 'b', 'c'], 'z')""",
+            "SELECT az, bz, cz",
+            {},
+        ),
+        (
+            """select * FROM @STAMPED('select a, b, c')""",
+            "SELECT * FROM (SELECT a, b, c, '2024-01-01' AS stamp)",
+            {},
+        ),
         (
             """@OR(NULL, TRUE)""",
             "TRUE",
@@ -359,12 +452,12 @@ def test_ast_correctness(macro_evaluator):
         ),
     ],
 )
-def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):
+def test_macro_functions(macro_evaluator: MacroEvaluator, assert_exp_eq, sql, expected, args):
     macro_evaluator.locals = args or {}
     assert_exp_eq(macro_evaluator.transform(parse_one(sql)), expected)
 
 
-def test_macro_returns_none(macro_evaluator):
+def test_macro_returns_none(macro_evaluator: MacroEvaluator):
     assert macro_evaluator.transform(parse_one("@NOOP()")) is None
 
 
@@ -372,3 +465,76 @@ def test_transform_no_template(mocker, macro_evaluator):
     spy = mocker.spy(macro_evaluator, "template")
     macro_evaluator.transform(parse_one("select a, b"))
     spy.assert_not_called()
+
+
+def test_macro_coercion(macro_evaluator: MacroEvaluator, assert_exp_eq):
+    coerce = macro_evaluator._coerce
+    assert coerce(exp.Literal.number(1), int) == 1
+    assert coerce(exp.Literal.number(1.1), float) == 1.1
+    assert coerce(exp.Literal.string("Hi mom"), str) == "Hi mom"
+    assert coerce(exp.true(), bool) is True
+
+    # Coercing a string literal to a column should return a column with the same name
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Column), exp.column("order"))
+    # Not possible to coerce this string literal Cast to an exp.Column node -- so it should just return the input
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), exp.Column), exp.Literal.string("order::date")
+    )
+    # This however, is correctly coercible since it's a cast
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), exp.Cast), exp.cast(exp.column("order"), "DATE")
+    )
+
+    # Here we resolve ambiguity via the user type hint
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Identifier), exp.to_identifier("order"))
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Table), exp.table_("order"))
+
+    # Resolve a union type hint by choosing the first one that works
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), t.Union[exp.Column, exp.Cast]),
+        exp.cast(exp.column("order"), "DATE"),
+    )
+
+    # Simply ask for a string, and always get a string
+    assert coerce(exp.column("order"), str) == "order"
+
+    # this is not legal since select is too complex
+    assert isinstance(coerce(parse_one("SELECT 1"), str), exp.Select)
+
+    # From a string literal to a Select should parse the string literal, and the inverse operation works as well
+    assert_exp_eq(
+        coerce(exp.Literal.string("SELECT 1 FROM a"), exp.Select), parse_one("SELECT 1 FROM a")
+    )
+    assert coerce(parse_one("SELECT 1 FROM a"), SQL) == "SELECT 1 FROM a"
+
+    # Get a list of exp directly instead of an exp.Array
+    assert coerce(parse_one("[1, 2, 3]"), list) == [
+        exp.Literal.number(1),
+        exp.Literal.number(2),
+        exp.Literal.number(3),
+    ]
+
+    # Generics work as well, recursively resolving inner types
+    assert coerce(parse_one("[1, 2, 3]"), t.List[int]) == [1, 2, 3]
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, int, float]) == (1, 2, 3.0)
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, ...]) == (1, 2, 3)
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, str, float]) == (1, "2", 3.0)
+    assert coerce(
+        parse_one("[1, 2, [3]]"), t.Tuple[int, str, t.Union[float, t.Tuple[float, ...]]]
+    ) == (1, "2", (3.0,))
+
+    # Using exp.Expression will always return the input expression
+    assert coerce(parse_one("order", into=exp.Column), exp.Expression) == exp.column("order")
+    assert coerce(exp.Literal.string("OK"), exp.Expression) == exp.Literal.string("OK")
+
+    # Strict flag allows raising errors and is used when recursively coercing expressions
+    # otherwise, in general, we want to be lenient and just warn the user when something is not possible
+    with pytest.raises(SQLMeshError):
+        coerce(exp.Literal.string("order"), exp.Select, strict=True)
+
+    with pytest.raises(SQLMeshError):
+        _ = coerce(
+            exp.Literal.string("order::date"),
+            t.Union[exp.Column, exp.Identifier],
+            strict=True,
+        )
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 46b47a02..bf936a56 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -517,6 +517,7 @@ def test_column_descriptions(sushi_context, assert_exp_eq):
         '"memory"."sushi"."customer_revenue_by_day"'
     ].column_descriptions == {
         "customer_id": "Customer id",
+        "country code": "Customer country code, used for testing spaces",
         "revenue": "Revenue from orders made by this customer",
         "event_date": "Date",
     }
@@ -1248,6 +1249,32 @@ def test_render_query(assert_exp_eq, sushi_context):
     with pytest.raises(SQLMeshError, match=r"Macro 'times4' does not exist.*"):
         model.render_query()
 
+    assert_exp_eq(
+        sushi_context.models['"memory"."sushi"."top_waiters"'].render_query().sql(),
+        """
+        WITH "test_macros" AS (
+          SELECT
+            2 AS "lit_two",
+            "waiter_revenue_by_day"."revenue" > 0 AS "sql_exp",
+            CAST("waiter_revenue_by_day"."revenue" AS TEXT) AS "sql_lit"
+          FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"
+        )
+        SELECT
+          CAST("waiter_revenue_by_day"."waiter_id" AS INT) AS "waiter_id",
+          CAST("waiter_revenue_by_day"."revenue" AS DOUBLE) AS "revenue"
+        FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"
+        WHERE
+          "waiter_revenue_by_day"."event_date" = (
+            SELECT
+              MAX("waiter_revenue_by_day"."event_date") AS "_col_0"
+            FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"
+          )
+        ORDER BY
+          "revenue" DESC
+        LIMIT 10
+        """,
+    )
+
 
 def test_time_column():
     expressions = d.parse(
diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py
index 824c221c..99689b93 100644
--- a/tests/integrations/jupyter/test_magics.py
+++ b/tests/integrations/jupyter/test_magics.py
@@ -160,22 +160,8 @@ def test_render(
     assert output.stdout == ""
     assert output.stderr == ""
     assert len(output.outputs) == 1
-    assert convert_all_html_output_to_text(output) == [
-        """SELECT                                                                                                             
-  CAST("waiter_revenue_by_day"."waiter_id" AS INT) AS "waiter_id",                                                 
-  CAST("waiter_revenue_by_day"."revenue" AS DOUBLE) AS "revenue"                                                   
-FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"                                           
-WHERE                                                                                                              
-  "waiter_revenue_by_day"."event_date" = (                                                                         
-    SELECT                                                                                                         
-      MAX("waiter_revenue_by_day"."event_date") AS "_col_0"                                                        
-    FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"                                       
-  )                                                                                                                
-ORDER BY                                                                                                           
-  "revenue" DESC                                                                                                   
-LIMIT 10"""
-    ]
-    assert convert_all_html_output_to_tags(output) == [["pre"] + (["span"] * 93)]
+    assert len(convert_all_html_output_to_text(output)[0]) == 2212
+    assert convert_all_html_output_to_tags(output) == [["pre"] + (["span"] * 153)]
 
 
 @pytest.mark.slow
@@ -188,22 +174,8 @@ def test_render_no_format(
     assert output.stdout == ""
     assert output.stderr == ""
     assert len(output.outputs) == 1
-    assert convert_all_html_output_to_text(output) == [
-        """SELECT
-  CAST("waiter_revenue_by_day"."waiter_id" AS INT) AS "waiter_id",
-  CAST("waiter_revenue_by_day"."revenue" AS DOUBLE) AS "revenue"
-FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"
-WHERE
-  "waiter_revenue_by_day"."event_date" = (
-    SELECT
-      MAX("waiter_revenue_by_day"."event_date") AS "_col_0"
-    FROM "memory"."sushi"."waiter_revenue_by_day" AS "waiter_revenue_by_day"
-  )
-ORDER BY
-  "revenue" DESC
-LIMIT 10"""
-    ]
-    assert convert_all_html_output_to_tags(output) == [["pre"] + (["span"] * 32)]
+    assert len(convert_all_html_output_to_text(output)[0]) == 698
+    assert convert_all_html_output_to_tags(output) == [["pre"] + (["span"] * 51)]
 
 
 @pytest.mark.slow
@@ -541,7 +513,7 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
     assert len(output.outputs) == 4
     assert convert_all_html_output_to_text(output) == [
         "Models: 16",
-        "Macros: 3",
+        "Macros: 5",
         "Data warehouse connection succeeded",
         "Test connection succeeded",
     ]
@@ -570,7 +542,7 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
                     h(
                         "span",
                         {"style": f"{NEUTRAL_STYLE}; font-weight: bold"},
-                        "3",
+                        "5",
                     )
                 ),
             )
diff --git a/tests/utils/test_metaprogramming.py b/tests/utils/test_metaprogramming.py
index a469b576..9b83a24e 100644
--- a/tests/utils/test_metaprogramming.py
+++ b/tests/utils/test_metaprogramming.py
@@ -140,19 +140,19 @@ def test_func_globals() -> None:
 def test_normalize_source() -> None:
     assert (
         normalize_source(main_func)
-        == """def main_func(y):
+        == """def main_func(y: int):
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
 
-    def closure(z):
+    def closure(z: int):
         return z + Z
     return closure(y) + other_func(Y)"""
     )
 
     assert (
         normalize_source(other_func)
-        == """def other_func(a):
+        == """def other_func(a: int):
     import sqlglot
     sqlglot.parse_one('1')
     pd.DataFrame([{'x': 1}])
@@ -179,12 +179,12 @@ def test_serialize_env() -> None:
             name="main_func",
             alias="MAIN",
             path="test_metaprogramming.py",
-            payload="""def main_func(y):
+            payload="""def main_func(y: int):
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
 
-    def closure(z):
+    def closure(z: int):
         return z + Z
     return closure(y) + other_func(Y)""",
         ),
@@ -236,7 +236,7 @@ class DataClass:
         "other_func": Executable(
             name="other_func",
             path="test_metaprogramming.py",
-            payload="""def other_func(a):
+            payload="""def other_func(a: int):
     import sqlglot
     sqlglot.parse_one('1')
     pd.DataFrame([{'x': 1}])