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 datetime
import uuid
from collections.abc import Sequence
from typing import Any

from sqlalchemy import Select, and_, delete, exists, func, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from src.domain.policy import VisibilityScope, legal_entities_visible
from src.models.access import AppUser
from src.models.audit_ex import Comment
from src.models.consolidation import ConsolidationDocument, ConsolidationVersion
from src.models.directories import LegalEntity
from src.models.enums import DocumentStatus, FileStatus, S3BucketRole
from src.models.files import Document, RecognitionJob, S3Object, StoredFile
from src.models.osv import OsvEntry
from src.repositories import statements as statements_repo

COMMENT_OBJECT_TYPE = "document"  # `comment.object_type` комментариев документа


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 comments_count_by_document(session: AsyncSession, document_ids: Sequence[uuid.UUID]) -> dict[uuid.UUID, int]:
    """Число комментариев документов одним запросом: `document_id → count` (документы без комментариев отсутствуют)."""
    if not document_ids:
        return {}
    stmt = (
        select(Comment.object_id, func.count(Comment.id))
        .where(Comment.object_type == COMMENT_OBJECT_TYPE, Comment.object_id.in_(document_ids))
        .group_by(Comment.object_id)
    )
    return {document_id: count for document_id, count in (await session.execute(stmt)).all()}


