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