Repository: TobikoData/sqlmesh
Pull request #2756 (2024-06-07 19:08:15)

Chore: Simplify notification targets by removing notify decorator


Patch:
diff --git a/sqlmesh/core/notification_target.py b/sqlmesh/core/notification_target.py
index 55269c10..7eaf1c3a 100644
--- a/sqlmesh/core/notification_target.py
+++ b/sqlmesh/core/notification_target.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import functools
 import smtplib
 import sys
 import typing as t
@@ -37,9 +36,6 @@ def _sqlmesh_version() -> str:
         return "0.0.0"
 
 
-NOTIFICATION_FUNCTIONS: t.Dict[NotificationEvent, str] = {}
-
-
 class NotificationStatus(str, Enum):
     SUCCESS = "success"
     FAILURE = "failure"
@@ -81,25 +77,13 @@ class NotificationEvent(str, Enum):
     MIGRATION_FAILURE = "migration_failure"
 
 
-def notify(event: NotificationEvent) -> t.Callable:
-    """Decorator used to register 'notify' methods and the events they correspond to."""
-
-    def decorator(f: t.Callable) -> t.Callable:
-        @functools.wraps(f)
-        def wrapper(*args: t.List[t.Any], **kwargs: t.Dict[str, t.Any]) -> None:
-            return f(*args, **kwargs)
-
-        NOTIFICATION_FUNCTIONS[event] = f.__name__
-        return wrapper
-
-    return decorator
-
-
 class BaseNotificationTarget(PydanticModel, frozen=True):
     """
     Base notification target model. Provides a command for sending notifications that is currently only used
     by the built-in scheduler. Other schedulers like Airflow use the configuration of the target itself
     to create the notification constructs appropriate for the scheduler.
+
+    Notification functions follow a naming convention of `notify_` + NotificationEvent value.
     """
 
     type_: str
@@ -113,7 +97,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
             msg: The message to send.
         """
 
-    @notify(NotificationEvent.APPLY_START)
     def notify_apply_start(self, environment: str, plan_id: str) -> None:
         """Notify when an apply starts.
 
@@ -126,7 +109,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
             f"Plan `{plan_id}` apply started for environment `{environment}`.",
         )
 
-    @notify(NotificationEvent.APPLY_END)
     def notify_apply_end(self, environment: str, plan_id: str) -> None:
         """Notify when an apply ends.
 
@@ -139,7 +121,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
             f"Plan `{plan_id}` apply finished for environment `{environment}`.",
         )
 
-    @notify(NotificationEvent.RUN_START)
     def notify_run_start(self, environment: str) -> None:
         """Notify when a SQLMesh run starts.
 
@@ -148,7 +129,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
         """
         self.send(NotificationStatus.INFO, f"SQLMesh run started for environment `{environment}`.")
 
-    @notify(NotificationEvent.RUN_END)
     def notify_run_end(self, environment: str) -> None:
         """Notify when a SQLMesh run ends.
 
@@ -159,17 +139,14 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
             NotificationStatus.SUCCESS, f"SQLMesh run finished for environment `{environment}`."
         )
 
-    @notify(NotificationEvent.MIGRATION_START)
     def notify_migration_start(self) -> None:
         """Notify when a SQLMesh migration starts."""
         self.send(NotificationStatus.INFO, "SQLMesh migration started.")
 
-    @notify(NotificationEvent.MIGRATION_END)
     def notify_migration_end(self) -> None:
         """Notify when a SQLMesh migration ends."""
         self.send(NotificationStatus.SUCCESS, "SQLMesh migration finished.")
 
-    @notify(NotificationEvent.APPLY_FAILURE)
     def notify_apply_failure(self, environment: str, plan_id: str, exc: str) -> None:
         """Notify in the case of an apply failure.
 
@@ -184,7 +161,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
             exc=exc,
         )
 
-    @notify(NotificationEvent.RUN_FAILURE)
     def notify_run_failure(self, exc: str) -> None:
         """Notify in the case of a run failure.
 
@@ -193,7 +169,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
         """
         self.send(NotificationStatus.FAILURE, "SQLMesh run failed.", exc=exc)
 
-    @notify(NotificationEvent.AUDIT_FAILURE)
     def notify_audit_failure(self, audit_error: AuditError) -> None:
         """Notify in the case of an audit failure.
 
@@ -202,7 +177,6 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
         """
         self.send(NotificationStatus.FAILURE, "Audit failure.", audit_error=audit_error)
 
-    @notify(NotificationEvent.MIGRATION_FAILURE)
     def notify_migration_failure(self, exc: str) -> None:
         """Notify in the case of a migration failure.
 
@@ -470,6 +444,5 @@ class NotificationTargetManager:
     def _get_notification_function(
         self, notification_target: NotificationTarget, event: NotificationEvent
     ) -> t.Callable:
-        """Lookup the registered function for a notification event"""
-        func_name = NOTIFICATION_FUNCTIONS[event]
-        return getattr(notification_target, func_name)
+        """Fetch the function for a notification event"""
+        return getattr(notification_target, f"notify_{event.value}")
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 4c9a9c5a..b3981e11 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -1259,13 +1259,16 @@ class DeployabilityIndex(PydanticModel, frozen=True):
             if deployable and node in snapshots:
                 snapshot = snapshots[node]
                 # Capture uncategorized snapshot which represents a forward-only model.
-                is_forward_only_model = (
-                    snapshot.previous_versions and snapshot.is_model and snapshot.model.forward_only
+                is_uncategorized_forward_only_model = (
+                    snapshot.change_category is None
+                    and snapshot.previous_versions
+                    and snapshot.is_model
+                    and snapshot.model.forward_only
                 )
                 if (
                     snapshot.is_forward_only
                     or snapshot.is_indirect_non_breaking
-                    or is_forward_only_model
+                    or is_uncategorized_forward_only_model
                 ):
                     # FORWARD_ONLY and INDIRECT_NON_BREAKING snapshots are not deployable by nature.
                     this_deployable = False
@@ -1275,7 +1278,8 @@ class DeployabilityIndex(PydanticModel, frozen=True):
                 else:
                     this_deployable = True
                 children_deployable = not (
-                    snapshot.is_paused and (snapshot.is_forward_only or is_forward_only_model)
+                    snapshot.is_paused
+                    and (snapshot.is_forward_only or is_uncategorized_forward_only_model)
                 )
             else:
                 this_deployable, children_deployable = False, False


Test patch:
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index 88860a3a..80724e01 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -1637,15 +1637,46 @@ def test_deployability_index_uncategorized_forward_only_model(make_snapshot):
     snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("SELECT 1")))
     snapshot_b.parents = (snapshot_a.snapshot_id,)
 
-    deplyability_index = DeployabilityIndex.create(
+    deployability_index = DeployabilityIndex.create(
         {s.snapshot_id: s for s in [snapshot_a, snapshot_b]}
     )
 
-    assert not deplyability_index.is_deployable(snapshot_a)
-    assert not deplyability_index.is_deployable(snapshot_b)
+    assert not deployability_index.is_deployable(snapshot_a)
+    assert not deployability_index.is_deployable(snapshot_b)
+
+    assert not deployability_index.is_representative(snapshot_a)
+    assert not deployability_index.is_representative(snapshot_b)
+
+
+def test_deployability_index_categorized_forward_only_model(make_snapshot):
+    model_a = SqlModel(
+        name="a",
+        query=parse_one("SELECT 1, ds"),
+        kind=IncrementalByTimeRangeKind(time_column="ds", forward_only=True),
+    )
+
+    snapshot_a_old = make_snapshot(model_a)
+    snapshot_a_old.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    snapshot_a = make_snapshot(model_a)
+    snapshot_a.previous_versions = snapshot_a_old.all_versions
+    snapshot_a.categorize_as(SnapshotChangeCategory.METADATA)
+
+    snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("SELECT 1")))
+    snapshot_b.parents = (snapshot_a.snapshot_id,)
+    snapshot_b.categorize_as(SnapshotChangeCategory.METADATA)
+
+    # The fact that the model is forward only should be ignored if an actual category
+    # has been assigned.
+    deployability_index = DeployabilityIndex.create(
+        {s.snapshot_id: s for s in [snapshot_a, snapshot_b]}
+    )
 
-    assert not deplyability_index.is_representative(snapshot_a)
-    assert not deplyability_index.is_representative(snapshot_b)
+    assert deployability_index.is_deployable(snapshot_a)
+    assert deployability_index.is_deployable(snapshot_b)
+
+    assert deployability_index.is_representative(snapshot_a)
+    assert deployability_index.is_representative(snapshot_b)
 
 
 def test_deployability_index_missing_parent(make_snapshot):