scholarr/app/services/runs/queue_service.py

from __future__ import annotations

from sqlalchemy import case, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import IngestionQueueItem
from app.services.ingestion import queue as queue_mutations
from app.services.runs.queue_queries import queue_item_select, queue_list_item_from_row
from app.services.runs.types import (
    QUEUE_STATUS_DROPPED,
    QUEUE_STATUS_QUEUED,
    QUEUE_STATUS_RETRYING,
    QueueClearResult,
    QueueListItem,
    QueueTransitionError,
)


async def list_queue_items_for_user(
    db_session: AsyncSession,
    *,
    user_id: int,
    limit: int = 200,
) -> list[QueueListItem]:
    result = await db_session.execute(
        queue_item_select(user_id=user_id)
        .order_by(
            case((IngestionQueueItem.status == QUEUE_STATUS_DROPPED, 1), else_=0).asc(),
            IngestionQueueItem.next_attempt_dt.asc(),
            IngestionQueueItem.id.asc(),
        )
        .limit(limit)
    )
    return [queue_list_item_from_row(row) for row in result.all()]


async def get_queue_item_for_user(
    db_session: AsyncSession,
    *,
    user_id: int,
    queue_item_id: int,
) -> QueueListItem | None:
    result = await db_session.execute(
        queue_item_select(user_id=user_id).where(IngestionQueueItem.id == queue_item_id).limit(1)
    )
    row = result.one_or_none()
    if row is None:
        return None
    return queue_list_item_from_row(row)


async def retry_queue_item_for_user(
    db_session: AsyncSession,
    *,
    user_id: int,
    queue_item_id: int,
) -> QueueListItem | None:
    result = await db_session.execute(
        select(IngestionQueueItem).where(
            IngestionQueueItem.id == queue_item_id,
            IngestionQueueItem.user_id == user_id,
        )
    )
    item = result.scalar_one_or_none()
    if item is None:
        return None
    if item.status == QUEUE_STATUS_QUEUED:
        raise QueueTransitionError(
            code="queue_item_already_queued",
            message="Queue item is already queued.",
            current_status=item.status,
        )
    if item.status == QUEUE_STATUS_RETRYING:
        raise QueueTransitionError(
            code="queue_item_retrying",
            message="Queue item is currently retrying.",
            current_status=item.status,
        )

    await queue_mutations.mark_queued_now(
        db_session,
        job_id=item.id,
        reason="manual_retry",
        reset_attempt_count=(item.status == QUEUE_STATUS_DROPPED),
    )
    await db_session.commit()
    return await get_queue_item_for_user(
        db_session,
        user_id=user_id,
        queue_item_id=queue_item_id,
    )


async def drop_queue_item_for_user(
    db_session: AsyncSession,
    *,
    user_id: int,
    queue_item_id: int,
) -> QueueListItem | None:
    result = await db_session.execute(
        select(IngestionQueueItem).where(
            IngestionQueueItem.id == queue_item_id,
            IngestionQueueItem.user_id == user_id,
        )
    )
    item = result.scalar_one_or_none()
    if item is None:
        return None
    if item.status == QUEUE_STATUS_DROPPED:
        raise QueueTransitionError(
            code="queue_item_already_dropped",
            message="Queue item is already dropped.",
            current_status=item.status,
        )

    await queue_mutations.mark_dropped(
        db_session,
        job_id=int(item.id),
        reason="manual_drop",
    )
    await db_session.commit()
    return await get_queue_item_for_user(
        db_session,
        user_id=user_id,
        queue_item_id=queue_item_id,
    )


async def clear_queue_item_for_user(
    db_session: AsyncSession,
    *,
    user_id: int,
    queue_item_id: int,
) -> QueueClearResult | None:
    result = await db_session.execute(
        select(IngestionQueueItem).where(
            IngestionQueueItem.id == queue_item_id,
            IngestionQueueItem.user_id == user_id,
        )
    )
    item = result.scalar_one_or_none()
    if item is None:
        return None
    if item.status != QUEUE_STATUS_DROPPED:
        raise QueueTransitionError(
            code="queue_item_not_dropped",
            message="Queue item can only be cleared after it is dropped.",
            current_status=item.status,
        )

    item_id = int(item.id)
    previous_status = str(item.status)
    deleted = await queue_mutations.delete_job_by_id(db_session, job_id=item_id)
    await db_session.commit()
    if not deleted:
        return None
    return QueueClearResult(queue_item_id=item_id, previous_status=previous_status)


async def queue_status_counts_for_user(
    db_session: AsyncSession,
    *,
    user_id: int,
) -> dict[str, int]:
    result = await db_session.execute(
        select(
            IngestionQueueItem.status,
            func.count(IngestionQueueItem.id),
        )
        .where(IngestionQueueItem.user_id == user_id)
        .group_by(IngestionQueueItem.status)
    )
    counts: dict[str, int] = {
        QUEUE_STATUS_QUEUED: 0,
        QUEUE_STATUS_RETRYING: 0,
        QUEUE_STATUS_DROPPED: 0,
    }
    for status, count in result.all():
        counts[str(status)] = int(count or 0)
    return counts