feat: add TTL-based cleanup strategy for DatabaseTriggers collection

python-sdk/exospherehost/_version.py

@@ -1 +1 @@
-version = "0.0.3b1"
+version = "0.0.3b2"

state-manager/app/config/settings.py

@@ -6,13 +6,14 @@
 
 class Settings(BaseModel):
     """Application settings loaded from environment variables."""
-    
+
     # MongoDB Configuration
     mongo_uri: str = Field(..., description="MongoDB connection URI" )
     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_workers: int = Field(default=1, description="Number of workers to run the trigger cron")
+    trigger_retention_days: int = Field(default=30, description="Number of days to retain completed/failed triggers before cleanup")
 
     @classmethod
     def from_env(cls) -> "Settings":
@@ -21,7 +22,8 @@ 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_workers=int(os.getenv("TRIGGER_WORKERS", 1)) # type: ignore
+            trigger_workers=int(os.getenv("TRIGGER_WORKERS", 1)), # type: ignore
+            trigger_retention_days=int(os.getenv("TRIGGER_RETENTION_DAYS", 30)) # type: ignore
         )
 
 

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

@@ -13,8 +13,9 @@ class DatabaseTriggers(Document):
     namespace: str = Field(..., description="Namespace of the graph")
     trigger_time: datetime = Field(..., description="Trigger time of the trigger")
     trigger_status: TriggerStatusEnum = Field(..., description="Status of the trigger")
+    expires_at: Optional[datetime] = Field(default=None, description="Expiration time for automatic cleanup of completed triggers")
 
-    class Settings: 
+    class Settings:
         indexes = [
             IndexModel(
                 [
@@ -32,5 +33,12 @@ class Settings:
                 ],
                 name="uniq_graph_type_expr_time",
                 unique=True
+            ),
+            IndexModel(
+                [
+                    ("expires_at", 1),
+                ],
+                name="ttl_expires_at",
+                expireAfterSeconds=0  # Delete immediately when expires_at is reached
             )
         ]

state-manager/app/tasks/trigger_cron.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
 from uuid import uuid4
 from app.models.db.trigger import DatabaseTriggers
 from app.models.trigger_models import TriggerStatusEnum, TriggerTypeEnum
@@ -34,10 +34,26 @@ async def call_trigger_graph(trigger: DatabaseTriggers):
         x_exosphere_request_id=str(uuid4())
     )
 
-async def mark_as_failed(trigger: DatabaseTriggers):
+async def mark_as_failed(trigger: DatabaseTriggers, retention_days: int):
+    expires_at = datetime.now(timezone.utc) + timedelta(days=retention_days)
+
     await DatabaseTriggers.get_pymongo_collection().update_one(
         {"_id": trigger.id},
-        {"$set": {"trigger_status": TriggerStatusEnum.FAILED}}
+        {"$set": {
+            "trigger_status": TriggerStatusEnum.FAILED,
+            "expires_at": expires_at
+        }}
+    )
+
+async def mark_as_cancelled(trigger: DatabaseTriggers, retention_days: int):
+    expires_at = datetime.now(timezone.utc) + timedelta(days=retention_days)
+
+    await DatabaseTriggers.get_pymongo_collection().update_one(
+        {"_id": trigger.id},
+        {"$set": {
+            "trigger_status": TriggerStatusEnum.CANCELLED,
+            "expires_at": expires_at
+        }}
     )
 
 async def create_next_triggers(trigger: DatabaseTriggers, cron_time: datetime):
@@ -65,24 +81,30 @@ async def create_next_triggers(trigger: DatabaseTriggers, cron_time: datetime):
         if next_trigger_time > cron_time:
             break
 
-async def mark_as_triggered(trigger: DatabaseTriggers):
+async def mark_as_triggered(trigger: DatabaseTriggers, retention_days: int):
+    expires_at = datetime.now(timezone.utc) + timedelta(days=retention_days)
+
     await DatabaseTriggers.get_pymongo_collection().update_one(
         {"_id": trigger.id},
-        {"$set": {"trigger_status": TriggerStatusEnum.TRIGGERED}}
+        {"$set": {
+            "trigger_status": TriggerStatusEnum.TRIGGERED,
+            "expires_at": expires_at
+        }}
     )
 
-async def handle_trigger(cron_time: datetime):
+async def handle_trigger(cron_time: datetime, retention_days: int):
     while(trigger:= await get_due_triggers(cron_time)):
         try:
             await call_trigger_graph(trigger)
