dify-1.8.1/api/controllers/console/datasets/document_schedule.py

"""
Document Schedule API
Manage scheduled document status changes
"""

from datetime import datetime, timezone

from flask import request
from flask_login import current_user
from flask_restx import Resource, fields, marshal_with, reqparse
from werkzeug.exceptions import BadRequest, Forbidden, NotFound

from controllers.console import api
from controllers.console.datasets.error import InvalidActionError
from controllers.console.wraps import (
    account_initialization_required,
    cloud_edition_billing_rate_limit_check,
    setup_required,
)
from libs.login import login_required
from services.dataset_operation_log_service import DatasetOperationLogService, DatasetOperationType
from services.dataset_service import DatasetService, DocumentService
from services.document_schedule_service import DocumentScheduleService, ScheduleAction

# Response fields for schedule
schedule_fields = {
    "schedule_id": fields.String,
    "document_id": fields.String,
    "dataset_id": fields.String,
    "action": fields.String,
    "scheduled_at": fields.String,
    "created_by": fields.String,
    "created_at": fields.String,
    "status": fields.String,
    "metadata": fields.Raw,
    "cancelled_by": fields.String,
    "cancelled_at": fields.String,
    "completed_at": fields.String,
    "error_message": fields.String,
}

schedule_list_fields = {
    "schedules": fields.List(fields.Nested(schedule_fields)),
    "total": fields.Integer,
}


class DocumentScheduleResource(Resource):
    """Base resource for document schedule operations"""

    pass


class DocumentScheduleListApi(DocumentScheduleResource):
    """API for listing and creating document schedules"""

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(schedule_list_fields)
    def get(self, dataset_id: str, document_id: str):
        """
        Get all schedules for a document

        Args:
            dataset_id: Dataset ID
            document_id: Document ID

        Returns:
            dict: List of schedules
        """
        dataset_id_str = str(dataset_id)
        document_id_str = str(document_id)

        # Get dataset and check permissions
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        DatasetService.check_dataset_permission(dataset, current_user)

        # Get document
        document = DocumentService.get_document(dataset_id_str, document_id_str)
        if not document:
            raise NotFound("Document not found.")

        # Parse query parameters
        parser = reqparse.RequestParser()
        parser.add_argument("include_completed", type=bool, default=False, location="args")
        args = parser.parse_args()

        # Get schedules from Redis
        schedules = DocumentScheduleService.get_document_schedules(
            document_id_str, include_completed=args["include_completed"]
        )

        return {"schedules": schedules, "total": len(schedules)}

    @setup_required
    @login_required
    @account_initialization_required
    @cloud_edition_billing_rate_limit_check("knowledge")
    @marshal_with(schedule_fields)
    def post(self, dataset_id: str, document_id: str):
        """
        Create a scheduled document status change

        Args:
            dataset_id: Dataset ID
            document_id: Document ID

        Request Body:
            {
                "action": "enable|disable|archive",
                "scheduled_at": "2024-01-01T00:00:00Z",
                "metadata": {"reason": "...", "notes": "..."}
            }

        Returns:
            dict: Created schedule details
        """
        dataset_id_str = str(dataset_id)
        document_id_str = str(document_id)

        # Get dataset and check permissions
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        if not current_user.is_dataset_editor:
            raise Forbidden()

        DatasetService.check_dataset_permission(dataset, current_user)
        DatasetService.check_dataset_model_setting(dataset)

        # Get document
        document = DocumentService.get_document(dataset_id_str, document_id_str)
        if not document:
            raise NotFound("Document not found.")

        # Parse request body
        data = request.get_json()
        if not data:
            raise BadRequest("Request body is required")

        action_str = data.get("action")
        scheduled_at_str = data.get("scheduled_at")
        metadata = data.get("metadata", {})

        # Validate action
        try:
            action = ScheduleAction(action_str)
        except ValueError:
            raise BadRequest(f"Invalid action: {action_str}. Must be one of: enable, disable, archive")

        # Validate scheduled_at
        if not scheduled_at_str:
            raise BadRequest("scheduled_at is required")

        try:
            scheduled_at = datetime.fromisoformat(scheduled_at_str.replace("Z", "+00:00"))
        except (ValueError, AttributeError):
            raise BadRequest("Invalid scheduled_at format. Use ISO 8601 format (e.g., 2024-01-01T00:00:00Z)")

        # Check if scheduled time is in the future
        if scheduled_at <= datetime.now(timezone.utc):
            raise BadRequest("scheduled_at must be in the future")

        # Create schedule
        try:
            schedule_id = DocumentScheduleService.add_schedule(
                document_id=document_id_str,
                dataset_id=dataset_id_str,
                tenant_id=str(dataset.tenant_id),
                action=action,
                scheduled_at=scheduled_at,
                created_by=str(current_user.id),
                metadata=metadata,
            )

            # Log operation
            DatasetOperationLogService.log_operation(
                tenant_id=str(dataset.tenant_id),
                account_id=str(current_user.id),
                action=DatasetOperationType.DOCUMENT_SCHEDULE_CREATE,
                dataset_id=dataset_id_str,
                resource_type="document",
                resource_id=document_id_str,
                metadata={
                    "schedule_id": schedule_id,
                    "action": action.value,
                    "scheduled_at": scheduled_at_str,
                },
            )

            # Get created schedule
            schedule = DocumentScheduleService.get_schedule(schedule_id)
            return schedule

        except Exception as e:
            raise InvalidActionError(str(e))


class DocumentScheduleDetailApi(DocumentScheduleResource):
    """API for managing a specific schedule"""

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(schedule_fields)
    def get(self, dataset_id: str, document_id: str, schedule_id: str):
        """
        Get schedule details

        Args:
            dataset_id: Dataset ID
            document_id: Document ID
            schedule_id: Schedule ID

        Returns:
            dict: Schedule details
        """
        dataset_id_str = str(dataset_id)

        # Get dataset and check permissions
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        DatasetService.check_dataset_permission(dataset, current_user)

        # Get schedule
        schedule = DocumentScheduleService.get_schedule(schedule_id)
        if not schedule:
            raise NotFound("Schedule not found.")

        # Verify the schedule belongs to this document
        if schedule.get("document_id") != str(document_id) or schedule.get("dataset_id") != dataset_id_str:
            raise NotFound("Schedule not found.")

        return schedule

    @setup_required
    @login_required
    @account_initialization_required
    @cloud_edition_billing_rate_limit_check("knowledge")
    def delete(self, dataset_id: str, document_id: str, schedule_id: str):
        """
        Cancel a scheduled task

        Args:
            dataset_id: Dataset ID
            document_id: Document ID
            schedule_id: Schedule ID

        Returns:
            dict: Success message
        """
        dataset_id_str = str(dataset_id)
        document_id_str = str(document_id)

        # Get dataset and check permissions
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        if not current_user.is_dataset_editor:
            raise Forbidden()

        DatasetService.check_dataset_permission(dataset, current_user)

        # Get schedule to verify it exists and belongs to this document
        schedule = DocumentScheduleService.get_schedule(schedule_id)
        if not schedule:
            raise NotFound("Schedule not found.")

        if schedule.get("document_id") != document_id_str or schedule.get("dataset_id") != dataset_id_str:
            raise NotFound("Schedule not found.")

        # Cancel schedule
        success = DocumentScheduleService.cancel_schedule(schedule_id, str(current_user.id))

        if not success:
            raise InvalidActionError("Failed to cancel schedule")

        # Log operation
        DatasetOperationLogService.log_operation(
            tenant_id=str(dataset.tenant_id),
            account_id=str(current_user.id),
            action=DatasetOperationType.DOCUMENT_SCHEDULE_CANCEL,
            dataset_id=dataset_id_str,
            resource_type="document",
            resource_id=document_id_str,
            metadata={
                "schedule_id": schedule_id,
                "action": schedule.get("action"),
                "scheduled_at": schedule.get("scheduled_at"),
            },
        )

        return {"result": "success", "message": "Schedule cancelled successfully"}, 200


class DocumentScheduleStatisticsApi(DocumentScheduleResource):
    """API for getting schedule statistics"""

    @setup_required
    @login_required
    @account_initialization_required
    def get(self, dataset_id: str):
        """
        Get schedule statistics for a dataset

        Args:
            dataset_id: Dataset ID

        Returns:
            dict: Statistics
        """
        dataset_id_str = str(dataset_id)

        # Get dataset and check permissions
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        DatasetService.check_dataset_permission(dataset, current_user)

        # Get statistics
        stats = DocumentScheduleService.get_statistics(tenant_id=str(dataset.tenant_id))

        return stats, 200


# Register API routes
api.add_resource(
    DocumentScheduleListApi,
    "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/schedules",
)
api.add_resource(
    DocumentScheduleDetailApi,
    "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/schedules/<string:schedule_id>",
)
api.add_resource(
    DocumentScheduleStatisticsApi,
    "/datasets/<uuid:dataset_id>/schedules/statistics",
)