gh-christianlouis-docuelevate/app/tasks/batch_tasks.py

"""
Scheduled batch processing tasks for DocuElevate.

This module provides three Celery tasks that can be scheduled via Celery Beat
and managed through the admin UI (``/admin/scheduled-jobs``):

- ``process_new_documents``    – Queue any documents that have never been processed.
- ``reprocess_failed_documents`` – Re-queue documents whose processing failed.
- ``cleanup_temp_files``       – Remove stale files from the ``workdir/tmp`` directory.

Each task records its execution result back to the ``ScheduledJob`` table so
the admin UI can display last-run times and statuses.
"""

import logging
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path

from app.celery_app import celery
from app.config import settings
from app.database import SessionLocal
from app.models import FileProcessingStep, FileRecord, ScheduledJob

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

_PROCESSING_STEPS = {
    "create_file_record",
    "check_text",
    "extract_text",
    "process_with_ocr",
    "process_with_azure_document_intelligence",
    "extract_metadata_with_gpt",
    "embed_metadata_into_pdf",
    "finalize_document_storage",
    "send_to_all_destinations",
}


def _update_job_status(job_name: str, status: str, detail: str) -> None:
    """Persist run status back to the ScheduledJob row for display in the UI."""
    try:
        with SessionLocal() as db:
            job = db.query(ScheduledJob).filter(ScheduledJob.name == job_name).first()
            if job:
                job.last_run_at = datetime.now(timezone.utc)
                job.last_run_status = status
                job.last_run_detail = detail
                db.commit()
    except Exception as exc:  # pragma: no cover – best-effort status update
        logger.warning("Could not update ScheduledJob status for %s: %s", job_name, exc)


# ---------------------------------------------------------------------------
# Task: process new (unprocessed) documents
# ---------------------------------------------------------------------------


@celery.task(name="app.tasks.batch_tasks.process_new_documents")
def process_new_documents() -> dict:
    """
    Queue all documents that have never been processed.

    A document is considered *new* when it has no ``FileProcessingStep`` rows
    that match the core pipeline steps.  The task loads each qualifying
    ``FileRecord``, verifies that the original file still exists on disk, and
    dispatches ``process_document`` for each one.

    Returns a summary dict with ``queued`` and ``skipped`` counts.
    """
    from app.tasks.process_document import process_document  # avoid circular import

    job_name = "process-new-documents"
    logger.info("[batch] Starting process_new_documents task")

    try:
        with SessionLocal() as db:
            # Files that already have at least one processing step recorded.
            processed_file_ids = (
                db.query(FileProcessingStep.file_id)
                .filter(FileProcessingStep.step_name.in_(_PROCESSING_STEPS))
                .distinct()
                .subquery()
            )

            # Candidate files: non-duplicate records with no processing steps yet.
            candidates = (
                db.query(FileRecord)
                .filter(FileRecord.is_duplicate.is_(False))
                .filter(~FileRecord.id.in_(db.query(processed_file_ids.c.file_id)))
                .all()
            )

        queued = 0
        skipped = 0
        for record in candidates:
            if not record.local_filename or not os.path.exists(record.local_filename):
                logger.warning(
                    "[batch] Skipping file_id=%s — local file not found: %s",
                    record.id,
                    record.local_filename,
                )
                skipped += 1
                continue
            process_document.delay(
                record.local_filename,
                original_filename=record.original_filename,
                file_id=record.id,
                owner_id=record.owner_id,
            )
            queued += 1

        detail = f"Queued {queued} document(s) for processing; skipped {skipped} (file not on disk)."
        logger.info("[batch] process_new_documents: %s", detail)
        _update_job_status(job_name, "success", detail)
        return {"queued": queued, "skipped": skipped}

    except Exception as exc:
        detail = f"Error: {exc}"
        logger.error("[batch] process_new_documents failed: %s", exc, exc_info=True)
        _update_job_status(job_name, "failed", detail)
        return {"queued": 0, "skipped": 0, "error": str(exc)}


# ---------------------------------------------------------------------------
# Task: reprocess failed documents
# ---------------------------------------------------------------------------


@celery.task(name="app.tasks.batch_tasks.reprocess_failed_documents")
def reprocess_failed_documents() -> dict:
    """
    Re-queue documents whose most-recent processing attempt failed.

    Only files that have at least one ``FileProcessingStep`` with
    ``status == "failure"`` **and** no currently ``in_progress`` steps are
    selected so that actively-running jobs are not interrupted.

    Returns a summary dict with ``queued`` and ``skipped`` counts.
    """
    from app.tasks.process_document import process_document  # avoid circular import

    job_name = "reprocess-failed-documents"
    logger.info("[batch] Starting reprocess_failed_documents task")

    try:
        with SessionLocal() as db:
            # Files with at least one failed step.
            failed_file_ids = (
                db.query(FileProcessingStep.file_id)
                .filter(FileProcessingStep.step_name.in_(_PROCESSING_STEPS))
                .filter(FileProcessingStep.status == "failure")
                .distinct()
                .subquery()
            )

            # Exclude files that are currently being processed.
            in_progress_file_ids = (
                db.query(FileProcessingStep.file_id)
                .filter(FileProcessingStep.status == "in_progress")
                .distinct()
                .subquery()
            )

            candidates = (
                db.query(FileRecord)
                .filter(FileRecord.is_duplicate.is_(False))
                .filter(FileRecord.id.in_(db.query(failed_file_ids.c.file_id)))
                .filter(~FileRecord.id.in_(db.query(in_progress_file_ids.c.file_id)))
                .all()
            )

        queued = 0
        skipped = 0
        for record in candidates:
            if not record.local_filename or not os.path.exists(record.local_filename):
                logger.warning(
                    "[batch] Skipping file_id=%s — local file not found: %s",
                    record.id,
                    record.local_filename,
                )
                skipped += 1
                continue
            process_document.delay(
                record.local_filename,
                original_filename=record.original_filename,
                file_id=record.id,
                owner_id=record.owner_id,
            )
            queued += 1

        detail = f"Re-queued {queued} failed document(s); skipped {skipped} (file not on disk)."
        logger.info("[batch] reprocess_failed_documents: %s", detail)
        _update_job_status(job_name, "success", detail)
        return {"queued": queued, "skipped": skipped}

    except Exception as exc:
        detail = f"Error: {exc}"
        logger.error("[batch] reprocess_failed_documents failed: %s", exc, exc_info=True)
        _update_job_status(job_name, "failed", detail)
        return {"queued": 0, "skipped": 0, "error": str(exc)}


