dify-1.8.1/api/schedule/document_schedule_check_task.py

"""
Document Schedule Check Task
Celery Beat task to check and execute scheduled document status changes
"""

import logging
import time

import click

import app
from extensions.ext_database import db
from models.dataset import Dataset, Document
from services.dataset_operation_log_service import DatasetOperationLogService, DatasetOperationType
from services.dataset_service import DocumentService
from services.document_schedule_service import DocumentScheduleService, ScheduleAction

logger = logging.getLogger(__name__)


@app.celery.task(queue="dataset")
def check_and_execute_document_schedules():
    """
    Check for pending scheduled document status changes and execute them

    This task should be run periodically (e.g., every 1-5 minutes) via Celery Beat
    """
    logger.info(click.style("Start checking scheduled document tasks", fg="green"))
    start_at = time.perf_counter()

    try:
        # Get all pending schedules that should be executed now
        pending_schedules = DocumentScheduleService.get_pending_schedules()

        if not pending_schedules:
            logger.info(click.style("No pending schedules found", fg="yellow"))
            return

        logger.info(click.style(f"Found {len(pending_schedules)} pending schedules", fg="green"))

        executed_count = 0
        failed_count = 0

        for schedule in pending_schedules:
            schedule_id = schedule.get("schedule_id")
            document_id = schedule.get("document_id")
            dataset_id = schedule.get("dataset_id")
            action = schedule.get("action")

            try:
                # Update status to executing
                DocumentScheduleService.update_schedule_status(schedule_id, "executing")

                # Execute the scheduled action
                success = _execute_schedule_action(schedule)

                if success:
                    # Update status to completed
                    DocumentScheduleService.update_schedule_status(schedule_id, "completed")
                    executed_count += 1
                    logger.info(
                        click.style(
                            f"Successfully executed schedule {schedule_id}: "
                            f"document={document_id}, action={action}",
                            fg="green",
                        )
                    )
                else:
                    # Update status to failed
                    DocumentScheduleService.update_schedule_status(
                        schedule_id, "failed", error_message="Execution failed"
                    )
                    failed_count += 1
                    logger.error(
                        click.style(
                            f"Failed to execute schedule {schedule_id}: "
                            f"document={document_id}, action={action}",
                            fg="red",
                        )
                    )

            except Exception as e:
                # Update status to failed with error message
                DocumentScheduleService.update_schedule_status(schedule_id, "failed", error_message=str(e))
                failed_count += 1
                logger.exception(
                    f"Error executing schedule {schedule_id}: "
                    f"document={document_id}, action={action}, error={str(e)}"
                )

        end_at = time.perf_counter()
        logger.info(
            click.style(
                f"Finished checking schedules: executed={executed_count}, "
                f"failed={failed_count}, latency={end_at - start_at:.2f}s",
                fg="green",
            )
        )

    except Exception as e:
        logger.exception(f"Error in check_and_execute_document_schedules: {str(e)}")
    finally:
        db.session.close()


def _execute_schedule_action(schedule: dict) -> bool:
    """
    Execute a scheduled action

    Args:
        schedule: Schedule details dictionary

    Returns:
        bool: True if successful, False otherwise
    """
    document_id = schedule.get("document_id")
    dataset_id = schedule.get("dataset_id")
    action = schedule.get("action")
    created_by = schedule.get("created_by")
    tenant_id = schedule.get("tenant_id")

    try:
        # Get dataset
        dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first()
        if not dataset:
            logger.error(f"Dataset not found: {dataset_id}")
            return False

        # Get document
        document = db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()
        if not document:
            logger.error(f"Document not found: {document_id}")
            return False

        # Create a mock user object for the action
        class ScheduledUser:
            def __init__(self, user_id: str):
                self.id = user_id

        user = ScheduledUser(created_by)

        # Execute the action
        if action == ScheduleAction.ENABLE.value:
            DocumentService.batch_update_document_status(dataset, [document_id], "enable", user)
            operation_type = DatasetOperationType.DOCUMENT_ENABLE
        elif action == ScheduleAction.DISABLE.value:
            DocumentService.batch_update_document_status(dataset, [document_id], "disable", user)
            operation_type = DatasetOperationType.DOCUMENT_DISABLE
        elif action == ScheduleAction.ARCHIVE.value:
            DocumentService.batch_update_document_status(dataset, [document_id], "archive", user)
            operation_type = DatasetOperationType.DOCUMENT_ARCHIVE
        else:
            logger.error(f"Unknown action: {action}")
            return False

        # Log the operation
        try:
            DatasetOperationLogService.log_operation(
                tenant_id=tenant_id,
                account_id=created_by,
                action=DatasetOperationType.DOCUMENT_SCHEDULE_EXECUTE,
                dataset_id=dataset_id,
                resource_type="document",
                resource_id=document_id,
                changes={
                    "action": action,
                    "operation_type": operation_type.value,
                },
                metadata={
                    "scheduled": True,
                    "schedule_id": schedule.get("schedule_id"),
                    "scheduled_at": schedule.get("scheduled_at"),
                },
            )
        except Exception as log_error:
            # Don't fail the entire operation if logging fails
            logger.error(f"Failed to log operation: {str(log_error)}")

        return True

    except Exception as e:
        logger.exception(f"Error executing schedule action: {str(e)}")
        db.session.rollback()
        return False