refactor to pass snapshot and kwargs unserialized

Author: themisvaltinos

sqlmesh/core/scheduler.py

@@ -49,7 +49,6 @@
     SQLMeshError,
     SignalEvalError,
 )
-from sqlmesh.utils.pydantic import serialize_expressions
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.context import ExecutionContext
@@ -64,27 +63,26 @@
 class SignalListener:
     def on_signal_register(
         self,
-        snapshot_id: SnapshotId,
-        signals: t.Dict[str, t.Dict[str, str]],
+        snapshots: t.List[Snapshot],
     ) -> None:
         pass
 
     def on_signal_start(
         self,
-        snapshot_id: SnapshotId,
+        snapshot: Snapshot,
         signal_name: str,
         signal_index: int,
         intervals: Intervals,
-        signal_kwargs: t.Dict[str, str],
+        signal_kwargs: t.Dict[str, t.Optional[exp.Expression]],
     ) -> None:
         pass
 
     def on_signal_end(
         self,
-        snapshot_id: SnapshotId,
+        snapshot: Snapshot,
         signal_name: str,
         signal_index: int,
-        signal_kwargs: t.Dict[str, str],
+        signal_kwargs: t.Dict[str, t.Optional[exp.Expression]],
         intervals: Intervals,
         ready_intervals: Intervals,
         error: t.Optional[Exception] = None,
@@ -339,6 +337,7 @@ def batch_intervals(
             )
             for snapshot, intervals in merged_intervals.items()
         }
+        self.signal_listener.on_signal_register(list(merged_intervals.keys()))
         snapshot_batches = {}
         all_unready_intervals: t.Dict[str, set[Interval]] = {}
         for snapshot_id in dag:
@@ -800,16 +799,6 @@ def _check_ready_intervals(
         if not (signals and signals.signals_to_kwargs):
             return intervals
 
-        signals_to_serialized_kwargs = {
-            signal_name: serialize_expressions(kwargs)
-            for signal_name, kwargs in signals.signals_to_kwargs.items()
-        }
-
-        self.signal_listener.on_signal_register(
-            snapshot_id=snapshot.snapshot_id,
-            signals=signals_to_serialized_kwargs,
-        )
-
         self.console.start_signal_progress(
             snapshot,
             self.default_catalog,
@@ -823,11 +812,11 @@ def _check_ready_intervals(
             signal_start_ts = time.perf_counter()
 
             self.signal_listener.on_signal_start(
-                snapshot_id=snapshot.snapshot_id,
+                snapshot=snapshot,
                 signal_name=signal_name,
                 signal_index=signal_idx,
                 intervals=intervals_to_check,
-                signal_kwargs=signals_to_serialized_kwargs[signal_name],
+                signal_kwargs=kwargs,
             )
 
             error = None
@@ -849,10 +838,10 @@ def _check_ready_intervals(
             finally:
                 ready_intervals = merge_intervals(intervals)
                 self.signal_listener.on_signal_end(
-                    snapshot_id=snapshot.snapshot_id,
+                    snapshot=snapshot,
                     signal_name=signal_name,
                     signal_index=signal_idx,
-                    signal_kwargs=signals_to_serialized_kwargs[signal_name],
+                    signal_kwargs=kwargs,
                     intervals=intervals_to_check,
                     ready_intervals=ready_intervals,
                     error=error,

sqlmesh/utils/pydantic.py

@@ -60,13 +60,6 @@ def _expression_encoder(e: exp.Expression) -> str:
     return e.meta.get("sql") or e.sql(dialect=e.meta.get("dialect"))
 
 
-def serialize_expressions(kwargs: t.Dict[str, t.Optional[exp.Expression]]) -> t.Dict[str, str]:
-    serialized_kwargs: t.Dict[str, str] = {}
-    for key, value in kwargs.items():
-        serialized_kwargs[key] = _expression_encoder(value) if value else str(value)
-    return serialized_kwargs
-
-
 AuditQueryTypes = t.Union[exp.Query, d.JinjaQuery]
 ModelQueryTypes = t.Union[exp.Query, d.JinjaQuery, d.MacroFunc]