# ---------------------------------------------------------------------------
# Task: clean up temporary files
# ---------------------------------------------------------------------------

#: Files in ``workdir/tmp`` that are older than this threshold are deleted.
_TEMP_FILE_MAX_AGE_HOURS: int = 24


@celery.task(name="app.tasks.batch_tasks.cleanup_temp_files")
def cleanup_temp_files(max_age_hours: int = _TEMP_FILE_MAX_AGE_HOURS) -> dict:
    """
    Delete stale files from the ``workdir/tmp`` directory.

    A file is considered stale when **both** of the following are true:

    1. Its modification time is older than *max_age_hours* (default 24 h).
    2. No ``FileRecord.local_filename`` points to it  **or**  the file is not
       referenced by any active in-progress processing step.

    This prevents accidental deletion of files that are being actively
    processed by the pipeline.

    Args:
        max_age_hours: Minimum age (in hours) before a temp file is eligible
                       for deletion.  Defaults to 24.

    Returns:
        A summary dict with ``deleted`` and ``skipped`` counts.
    """
    job_name = "cleanup-temp-files"
    logger.info("[batch] Starting cleanup_temp_files (max_age_hours=%s)", max_age_hours)

    tmp_dir = Path(settings.workdir) / "tmp"
    if not tmp_dir.exists():
        detail = "workdir/tmp does not exist; nothing to clean."
        logger.info("[batch] cleanup_temp_files: %s", detail)
        _update_job_status(job_name, "success", detail)
        return {"deleted": 0, "skipped": 0}

    cutoff = datetime.now(timezone.utc) - timedelta(hours=max_age_hours)

    deleted = 0
    skipped = 0
    errors = 0

    try:
        with SessionLocal() as db:
            # Collect filenames actively referenced by in-progress processing steps.
            in_progress_filenames: set[str] = set()
            in_progress_records = (
                db.query(FileRecord.local_filename)
                .join(FileProcessingStep, FileProcessingStep.file_id == FileRecord.id)
                .filter(FileProcessingStep.status == "in_progress")
                .distinct()
                .all()
            )
            for row in in_progress_records:
                if row.local_filename:
                    in_progress_filenames.add(os.path.basename(row.local_filename))

            # Also collect all filenames referenced by FileRecord.local_filename
            # that point into workdir/tmp (files still in the tmp pipeline).
            active_tmp_filenames: set[str] = set()
            tmp_dir_str = str(tmp_dir.resolve())
            active_records = (
                db.query(FileRecord.local_filename)
                .filter(FileRecord.local_filename.like(f"{tmp_dir_str}%"))
                .all()
            )
            for row in active_records:
                if row.local_filename:
                    active_tmp_filenames.add(os.path.basename(row.local_filename))

        protected_basenames = in_progress_filenames | active_tmp_filenames

        for entry in tmp_dir.iterdir():
            if not entry.is_file():
                continue

            # Check modification time.
            try:
                mtime = datetime.fromtimestamp(entry.stat().st_mtime, tz=timezone.utc)
            except OSError:
                skipped += 1
                continue

            if mtime >= cutoff:
                skipped += 1
                continue

            if entry.name in protected_basenames:
                logger.debug("[batch] cleanup_temp_files: keeping protected file %s", entry.name)
                skipped += 1
                continue

            try:
                entry.unlink()
                logger.debug("[batch] cleanup_temp_files: deleted %s", entry)
                deleted += 1
            except OSError as exc:
                logger.warning("[batch] cleanup_temp_files: could not delete %s: %s", entry, exc)
                errors += 1

        detail = (
            f"Deleted {deleted} stale temp file(s); "
            f"skipped {skipped} (too new or protected); "
            f"{errors} error(s)."
        )
        status = "failed" if errors and not deleted else "success"
        logger.info("[batch] cleanup_temp_files: %s", detail)
        _update_job_status(job_name, status, detail)
        return {"deleted": deleted, "skipped": skipped, "errors": errors}

    except Exception as exc:
        detail = f"Error: {exc}"
        logger.error("[batch] cleanup_temp_files failed: %s", exc, exc_info=True)
        _update_job_status(job_name, "failed", detail)
        return {"deleted": 0, "skipped": 0, "errors": 1, "error": str(exc)}