async def list_comments(session: AsyncSession, document_id: uuid.UUID) -> list[tuple[Comment, str]]:
    """Комментарии документа по `created_at` (затем `id`) вместе с ФИО автора."""
    stmt = (
        select(Comment, AppUser.full_name)
        .join(AppUser, AppUser.id == Comment.author_id)
        .where(Comment.object_type == COMMENT_OBJECT_TYPE, Comment.object_id == document_id)
        .order_by(Comment.created_at, Comment.id)
    )
    return [(comment, full_name) for comment, full_name in (await session.execute(stmt)).all()]


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 NO KEY UPDATE` с перечитыванием; вызывать после `lock_file` его файла.

    R-3 (design №13): `FOR NO KEY UPDATE`, а не `FOR UPDATE` — другие писатели документа ждут так же, но вставка
    `consolidation_document` (FK берёт `FOR KEY SHARE` документа) не ждёт: иначе подтверждение (держит документ, ждёт
    версию) и подхват (держит версию, вставляет привязку) взаимоблокировались бы."""
    stmt = (
        select(Document)
        .where(Document.id == document_id)
        .with_for_update(key_share=True)
        .execution_options(populate_existing=True)
    )
    return (await session.execute(stmt)).scalar_one_or_none()


async def lock_jobs(session: AsyncSession, file_id: uuid.UUID) -> list[RecognitionJob]:
    """`SELECT recognition_job … FOR UPDATE` всех заданий файла (по `id`) с перечитыванием. Порядок блокировок обмена и
    маппинга — «задание → файл», поэтому вызывать до `lock_file`."""
    stmt = (
        select(RecognitionJob)
        .where(RecognitionJob.file_id == file_id)
        .order_by(RecognitionJob.id)
        .with_for_update()
        .execution_options(populate_existing=True)
    )
    return list((await session.execute(stmt)).scalars().all())


async def lock_file_documents(session: AsyncSession, file_id: uuid.UUID) -> list[Document]:
    """`SELECT document … FOR NO KEY UPDATE` всех документов файла (по `id`) с перечитыванием; после `lock_file`.

    Режим — как у `lock_document` (R-3): совместим с `FOR KEY SHARE` вставки привязки; `DELETE` документа (удаление
    файла) сам повышает блокировку до `FOR UPDATE`."""
    stmt = (
        select(Document)
        .where(Document.stored_file_id == file_id)
        .order_by(Document.id)
        .with_for_update(key_share=True)
        .execution_options(populate_existing=True)
    )
    return list((await session.execute(stmt)).scalars().all())


async def confirmed_twin(session: AsyncSession, document: Document) -> uuid.UUID | None:
    """Подтверждённая форма того же файла, типа, года и квартала (другой `file_ordinal`) — для `DUPLICATE_FORM`."""
    if document.period_year is None or document.period_quarter is None:
        return None
    stmt = (
        select(Document.id)
        .where(
            Document.stored_file_id == document.stored_file_id,
            Document.document_type == document.document_type,
            Document.period_year == document.period_year,
            Document.period_quarter == document.period_quarter,
            Document.status == DocumentStatus.CONFIRMED,
            Document.id != document.id,
        )
        .order_by(Document.file_ordinal)
        .limit(1)
    )
    return (await session.execute(stmt)).scalar_one_or_none()


async def supersedes_others(session: AsyncSession, file_id: uuid.UUID) -> bool:
    """На файл ссылаются `superseded_by_file_id` других файлов (файл заменил другие версии)."""
    stmt = select(select(StoredFile.id).where(StoredFile.superseded_by_file_id == file_id).exists())
    return bool((await session.execute(stmt)).scalar_one())


async def delete_file_rows(session: AsyncSession, file: StoredFile) -> None:
    """Удаляет формы ББ/ОФР (шапки, колонки, строки, значения), комментарии документов, строки ОСВ, документы, все
    задания и сам файл (формы и комментарии — до документов и заданий: на них FK без CASCADE). FK-цикл
    `stored_file.current_recognition_job_id` ↔ `recognition_job.file_id` разрывается обнулением ссылки до удаления
    заданий. Блокировки — у вызывающего."""
    file_id = file.id
    file.current_recognition_job_id = None
    await session.flush()
    document_ids = select(Document.id).where(Document.stored_file_id == file_id)
    await statements_repo.delete_file_statements(session, file_id)
    await session.execute(
        delete(Comment).where(Comment.object_type == COMMENT_OBJECT_TYPE, Comment.object_id.in_(document_ids))
    )
    await session.execute(delete(OsvEntry).where(OsvEntry.document_id.in_(document_ids)))
    await session.execute(delete(Document).where(Document.stored_file_id == file_id))
    await session.execute(delete(RecognitionJob).where(RecognitionJob.file_id == file_id))
    await session.execute(delete(StoredFile).where(StoredFile.id == file_id))


async def expire_s3_objects(
    session: AsyncSession,
    object_ids: Sequence[uuid.UUID],
    *,
    owned_prefixes: Sequence[tuple[S3BucketRole, str]],
    now: datetime.datetime,
) -> None:
    """`s3_object.expires_at = now` (ретеншн бакета заберёт объект; физически объект не удаляется — OQ-22) — только
    объектам из `object_ids`, чья пара (роль бакета, начало ключа) входит в `owned_prefixes`. Остальные (например,
    ссылка задания на объект другого файла) не трогаются и ошибкой не считаются (M-3)."""
    if not object_ids or not owned_prefixes:
        return
    owned = or_(
        *(
            and_(S3Object.bucket_role == role, S3Object.key.startswith(prefix, autoescape=True))
            for role, prefix in owned_prefixes
        )
    )
    await session.execute(update(S3Object).where(S3Object.id.in_(object_ids), owned).values(expires_at=now))


async def list_desk_files(  # noqa: PLR0913 - фильтр и пагинация списка передаются явными именованными аргументами
    session: AsyncSession,
    scope: VisibilityScope,
    entity_id: uuid.UUID,
    consolidation_id: uuid.UUID,
    *,
    include_superseded: bool,
    limit: int,
    offset: int,
) -> tuple[list[StoredFile], int]:
    """Рабочий стол консолидации для ЮЛ (spec documents-api «Рабочий стол консолидации для ЮЛ»; §9.5 п. 11; D-6):
    файлы ЮЛ с `consolidation_id` этой консолидации ИЛИ с документом, привязанным к любой её версии. Порядок и
    `include_superseded` — как в `list_files`; страница и общее число тем же фильтром."""
    bound = exists().where(
        Document.stored_file_id == StoredFile.id,
        ConsolidationDocument.document_id == Document.id,
        ConsolidationVersion.id == ConsolidationDocument.consolidation_version_id,
        ConsolidationVersion.consolidation_id == consolidation_id,
    )

    def filtered(stmt: Select[Any]) -> Select[Any]:
        stmt = _visible_files(stmt, scope).where(
            StoredFile.legal_entity_id == entity_id,
            or_(StoredFile.consolidation_id == consolidation_id, bound),
        )
        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