dify-1.8.1/api/services/document_schedule_service.py

"""
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