adding triggers

Author: NiveditJain

state-manager/app/config/settings.py

@@ -12,6 +12,7 @@ class Settings(BaseModel):
     mongo_database_name: str = Field(default="exosphere-state-manager", description="MongoDB database name")
     state_manager_secret: str = Field(..., description="Secret key for API authentication")
     secrets_encryption_key: str = Field(..., description="Key for encrypting secrets")
+    trigger_ahead_time: int = Field(default=10, description="Time in minutes to trigger the graph ahead of the current time")
 
     @classmethod
     def from_env(cls) -> "Settings":
@@ -20,6 +21,7 @@ def from_env(cls) -> "Settings":
             mongo_database_name=os.getenv("MONGO_DATABASE_NAME", "exosphere-state-manager"), # type: ignore
             state_manager_secret=os.getenv("STATE_MANAGER_SECRET"), # type: ignore
             secrets_encryption_key=os.getenv("SECRETS_ENCRYPTION_KEY"), # type: ignore
+            trigger_ahead_time=int(os.getenv("TRIGGER_AHEAD_TIME", 10)) # type: ignore
         )
 
 

state-manager/app/controller/upsert_graph_template.py

@@ -10,6 +10,9 @@
 
 async def upsert_graph_template(namespace_name: str, graph_name: str, body: UpsertGraphTemplateRequest, x_exosphere_request_id: str, background_tasks: BackgroundTasks) -> UpsertGraphTemplateResponse:
     try:
+
+        old_triggers = []
+
         graph_template = await GraphTemplate.find_one(
             GraphTemplate.name == graph_name,
             GraphTemplate.namespace == namespace_name
@@ -21,6 +24,7 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
                     "Graph template already exists in namespace", graph_template=graph_template,
                     namespace_name=namespace_name,
                     x_exosphere_request_id=x_exosphere_request_id)
+                old_triggers = graph_template.triggers
 
                 graph_template.set_secrets(body.secrets)
                 graph_template.validation_status = GraphTemplateValidationStatus.PENDING
