Repository: TobikoData/sqlmesh
Pull request #4236 (2025-04-23 19:34:18)

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


Patch:
diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py
index fd9c9e89..8febdec8 100644
--- a/sqlmesh/core/console.py
+++ b/sqlmesh/core/console.py
@@ -891,7 +891,7 @@ class TerminalConsole(Console):
                     self.evaluation_column_widths["duration"]
                 )
 
-                msg = f"{batch} {display_name}  {annotation} {duration}".replace(
+                msg = f"{batch} {display_name}   {annotation}   {duration}".replace(
                     CHECK_MARK, GREEN_CHECK_MARK
                 )
 
@@ -3359,54 +3359,56 @@ def _create_evaluation_model_annotation(snapshot: Snapshot, interval_info: t.Opt
     return interval_info if interval_info else ""
 
 
-def _calculate_interval_str_len(batched_intervals: t.Dict[Snapshot, t.List[Interval]]) -> int:
+def _calculate_interval_str_len(snapshot: Snapshot, intervals: t.List[Interval]) -> int:
     interval_str_len = 0
-    for snapshot, intervals in batched_intervals.items():
-        for interval in intervals:
-            interval_str_len = max(
-                interval_str_len,
-                len(
-                    _create_evaluation_model_annotation(
-                        snapshot, _format_evaluation_model_interval(snapshot, interval)
-                    )
-                ),
-            )
+    for interval in intervals:
+        interval_str_len = max(
+            interval_str_len,
+            len(
+                _create_evaluation_model_annotation(
+                    snapshot, _format_evaluation_model_interval(snapshot, interval)
+                )
+            ),
+        )
     return interval_str_len
 
 
-def _calculate_audit_str_len(batched_intervals: t.Dict[Snapshot, t.List[Interval]]) -> int:
+def _calculate_audit_str_len(snapshot: Snapshot) -> int:
     # The annotation includes audit results. We cannot build the audits result string
     # until after evaluation occurs, but we must determine the annotation column width here.
     # Therefore, we add enough padding for the longest possible audits result string.
     audit_str_len = 0
     audit_base_str_len = len(f", audits ") + 1  # +1 for check/X
-    for snapshot in batched_intervals:
-        if snapshot.is_audit:
+    if snapshot.is_audit:
+        # +1 for "1" audit count, +1 for red X
+        audit_str_len = max(
+            audit_str_len, audit_base_str_len + (2 if not snapshot.audit.blocking else 1)
+        )
+    if snapshot.is_model and snapshot.model.audits:
+        num_audits = len(snapshot.model.audits_with_args)
+        num_nonblocking_audits = sum(
+            1
+            for audit in snapshot.model.audits_with_args
+            if not audit[0].blocking
+            or ("blocking" in audit[1] and audit[1]["blocking"] == exp.false())
+        )
+        if num_audits == 1:
             # +1 for "1" audit count, +1 for red X
-            audit_str_len = max(
-                audit_str_len, audit_base_str_len + (2 if not snapshot.audit.blocking else 1)
-            )
-        if snapshot.is_model and snapshot.model.audits:
-            num_audits = len(snapshot.model.audits_with_args)
-            num_nonblocking_audits = sum(
-                1
-                for audit in snapshot.model.audits_with_args
-                if not audit[0].blocking
-                or ("blocking" in audit[1] and audit[1]["blocking"] == exp.false())
-            )
-            if num_audits == 1:
-                # +1 for "1" audit count, +1 for red X
-                audit_len = audit_base_str_len + (2 if num_nonblocking_audits else 1)
-            else:
-                audit_len = audit_base_str_len + len(str(num_audits))
-                if num_nonblocking_audits:
-                    # +1 for space, +1 for red X
-                    audit_len += len(str(num_nonblocking_audits)) + 2
-            audit_str_len = max(audit_str_len, audit_len)
+            audit_len = audit_base_str_len + (2 if num_nonblocking_audits else 1)
+        else:
+            audit_len = audit_base_str_len + len(str(num_audits))
+            if num_nonblocking_audits:
+                # +1 for space, +1 for red X
+                audit_len += len(str(num_nonblocking_audits)) + 2
+        audit_str_len = max(audit_str_len, audit_len)
     return audit_str_len
 
 
 def _calculate_annotation_str_len(batched_intervals: t.Dict[Snapshot, t.List[Interval]]) -> int:
-    return _calculate_interval_str_len(batched_intervals) + _calculate_audit_str_len(
-        batched_intervals
-    )
+    annotation_str_len = 0
+    for snapshot, intervals in batched_intervals.items():
+        annotation_str_len = max(
+            annotation_str_len,
+            _calculate_interval_str_len(snapshot, intervals) + _calculate_audit_str_len(snapshot),
+        )
+    return annotation_str_len
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 88728457..5bbcbf21 100644
--- a/sqlmesh/core/macros.py
+++ b/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 = {}
 


Test patch:
diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
index 12b9e834..cf911c5f 100644
--- a/tests/cli/test_cli.py
+++ b/tests/cli/test_cli.py
@@ -182,7 +182,7 @@ def test_plan(runner, tmp_path):
     assert_plan_success(result)
     # 'Models needing backfill' section and eval progress bar should display the same inclusive intervals
     assert "sqlmesh_example.incremental_model: [2020-01-01 - 2022-12-31]" in result.output
-    assert "sqlmesh_example.incremental_model  [insert 2020-01-01 - 2022-12-31]" in result.output
+    assert "sqlmesh_example.incremental_model   [insert 2020-01-01 - 2022-12-31]" in result.output
 
 
 def test_plan_skip_tests(runner, tmp_path):
@@ -243,7 +243,7 @@ def test_plan_restate_model(runner, tmp_path):
     assert result.exit_code == 0
     assert_duckdb_test(result)
     assert "No changes to plan: project files match the `prod` environment" in result.output
-    assert "sqlmesh_example.full_model  [full refresh" in result.output
+    assert "sqlmesh_example.full_model   [full refresh" in result.output
     assert_model_batches_executed(result)
     assert_virtual_layer_updated(result)
 
@@ -553,6 +553,56 @@ def test_plan_dev_no_changes(runner, tmp_path):
     assert_virtual_layer_updated(result)
 
 
+def test_plan_dev_longnames(runner, tmp_path):
+    create_example_project(tmp_path)
+
+    long_model_names = {
+        "full": f"full_{'a' * 80}",
+        "incremental": f"incremental_{'b' * 80}",
+        "seed": f"seed_{'c' * 80}",
+    }
+    for model_name in long_model_names:
+        with open(tmp_path / "models" / f"{model_name}_model.sql", "r") as f:
+            model_text = f.read()
+            for more_model_names in long_model_names:
+                model_text = model_text.replace(
+                    f"sqlmesh_example.{more_model_names}_model",
+                    f"sqlmesh_example.{long_model_names[more_model_names]}_model",
+                )
+        with open(tmp_path / "models" / f"{model_name}_model.sql", "w") as f:
+            f.write(model_text)
+
+    # Input: `y` to apply and backfill
+    result = runner.invoke(
+        cli,
+        [
+            "--log-file-dir",
+            tmp_path,
+            "--paths",
+            tmp_path,
+            "plan",
+            "dev_butamuchlongerenvironmentname",
+            "--skip-tests",
+            "--no-prompts",
+            "--auto-apply",
+        ],
+    )
+    assert result.exit_code == 0
+    assert (
+        "sqlmesh_example__dev_butamuchlongerenvironmentname.seed_cccccccccccccccccccccccc\ncccccccccccccccccccccccccccccccccccccccccccccccccccccccc_model          [insert \nseed file]"
+        in result.output
+    )
+    assert (
+        "sqlmesh_example__dev_butamuchlongerenvironmentname.incremental_bbbbbbbbbbbbbbbbb\nbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb_model   [insert "
+        in result.output
+    )
+    assert (
+        "sqlmesh_example__dev_butamuchlongerenvironmentname.full_aaaaaaaaaaaaaaaaaaaaaaaa\naaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_model          [full \nrefresh"
+        in result.output
+    )
+    assert_backfill_success(result)
+
+
 def test_plan_nonbreaking(runner, tmp_path):
     create_example_project(tmp_path)
     init_prod_and_backfill(runner, tmp_path)
@@ -568,8 +618,8 @@ def test_plan_nonbreaking(runner, tmp_path):
     assert "+  'a' AS new_col" in result.output
     assert "Directly Modified: sqlmesh_example.incremental_model (Non-breaking)" in result.output
     assert "sqlmesh_example.full_model (Indirect Non-breaking)" in result.output
-    assert "sqlmesh_example.incremental_model  [insert" in result.output
-    assert "sqlmesh_example.full_model  [full refresh" not in result.output
+    assert "sqlmesh_example.incremental_model   [insert" in result.output
+    assert "sqlmesh_example.full_model   [full refresh" not in result.output
     assert_backfill_success(result)
 
 
@@ -626,8 +676,8 @@ def test_plan_breaking(runner, tmp_path):
     assert result.exit_code == 0
     assert "+  item_id + 1 AS item_id," in result.output
     assert "Directly Modified: sqlmesh_example.full_model (Breaking)" in result.output
-    assert "sqlmesh_example.full_model  [full refresh" in result.output
-    assert "sqlmesh_example.incremental_model  [insert" not in result.output
+    assert "sqlmesh_example.full_model   [full refresh" in result.output
+    assert "sqlmesh_example.incremental_model   [insert" not in result.output
     assert_backfill_success(result)
 
 
@@ -665,8 +715,8 @@ def test_plan_dev_select(runner, tmp_path):
     assert "+  item_id + 1 AS item_id," not in result.output
     assert "Directly Modified: sqlmesh_example__dev.full_model (Breaking)" not in result.output
     # only incremental_model backfilled
-    assert "sqlmesh_example__dev.incremental_model  [insert" in result.output
-    assert "sqlmesh_example__dev.full_model  [full refresh" not in result.output
+    assert "sqlmesh_example__dev.incremental_model   [insert" in result.output
+    assert "sqlmesh_example__dev.full_model   [full refresh" not in result.output
     assert_backfill_success(result)
 
 
@@ -704,8 +754,8 @@ def test_plan_dev_backfill(runner, tmp_path):
         "Directly Modified: sqlmesh_example__dev.incremental_model (Non-breaking)" in result.output
     )
     # only incremental_model backfilled
-    assert "sqlmesh_example__dev.incremental_model  [insert" in result.output
-    assert "sqlmesh_example__dev.full_model  [full refresh" not in result.output
+    assert "sqlmesh_example__dev.incremental_model   [insert" in result.output
+    assert "sqlmesh_example__dev.full_model   [full refresh" not in result.output
     assert_backfill_success(result)
 
 
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index 370d70b0..720a3563 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -21,6 +21,7 @@ from sqlmesh.core.config import (
 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 test_render_signal(make_snapshot, mocker):
             """
         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;
         """