arm-core/api/src/repositories/files.py

"""Репозиторий первичных файлов, документов и строк ОСВ (spec documents-api; design D-11).

SQL без бизнес-правил; видимость — через ЮЛ-владельца (`src.domain.policy.legal_entities_visible`): файл или документ
вне видимости неотличим от несуществующего (`None`). Связей `relationship()` в моделях нет, поэтому задания и документы
страницы файлов читаются пакетно (`current_jobs`, `documents_by_file`) — число запросов не зависит от числа файлов.

Функции не делают `commit`: транзакцией запроса управляет `SessionDep`.
"""

import uuid
from collections.abc import Sequence
from typing import Any

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

from src.domain.policy import VisibilityScope, legal_entities_visible
from src.models.directories import LegalEntity
from src.models.enums import FileStatus
from src.models.files import Document, RecognitionJob, StoredFile
from src.models.osv import OsvEntry


def _visible_files(stmt: Select[Any], scope: VisibilityScope) -> Select[Any]:
    """Файлы, чьё ЮЛ-владелец в области видимости (`stmt` ссылается на `StoredFile`)."""
    return legal_entities_visible(stmt.join(LegalEntity, LegalEntity.id == StoredFile.legal_entity_id), scope)


def _visible_documents(stmt: Select[Any], scope: VisibilityScope) -> Select[Any]:
    """Документы, чьё ЮЛ-владелец в области видимости (`stmt` ссылается на `Document`)."""
    return legal_entities_visible(stmt.join(LegalEntity, LegalEntity.id == Document.legal_entity_id), scope)


async def list_files(  # noqa: PLR0913 - фильтр и пагинация списка передаются явными именованными аргументами
    session: AsyncSession,
    scope: VisibilityScope,
    entity_id: uuid.UUID,
    *,
    include_superseded: bool,
    limit: int,
    offset: int,
) -> tuple[list[StoredFile], int]:
    """Страница видимых файлов ЮЛ по убыванию `uploaded_at` (затем `id`) и общее число по фильтру; `SUPERSEDED` — только
    при `include_superseded`."""

    def filtered(stmt: Select[Any]) -> Select[Any]:
        stmt = _visible_files(stmt, scope).where(StoredFile.legal_entity_id == entity_id)
        if not include_superseded:
            stmt = stmt.where(StoredFile.status != FileStatus.SUPERSEDED)
        return stmt

    total = (await session.execute(filtered(select(func.count(StoredFile.id))))).scalar_one()
    page_stmt = (
        filtered(select(StoredFile))
        .order_by(StoredFile.uploaded_at.desc().nulls_last(), StoredFile.id)
        .limit(limit)
        .offset(offset)
    )
    files = list((await session.execute(page_stmt)).scalars().all())
    return files, total


async def get_file(session: AsyncSession, scope: VisibilityScope, file_id: uuid.UUID) -> StoredFile | None:
    """Файл по id; `None` и для несуществующего, и для файла ЮЛ вне видимости."""
    stmt = _visible_files(select(StoredFile).where(StoredFile.id == file_id), scope)
    return (await session.execute(stmt)).scalar_one_or_none()


async def current_jobs(session: AsyncSession, files: Sequence[StoredFile]) -> dict[uuid.UUID, RecognitionJob]:
    """Текущие задания файлов (`stored_file.current_recognition_job_id`) одним запросом: `file_id → задание`."""
    job_ids = [file.current_recognition_job_id for file in files if file.current_recognition_job_id is not None]
    if not job_ids:
        return {}
    jobs = (await session.execute(select(RecognitionJob).where(RecognitionJob.id.in_(job_ids)))).scalars().all()
    return {job.file_id: job for job in jobs}


async def documents_by_file(session: AsyncSession, file_ids: Sequence[uuid.UUID]) -> dict[uuid.UUID, list[Document]]:
    """Документы файлов одним запросом: `file_id → документы` (по типу, счёту, `locator`)."""
    if not file_ids:
        return {}
    stmt = (
        select(Document)
        .where(Document.stored_file_id.in_(file_ids))
        .order_by(Document.document_type, Document.osv_account_code, Document.locator, Document.id)
    )
    grouped: dict[uuid.UUID, list[Document]] = {file_id: [] for file_id in file_ids}
    for document in (await session.execute(stmt)).scalars():
        grouped[document.stored_file_id].append(document)
    return grouped


async def get_document(session: AsyncSession, scope: VisibilityScope, document_id: uuid.UUID) -> Document | None:
    """Документ по id; `None` и для несуществующего, и для документа ЮЛ вне видимости."""
    stmt = _visible_documents(select(Document).where(Document.id == document_id), scope)
    return (await session.execute(stmt)).scalar_one_or_none()


async def list_entries(
    session: AsyncSession, document_id: uuid.UUID, *, limit: int, offset: int
) -> tuple[list[OsvEntry], int]:
    """Страница строк документа по возрастанию `row_number` и их общее число (видимость документа — у вызывающего)."""
    total = (
        await session.execute(select(func.count(OsvEntry.id)).where(OsvEntry.document_id == document_id))
    ).scalar_one()
    stmt = (
        select(OsvEntry)
        .where(OsvEntry.document_id == document_id)
        .order_by(OsvEntry.row_number, OsvEntry.id)
        .limit(limit)
        .offset(offset)
    )
    return list((await session.execute(stmt)).scalars().all()), total


async def lock_file(session: AsyncSession, file_id: uuid.UUID) -> StoredFile | None:
    """`SELECT stored_file … FOR UPDATE` с перечитыванием (объект в identity map мог устареть)."""
    stmt = (
        select(StoredFile).where(StoredFile.id == file_id).with_for_update().execution_options(populate_existing=True)
    )
    return (await session.execute(stmt)).scalar_one_or_none()


async def lock_document(session: AsyncSession, document_id: uuid.UUID) -> Document | None:
    """`SELECT document … FOR UPDATE` с перечитыванием; вызывать после `lock_file` его файла (порядок блокировок)."""
    stmt = (
        select(Document).where(Document.id == document_id).with_for_update().execution_options(populate_existing=True)
    )
    return (await session.execute(stmt)).scalar_one_or_none()