Small schedule fixes (#4561)

* Fix schedule snapshot detach

* Remove some entitlement checks

* Move analytics tracking

* Fixes for triggers (#4562)

* Update metadata instead of overwriting

* Linting

---------

Co-authored-by: Json-Andriopoulos <Json-Andriopoulos@users.noreply.github.com>

Author: schustmi

src/zenml/analytics/utils.py

@@ -181,7 +181,9 @@ def inner_func(*args: Any, **kwargs: Any) -> Any:
                 try:
                     for obj in list(args) + list(kwargs.values()):
                         if isinstance(obj, AnalyticsTrackedModelMixin):
-                            handler.metadata = obj.get_analytics_metadata()
+                            handler.metadata.update(
+                                obj.get_analytics_metadata()
+                            )
                             break
                 except Exception as e:
                     logger.debug(f"Analytics tracking failure for {func}: {e}")
@@ -190,7 +192,9 @@ def inner_func(*args: Any, **kwargs: Any) -> Any:
 
                 try:
                     if isinstance(result, AnalyticsTrackedModelMixin):
-                        handler.metadata = result.get_analytics_metadata()
+                        handler.metadata.update(
+                            result.get_analytics_metadata()
+                        )
                 except Exception as e:
                     logger.debug(f"Analytics tracking failure for {func}: {e}")
 

src/zenml/client.py

@@ -4024,7 +4024,7 @@ def list_schedule_triggers(
         return self.zen_store.list_triggers(
             triggers_filter_model=TriggerFilter(
                 project=project or self.active_project.id,
-                user=user or self.active_user.id,
+                user=user,
                 id=id,
                 created=created,
                 updated=updated,

src/zenml/models/v2/core/triggers.py

@@ -15,7 +15,7 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Type
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Literal, Optional, Type
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
@@ -80,6 +80,17 @@ def get_extra_fields(self) -> dict[str, Any]:
         """
         pass
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Get the analytics metadata for the model.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+        metadata["type"] = self.type.value
+        metadata["flavor"] = self.flavor.value
+        return metadata
+
 
 class TriggerUpdate(TriggerBase, BaseUpdate, ABC):
     """Base trigger update class."""
@@ -434,6 +445,13 @@ class ScheduleTriggerResponseBody(ScheduleTrigger, TriggerResponseBody):
 
     next_occurrence: datetime | None = None
 
+    @model_validator(mode="after")
+    def _next_occurrence_for_inactive(self) -> "ScheduleTriggerResponseBody":
+        if not (self.active and not self.is_archived):
+            self.next_occurrence = None
+
+        return self
+
     def get_extra_fields(self) -> list[str]:
         """Returns the extra fields (e.g. next occurrence).
 

src/zenml/pub_sub/models.py

@@ -28,6 +28,7 @@ class SnapshotExecutionPayload(BaseModel):
     snapshot_id: UUID
     trigger_type: Literal["REST", "SCHEDULE", "WEBHOOK"]
     trigger_id: UUID | None = None
+    scheduled_at: datetime | None = None
 
 
 class CriticalEventType(StrEnum):

src/zenml/zen_server/routers/trigger_endpoints.py

@@ -17,8 +17,6 @@
 
 from fastapi import APIRouter, Depends, Security
 
-from zenml.analytics.enums import AnalyticsEvent
-from zenml.analytics.utils import track_handler
 from zenml.constants import (
     API,
     PIPELINE_SNAPSHOTS,
@@ -80,18 +78,10 @@ def create_trigger(
     """
     check_entitlement(feature=SCHEDULE_FEATURE)
 
-    with track_handler(
-        event=AnalyticsEvent.CREATED_TRIGGER
-    ) as analytics_handler:
-        analytics_handler.metadata = {
-            "type": trigger.type.value,
-            "flavor": trigger.flavor.value,
-        }
-
-        return verify_permissions_and_create_entity(
-            request_model=trigger,
-            create_method=zen_store().create_trigger,
-        )
+    return verify_permissions_and_create_entity(
+        request_model=trigger,
+        create_method=zen_store().create_trigger,
+    )
 
 
 @router.get(
@@ -117,8 +107,6 @@ def list_triggers(
     Returns:
         List of trigger objects.
     """
-    check_entitlement(feature=SCHEDULE_FEATURE)
-
     return verify_permissions_and_list_entities(
         filter_model=trigger_filter_model,
         resource_type=ResourceType.TRIGGER,
@@ -147,8 +135,6 @@ def get_trigger(
     Returns:
         A specific trigger object.
     """
-    check_entitlement(feature=SCHEDULE_FEATURE)
-
     return verify_permissions_and_get_entity(
         id=trigger_id,
         get_method=zen_store().get_trigger,
@@ -201,15 +187,12 @@ def delete_trigger(
         trigger_id: ID of the trigger to delete.
         soft: Soft deletion will archive the trigger.
     """
-    check_entitlement(feature=SCHEDULE_FEATURE)
-
-    with track_handler(event=AnalyticsEvent.DELETED_TRIGGER):
-        verify_permissions_and_delete_entity(
-            id=trigger_id,
-            get_method=zen_store().get_trigger,
-            delete_method=zen_store().delete_trigger,
-            soft=soft,
-        )
+    verify_permissions_and_delete_entity(
+        id=trigger_id,
+        get_method=zen_store().get_trigger,
+        delete_method=zen_store().delete_trigger,
+        soft=soft,
+    )
 
 
 @router.put(
@@ -280,8 +263,6 @@ def detach_trigger_from_snapshot(
         action=Action.READ,
     )
 
-    check_entitlement(feature=SCHEDULE_FEATURE)
-
     zen_store().detach_trigger_from_snapshot(
         trigger_id=trigger_id,
         snapshot_id=snapshot_id,

src/zenml/zen_stores/rest_zen_store.py

@@ -2282,7 +2282,7 @@ def detach_trigger_from_snapshot(
             snapshot_id: The ID of the snapshot.
         """
         self.delete(
-            path=f"{TRIGGERS}/{trigger_id}/{PIPELINE_SNAPSHOTS}/{snapshot_id}",
+            path=f"{TRIGGERS}/{trigger_id}{PIPELINE_SNAPSHOTS}/{snapshot_id}",
             timeout=5,
         )
 

src/zenml/zen_stores/sql_zen_store.py

@@ -7002,6 +7002,7 @@ def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
 
     # -------------------- Triggers ---------------------
 
+    @track_decorator(AnalyticsEvent.CREATED_TRIGGER)
     def create_trigger(
         self, trigger: TriggerRequest
     ) -> TRIGGER_RETURN_TYPE_UNION:
@@ -7149,6 +7150,7 @@ def update_trigger(
                 include_metadata=True, include_resources=True
             )
 
+    @track_decorator(event=AnalyticsEvent.DELETED_TRIGGER)
     def delete_trigger(self, trigger_id: UUID, soft: bool = True) -> None:
         """Deletes a trigger.