-            await mark_as_triggered(trigger)
+            await mark_as_triggered(trigger, retention_days)
         except Exception as e:
-            await mark_as_failed(trigger)
+            await mark_as_failed(trigger, retention_days)
             logger.error(f"Error calling trigger graph: {e}")
         finally:
             await create_next_triggers(trigger, cron_time)
 
 async def trigger_cron():
     cron_time = datetime.now()
+    settings = get_settings()
     logger.info(f"starting trigger_cron: {cron_time}")
-    await asyncio.gather(*[handle_trigger(cron_time) for _ in range(get_settings().trigger_workers)])
+    await asyncio.gather(*[handle_trigger(cron_time, settings.trigger_retention_days) for _ in range(settings.trigger_workers)])

state-manager/tests/unit/tasks/test_trigger_ttl.py

@@ -0,0 +1,85 @@
+"""
+Tests for trigger TTL (Time To Live) expiration logic.
+Verifies that completed/failed triggers are properly marked for cleanup.
+"""
+import pytest
+from unittest.mock import MagicMock, AsyncMock, patch
+from datetime import datetime, timedelta, timezone
+
+from app.tasks.trigger_cron import mark_as_triggered, mark_as_failed, mark_as_cancelled
+from app.models.db.trigger import DatabaseTriggers
+from app.models.trigger_models import TriggerStatusEnum
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mark_function,expected_status", [
+    (mark_as_triggered, TriggerStatusEnum.TRIGGERED),
+    (mark_as_failed, TriggerStatusEnum.FAILED),
+    (mark_as_cancelled, TriggerStatusEnum.CANCELLED),
+])
+async def test_mark_trigger_sets_expires_at(mark_function, expected_status):
+    """Test that marking a trigger sets the expires_at field correctly"""
+    # Create a mock trigger
+    trigger = MagicMock(spec=DatabaseTriggers)
+    trigger.id = "test_trigger_id"
+
+    # Mock the database update
+    with patch.object(DatabaseTriggers, 'get_pymongo_collection') as mock_collection:
+        mock_collection.return_value.update_one = AsyncMock()
+
+        # Call the function with retention_days parameter
+        await mark_function(trigger, retention_days=30)
+
+        # Verify update_one was called
+        assert mock_collection.return_value.update_one.called
+        call_args = mock_collection.return_value.update_one.call_args
+
+        # Verify the filter (first argument)
+        assert call_args[0][0] == {"_id": trigger.id}
+
+        # Verify the update includes both status and expires_at
+        update_dict = call_args[0][1]["$set"]
+        assert update_dict["trigger_status"] == expected_status
+        assert "expires_at" in update_dict
+
+        # Verify expires_at is approximately 30 days from now (UTC)
+        expires_at = update_dict["expires_at"]
+        expected_expiry = datetime.now(timezone.utc) + timedelta(days=30)
+        time_diff = abs((expires_at - expected_expiry).total_seconds())
+        assert time_diff < 2  # Within 2 seconds tolerance
+
+        # Verify expires_at is timezone-aware UTC
+        assert expires_at.tzinfo is not None
+        assert expires_at.tzinfo == timezone.utc
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mark_function,retention_days", [
+    (mark_as_triggered, 7),
+    (mark_as_failed, 14),
+    (mark_as_cancelled, 21),
+])
+async def test_mark_trigger_uses_custom_retention_period(mark_function, retention_days):
+    """Test that custom retention period is respected"""
+    # Create a mock trigger
+    trigger = MagicMock(spec=DatabaseTriggers)
+    trigger.id = "test_trigger_id"
+
+    # Mock the database update
+    with patch.object(DatabaseTriggers, 'get_pymongo_collection') as mock_collection:
+        mock_collection.return_value.update_one = AsyncMock()
+
+        # Call the function with custom retention period
+        await mark_function(trigger, retention_days=retention_days)
+
+        # Verify expires_at is approximately retention_days from now (UTC)
+        call_args = mock_collection.return_value.update_one.call_args
+        update_dict = call_args[0][1]["$set"]
+        expires_at = update_dict["expires_at"]
+        expected_expiry = datetime.now(timezone.utc) + timedelta(days=retention_days)
+        time_diff = abs((expires_at - expected_expiry).total_seconds())
+        assert time_diff < 2  # Within 2 seconds tolerance
+
+        # Verify expires_at is timezone-aware UTC
+        assert expires_at.tzinfo is not None
+        assert expires_at.tzinfo == timezone.utc