""" 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//documents//schedules", ) api.add_resource( DocumentScheduleDetailApi, "/datasets//documents//schedules/", ) api.add_resource( DocumentScheduleStatisticsApi, "/datasets//schedules/statistics", )