@@ -54,7 +58,7 @@ async def upsert_graph_template(namespace_name: str, graph_name: str, body: Upse
             logger.error("Error validating graph template", error=e, x_exosphere_request_id=x_exosphere_request_id)
             raise HTTPException(status_code=400, detail=f"Error validating graph template: {str(e)}")
 
-        background_tasks.add_task(verify_graph, graph_template)
+        background_tasks.add_task(verify_graph, graph_template, old_triggers)
 
         return UpsertGraphTemplateResponse(
             nodes=graph_template.nodes,

state-manager/app/models/db/trigger.py

@@ -0,0 +1,24 @@
+from pydantic import Field
+from beanie import Document
+from typing import Optional
+
+from pymongo import IndexModel
+from ..trigger_models import TriggerTypeEnum, TriggerStatusEnum
+from datetime import datetime
+
+class Triggers(Document):
+    type: TriggerTypeEnum = Field(..., description="Type of the trigger")
+    expression: Optional[str] = Field(default=None, description="Expression of the trigger")
+    graph_name: str = Field(..., description="Name of the graph")
+    trigger_time: datetime = Field(..., description="Trigger time of the trigger")
+    trigger_status: TriggerStatusEnum = Field(..., description="Status of the trigger")
+
+    class Settings: 
+        indexes = [
+            IndexModel(
+                [
+                    ("trigger_time", -1),
+                ],
+                name="idx_trigger_time"
+            )
+        ]

state-manager/app/models/trigger_model.py …anager/app/models/trigger_graph_model.pystate-manager/app/models/trigger_model.py renamed to state-manager/app/models/trigger_graph_model.py state-manager/app/models/trigger_model.py renamed to state-manager/app/models/trigger_graph_model.py

@@ -6,6 +6,12 @@
 class TriggerTypeEnum(str, Enum):
     CRON = "CRON"
 
+class TriggerStatusEnum(str, Enum):
+    PENDING = "PENDING"
+    FAILED = "FAILED"
+    CANCELLED = "CANCELLED"
+    TRIGGERED = "TRIGGERED"
+
 class CronTrigger(BaseModel):
     expression: str = Field(..., description="Cron expression for the trigger")
 

state-manager/app/models/trigger_models.py

@@ -9,7 +9,7 @@
 from .models.enqueue_request import EnqueueRequestModel
 from .controller.enqueue_states import enqueue_states
 
-from .models.trigger_model import TriggerGraphRequestModel, TriggerGraphResponseModel
+from .models.trigger_graph_model import TriggerGraphRequestModel, TriggerGraphResponseModel
 from .controller.trigger_graph import trigger_graph
 
 from .models.executed_models import ExecutedRequestModel, ExecutedResponseModel

state-manager/app/routes.py

@@ -1,10 +1,17 @@
 import asyncio
+import croniter
+
+from datetime import datetime, timedelta
+from beanie.operators import In
+from json_schema_to_pydantic import create_model
 
 from app.models.db.graph_template_model import GraphTemplate
 from app.models.graph_template_validation_status import GraphTemplateValidationStatus
 from app.models.db.registered_node import RegisteredNode
 from app.singletons.logs_manager import LogsManager
-from json_schema_to_pydantic import create_model
+from app.models.trigger_models import Trigger, CronTrigger, TriggerStatusEnum, TriggerTypeEnum
+from app.models.db.trigger import Triggers as DatabaseTriggers
+from config.settings import get_settings
 
 logger = LogsManager().get_logger()
 
@@ -95,7 +102,54 @@ async def verify_inputs(graph_template: GraphTemplate, registered_nodes: list[Re
 
     return errors
 
-async def verify_graph(graph_template: GraphTemplate):
+async def cancel_crons(graph_template: GraphTemplate, old_triggers: list[Trigger]):
+    old_crons = set([CronTrigger(**trigger.value) for trigger in old_triggers if trigger.type == TriggerTypeEnum.CRON])
+    new_crons = set([CronTrigger(**trigger.value) for trigger in graph_template.triggers if trigger.type == TriggerTypeEnum.CRON])
+
+    removed = old_crons - new_crons
+
+    await DatabaseTriggers.find(
+        DatabaseTriggers.graph_name == graph_template.name,
+        DatabaseTriggers.trigger_status == TriggerStatusEnum.PENDING,
+        DatabaseTriggers.type == TriggerTypeEnum.CRON,
+        In(DatabaseTriggers.expression, [cron.expression for cron in removed])
+    ).update(
+        {
+            "$set": {
+                "trigger_status": TriggerStatusEnum.CANCELLED
+            }
+        }
+    ) # type: ignore
+
+async def create_crons(graph_template: GraphTemplate, old_triggers: list[Trigger]):
+    old_crons = set([CronTrigger(**trigger.value) for trigger in old_triggers if trigger.type == TriggerTypeEnum.CRON])
+    new_crons = set([CronTrigger(**trigger.value) for trigger in graph_template.triggers if trigger.type == TriggerTypeEnum.CRON])
+
+    crons_to_create = new_crons - old_crons
+
+    trigger_ahead_time = get_settings().trigger_ahead_time
+    current_time = datetime.now()
+    limit_time = datetime.now() + timedelta(minutes=trigger_ahead_time)
+    
+    new_db_triggers = []
+    for cron in crons_to_create:
+        iter = croniter.croniter(cron.expression, current_time)
+        next_trigger_time = iter.get_next(datetime)
+
+        while (next_trigger_time < limit_time):
+            new_db_triggers.append(
+                DatabaseTriggers(
+                    type=TriggerTypeEnum.CRON,
+                    expression=cron.expression,
+                    graph_name=graph_template.name,
+                    trigger_status=TriggerStatusEnum.PENDING,
+                    trigger_time=next_trigger_time
+                )
+            )
+            next_trigger_time = iter.get_next(datetime)
+    await DatabaseTriggers.insert_many(new_db_triggers)
+
+async def verify_graph(graph_template: GraphTemplate, old_triggers: list[Trigger]):
     try:
         errors = []
         registered_nodes = await RegisteredNode.list_nodes_by_templates(graph_template.nodes)
@@ -118,6 +172,9 @@ async def verify_graph(graph_template: GraphTemplate):
 
         graph_template.validation_status = GraphTemplateValidationStatus.VALID
         graph_template.validation_errors = []
+
+        await asyncio.gather(*[cancel_crons(graph_template, old_triggers), create_crons(graph_template, old_triggers)])
+
         await graph_template.save()
 
     except Exception as e: