""" Document Schedule Service Manages scheduled document status changes using Redis """ import json import logging from datetime import datetime from enum import Enum from typing import Any, Optional from extensions.ext_redis import redis_client logger = logging.getLogger(__name__) class ScheduleAction(str, Enum): """Scheduled action types""" ENABLE = "enable" DISABLE = "disable" ARCHIVE = "archive" class DocumentScheduleService: """Service for managing scheduled document status changes""" # Redis key prefixes SCHEDULE_KEY_PREFIX = "document_schedule:" # Sorted Set for scheduled tasks SCHEDULE_DETAIL_PREFIX = "document_schedule_detail:" # Hash for task details @classmethod def add_schedule( cls, document_id: str, dataset_id: str, tenant_id: str, action: ScheduleAction, scheduled_at: datetime, created_by: str, metadata: Optional[dict] = None, ) -> str: """ Add a scheduled document status change task Args: document_id: Document ID dataset_id: Dataset ID tenant_id: Tenant ID action: Action to perform (enable/disable/archive) scheduled_at: When to execute the action created_by: User ID who created the schedule metadata: Additional metadata (e.g., reason, notes) Returns: schedule_id: Unique identifier for this scheduled task """ # Generate unique schedule ID schedule_id = f"{document_id}:{action.value}:{int(scheduled_at.timestamp())}" # Prepare schedule details schedule_detail = { "schedule_id": schedule_id, "document_id": document_id, "dataset_id": dataset_id, "tenant_id": tenant_id, "action": action.value, "scheduled_at": scheduled_at.isoformat(), "created_by": created_by, "created_at": datetime.utcnow().isoformat(), "status": "pending", # pending, executing, completed, cancelled, failed "metadata": metadata or {}, } # Store in Redis # 1. Add to Sorted Set (for efficient time-based queries) schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}all" redis_client.zadd(schedule_key, {schedule_id: scheduled_at.timestamp()}) # 2. Store detailed information in Hash detail_key = f"{cls.SCHEDULE_DETAIL_PREFIX}{schedule_id}" redis_client.hset(detail_key, mapping={k: json.dumps(v) for k, v in schedule_detail.items()}) # 3. Set expiration (keep for 30 days after execution) expire_seconds = int((scheduled_at.timestamp() - datetime.utcnow().timestamp()) + 30 * 24 * 3600) if expire_seconds > 0: redis_client.expire(detail_key, expire_seconds) # 4. Add to document-specific index for quick lookup doc_schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}doc:{document_id}" redis_client.zadd(doc_schedule_key, {schedule_id: scheduled_at.timestamp()}) redis_client.expire(doc_schedule_key, expire_seconds if expire_seconds > 0 else 86400) logger.info( f"Added schedule: {schedule_id} for document {document_id}, " f"action={action.value}, scheduled_at={scheduled_at}" ) return schedule_id @classmethod def cancel_schedule(cls, schedule_id: str, cancelled_by: str) -> bool: """ Cancel a scheduled task Args: schedule_id: Schedule ID to cancel cancelled_by: User ID who cancelled the schedule Returns: bool: True if cancelled successfully, False if not found """ detail_key = f"{cls.SCHEDULE_DETAIL_PREFIX}{schedule_id}" # Check if schedule exists if not redis_client.exists(detail_key): return False # Update status to cancelled redis_client.hset( detail_key, mapping={ "status": json.dumps("cancelled"), "cancelled_by": json.dumps(cancelled_by), "cancelled_at": json.dumps(datetime.utcnow().isoformat()), }, ) # Remove from sorted set schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}all" redis_client.zrem(schedule_key, schedule_id) # Extract document_id from schedule_id document_id = schedule_id.split(":")[0] doc_schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}doc:{document_id}" redis_client.zrem(doc_schedule_key, schedule_id) logger.info(f"Cancelled schedule: {schedule_id} by user {cancelled_by}") return True @classmethod def get_schedule(cls, schedule_id: str) -> Optional[dict]: """ Get schedule details Args: schedule_id: Schedule ID Returns: dict: Schedule details or None if not found """ detail_key = f"{cls.SCHEDULE_DETAIL_PREFIX}{schedule_id}" if not redis_client.exists(detail_key): return None # Get all fields from hash raw_data = redis_client.hgetall(detail_key) if not raw_data: return None # Decode and parse JSON schedule_detail = {} for key, value in raw_data.items(): key_str = key.decode("utf-8") if isinstance(key, bytes) else key value_str = value.decode("utf-8") if isinstance(value, bytes) else value try: schedule_detail[key_str] = json.loads(value_str) except json.JSONDecodeError: schedule_detail[key_str] = value_str return schedule_detail @classmethod def get_document_schedules(cls, document_id: str, include_completed: bool = False) -> list[dict]: """ Get all schedules for a specific document Args: document_id: Document ID include_completed: Whether to include completed/cancelled schedules Returns: list: List of schedule details """ doc_schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}doc:{document_id}" # Get all schedule IDs for this document schedule_ids = redis_client.zrange(doc_schedule_key, 0, -1) if not schedule_ids: return [] schedules = [] for schedule_id_bytes in schedule_ids: schedule_id = schedule_id_bytes.decode("utf-8") if isinstance(schedule_id_bytes, bytes) else schedule_id_bytes schedule = cls.get_schedule(schedule_id) if schedule: # Filter based on status if include_completed or schedule.get("status") == "pending": schedules.append(schedule) # Sort by scheduled_at schedules.sort(key=lambda x: x.get("scheduled_at", "")) return schedules @classmethod def get_pending_schedules(cls, before_timestamp: Optional[float] = None) -> list[dict]: """ Get all pending schedules that should be executed Args: before_timestamp: Only get schedules before this timestamp (default: now) Returns: list: List of schedule details ready for execution """ if before_timestamp is None: before_timestamp = datetime.utcnow().timestamp() schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}all" # Get schedule IDs with score (timestamp) <= before_timestamp schedule_ids = redis_client.zrangebyscore(schedule_key, 0, before_timestamp) if not schedule_ids: return [] pending_schedules = [] for schedule_id_bytes in schedule_ids: schedule_id = schedule_id_bytes.decode("utf-8") if isinstance(schedule_id_bytes, bytes) else schedule_id_bytes schedule = cls.get_schedule(schedule_id) if schedule and schedule.get("status") == "pending": pending_schedules.append(schedule) return pending_schedules @classmethod def update_schedule_status( cls, schedule_id: str, status: str, error_message: Optional[str] = None ) -> bool: """ Update schedule execution status Args: schedule_id: Schedule ID status: New status (executing, completed, failed) error_message: Error message if failed Returns: bool: True if updated successfully """ detail_key = f"{cls.SCHEDULE_DETAIL_PREFIX}{schedule_id}" if not redis_client.exists(detail_key): return False update_data = { "status": json.dumps(status), "updated_at": json.dumps(datetime.utcnow().isoformat()), } if error_message: update_data["error_message"] = json.dumps(error_message) if status == "completed": update_data["completed_at"] = json.dumps(datetime.utcnow().isoformat()) redis_client.hset(detail_key, mapping=update_data) # Remove from sorted set if completed/failed if status in ["completed", "failed"]: schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}all" redis_client.zrem(schedule_key, schedule_id) # Extract document_id from schedule_id document_id = schedule_id.split(":")[0] doc_schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}doc:{document_id}" redis_client.zrem(doc_schedule_key, schedule_id) return True @classmethod def clear_expired_schedules(cls, days: int = 30) -> int: """ Clear old completed/cancelled schedules Args: days: Clear schedules older than this many days Returns: int: Number of schedules cleared """ # This is handled by Redis TTL, but can be called manually # to clean up if needed cutoff_timestamp = datetime.utcnow().timestamp() - (days * 24 * 3600) schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}all" removed_count = redis_client.zremrangebyscore(schedule_key, 0, cutoff_timestamp) logger.info(f"Cleared {removed_count} expired schedules older than {days} days") return removed_count @classmethod def get_statistics(cls, tenant_id: Optional[str] = None) -> dict[str, Any]: """ Get schedule statistics Args: tenant_id: Optional tenant ID to filter by Returns: dict: Statistics including counts by status """ schedule_key = f"{cls.SCHEDULE_KEY_PREFIX}all" all_schedule_ids = redis_client.zrange(schedule_key, 0, -1) stats = { "total": 0, "pending": 0, "executing": 0, "completed": 0, "cancelled": 0, "failed": 0, } for schedule_id_bytes in all_schedule_ids: schedule_id = schedule_id_bytes.decode("utf-8") if isinstance(schedule_id_bytes, bytes) else schedule_id_bytes schedule = cls.get_schedule(schedule_id) if schedule: # Filter by tenant if specified if tenant_id and schedule.get("tenant_id") != tenant_id: continue stats["total"] += 1 status = schedule.get("status", "unknown") if status in stats: stats[status] += 1 return stats