fix!: add exp and SQL to type coerce so that it works for signals (#4236)

Author: tobymao

sqlmesh/core/macros.py

@@ -75,14 +75,24 @@ class MacroStrTemplate(Template):
 EXPRESSIONS_NAME_MAP = {}
 SQL = t.NewType("SQL", str)
 
-SUPPORTED_TYPES = {
-    "t": t,
-    "typing": t,
-    "List": t.List,
-    "Tuple": t.Tuple,
-    "Union": t.Union,
-    "DatetimeRanges": DatetimeRanges,
-}
+
+@lru_cache()
+def get_supported_types() -> t.Dict[str, t.Any]:
+    from sqlmesh.core.context import ExecutionContext
+
+    return {
+        "t": t,
+        "typing": t,
+        "List": t.List,
+        "Tuple": t.Tuple,
+        "Union": t.Union,
+        "DatetimeRanges": DatetimeRanges,
+        "exp": exp,
+        "SQL": SQL,
+        "MacroEvaluator": MacroEvaluator,
+        "ExecutionContext": ExecutionContext,
+    }
+
 
 for klass in sqlglot.Parser.EXPRESSION_PARSERS:
     name = klass if isinstance(klass, str) else klass.__name__  # type: ignore
@@ -1305,7 +1315,7 @@ def call_macro(
     bound.apply_defaults()
 
     try:
-        annotations = t.get_type_hints(func, localns=SUPPORTED_TYPES)
+        annotations = t.get_type_hints(func, localns=get_supported_types())
     except (NameError, TypeError):  # forward references aren't handled
         annotations = {}
 

tests/core/test_snapshot.py

@@ -21,6 +21,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import parse, parse_one
 from sqlmesh.core.environment import EnvironmentNamingInfo
+from sqlmesh.core.macros import SQL
 from sqlmesh.core.model import (
     FullKind,
     IncrementalByTimeRangeKind,
@@ -2909,8 +2910,14 @@ def test_apply_auto_restatements_disable_restatement_downstream(make_snapshot):
 
 def test_render_signal(make_snapshot, mocker):
     @signal()
-    def check_types(batch, env: str, default: int = 0):
-        if env != "in_memory" or not default == 0:
+    def check_types(batch, env: str, sql: list[SQL], table: exp.Table, default: int = 0):
+        if not (
+            env == "in_memory"
+            and default == 0
+            and isinstance(sql, list)
+            and isinstance(sql[0], str)
+            and isinstance(table, exp.Table)
+        ):
             raise
         return True
 
@@ -2919,7 +2926,7 @@ def check_types(batch, env: str, default: int = 0):
             """
         MODEL (
             name test_schema.test_model,
-            signals check_types(env := @gateway)
+            signals check_types(env := @gateway, sql := [a.b], table := b.c)
         );
         SELECT a FROM tbl;
         """