ArchiveBox/archivebox/api/v1_workers.py

__package__ = 'archivebox.api'

from uuid import UUID
from typing import List, Any
from datetime import datetime

from ninja import Router, Schema


router = Router(tags=['Workers and Tasks'])


class QueueItemSchema(Schema):
    """Schema for a single item in a worker's queue."""
    TYPE: str
    id: UUID
    status: str
    retry_at: datetime | None
    created_at: datetime
    modified_at: datetime
    description: str

    @staticmethod
    def resolve_TYPE(obj) -> str:
        return f'{obj._meta.app_label}.{obj._meta.model_name}'

    @staticmethod
    def resolve_description(obj) -> str:
        return str(obj)


class WorkerSchema(Schema):
    """Schema for a Worker type."""
    name: str
    model: str
    max_tick_time: int
    max_concurrent_tasks: int
    poll_interval: float
    idle_timeout: int
    running_count: int
    running_workers: List[dict[str, Any]]
    queue_count: int
    queue: List[QueueItemSchema]

    @staticmethod
    def resolve_model(obj) -> str:
        Model = obj.get_model()
        return f'{Model._meta.app_label}.{Model._meta.model_name}'

    @staticmethod
    def resolve_max_tick_time(obj) -> int:
        return obj.MAX_TICK_TIME

    @staticmethod
    def resolve_max_concurrent_tasks(obj) -> int:
        return obj.MAX_CONCURRENT_TASKS

    @staticmethod
    def resolve_poll_interval(obj) -> float:
        return obj.POLL_INTERVAL

    @staticmethod
    def resolve_idle_timeout(obj) -> int:
        return obj.IDLE_TIMEOUT

    @staticmethod
    def resolve_running_count(obj) -> int:
        return len(obj.get_running_workers())

    @staticmethod
    def resolve_running_workers(obj) -> List[dict[str, Any]]:
        return obj.get_running_workers()

    @staticmethod
    def resolve_queue_count(obj) -> int:
        return obj.get_queue().count()

    @staticmethod
    def resolve_queue(obj) -> List[QueueItemSchema]:
        return list(obj.get_queue()[:50])  # Limit to 50 items


class OrchestratorSchema(Schema):
    """Schema for the Orchestrator."""
    is_running: bool
    poll_interval: float
    idle_timeout: int
    max_workers_per_type: int
    max_total_workers: int
    total_worker_count: int
    workers: List[WorkerSchema]


@router.get("/orchestrator", response=OrchestratorSchema, url_name="get_orchestrator")
def get_orchestrator(request):
    """Get the orchestrator status and all worker queues."""
    from archivebox.workers.orchestrator import Orchestrator
    from workers.worker import CrawlWorker, SnapshotWorker, ArchiveResultWorker

    orchestrator = Orchestrator()

    # Create temporary worker instances to query their queues
    workers = [
        CrawlWorker(worker_id=-1),
        SnapshotWorker(worker_id=-1),
        ArchiveResultWorker(worker_id=-1),
    ]

    return {
        'is_running': orchestrator.is_running(),
        'poll_interval': orchestrator.POLL_INTERVAL,
        'idle_timeout': orchestrator.IDLE_TIMEOUT,
        'max_workers_per_type': orchestrator.MAX_WORKERS_PER_TYPE,
        'max_total_workers': orchestrator.MAX_TOTAL_WORKERS,
        'total_worker_count': orchestrator.get_total_worker_count(),
        'workers': workers,
    }


@router.get("/workers", response=List[WorkerSchema], url_name="get_workers")
def get_workers(request):
    """List all worker types and their current status."""
    from workers.worker import CrawlWorker, SnapshotWorker, ArchiveResultWorker

    # Create temporary instances to query their queues
    return [
        CrawlWorker(worker_id=-1),
        SnapshotWorker(worker_id=-1),
        ArchiveResultWorker(worker_id=-1),
    ]


@router.get("/worker/{worker_name}", response=WorkerSchema, url_name="get_worker")
def get_worker(request, worker_name: str):
    """Get status and queue for a specific worker type."""
    from workers.worker import WORKER_TYPES

    if worker_name not in WORKER_TYPES:
        from ninja.errors import HttpError
        raise HttpError(404, f"Unknown worker type: {worker_name}. Valid types: {list(WORKER_TYPES.keys())}")

    WorkerClass = WORKER_TYPES[worker_name]
    return WorkerClass(worker_id=-1)


@router.get("/worker/{worker_name}/queue", response=List[QueueItemSchema], url_name="get_worker_queue")
def get_worker_queue(request, worker_name: str, limit: int = 100):
    """Get the current queue for a specific worker type."""
    from workers.worker import WORKER_TYPES

    if worker_name not in WORKER_TYPES:
        from ninja.errors import HttpError
        raise HttpError(404, f"Unknown worker type: {worker_name}. Valid types: {list(WORKER_TYPES.keys())}")

    WorkerClass = WORKER_TYPES[worker_name]
    worker = WorkerClass(worker_id=-1)
    return list(worker.get_queue()[:limit])


# Progress endpoint moved to core.views.live_progress_view for simplicity