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

"""Таблицы `financial_statement_*` (Срез 3b): вставка формы при маппинге, очистка файла, чтение сумм и флагов значений.
Только SQL, без правил: правила — `domain/statement_mapping.py`, `domain/statement_checks.py`."""

import datetime
import uuid
from collections.abc import Callable, Collection, Sequence
from dataclasses import dataclass
from decimal import Decimal
from typing import Any

from sqlalchemy import delete, func, insert, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql.expression import ScalarSelect

from src.domain.statement_aggregate import AggLine, CodeAggregate, aggregate_by_code
from src.domain.statement_issues import IssueDraft
from src.domain.statement_mapping import StatementDraft
from src.models.enums import (
    DocumentIssueKind,
    DocumentStatus,
    DocumentType,
    FileStatus,
    IssueSeverity,
    IssueSource,
    PeriodQuarter,
    StatementLineStatus,
)
from src.models.files import Document, StoredFile
from src.models.statements import (
    DocumentIssue,
    FinancialStatementColumn,
    FinancialStatementLine,
    FinancialStatementMeta,
    FinancialStatementValue,
)

__all__ = [
    "BATCH",
    "add_issues",
    "add_line",
    "add_value",
    "code_amount",
    "column_ids",
    "delete_documents_statements",
    "delete_file_statements",
    "delete_issues",
    "IssueCounts",
    "form_aggregate",
    "form_lines",
    "get_meta",
    "insert_statement",
    "issue_counts_by_document",
    "issues",
    "load_form",
    "lock_cells",
    "file_form_periods",
    "pair_amounts",
    "replace_issues",
    "value_flags",
]

# строк на один INSERT: у строки формы 15 колонок, у значения 10, у `document_issue` 16; 16 × 1000 < 32767 параметров
# asyncpg
BATCH = 1000


@dataclass(frozen=True)
class IssueCounts:
    """Число записей `document_issue` документа по `severity`."""

    error: int = 0
    info: int = 0


async def insert_statement(
    session: AsyncSession, *, document_id: uuid.UUID, job_id: uuid.UUID, draft: StatementDraft
) -> FinancialStatementMeta:
    """Шапка, колонки, строки и значения формы из черновика (без commit); возвращает шапку."""
    meta = FinancialStatementMeta(
        id=uuid.uuid4(),
        document_id=document_id,
        recognition_job_id=job_id,
        statement_type=draft.statement_type,
        **draft.header,
    )
    session.add(meta)
    await session.flush()
    column_ids: dict[int, uuid.UUID] = {}
    columns: list[dict[str, Any]] = []
    for column in draft.columns:
        column_ids[column.column_index] = uuid.uuid4()
        columns.append(
            {
                "id": column_ids[column.column_index],
                "meta_id": meta.id,
                "column_index": column.column_index,
                "label": column.label,
                "period_date": column.period_date,
                "is_current": column.is_current,
                "period_kind": column.period_kind,
                "label_confidence": column.label_confidence,
            }
        )
    if columns:
        await session.execute(insert(FinancialStatementColumn).values(columns))
    line_ids: dict[int, uuid.UUID] = {}  # по `row_index`: код строки может повторяться
    lines: list[dict[str, Any]] = []
    for line in draft.lines:
        line_ids[line.row_index] = uuid.uuid4()
        lines.append(
            {
                "id": line_ids[line.row_index],
                "meta_id": meta.id,
                "row_index": line.row_index,
                "source_ref": line.source_ref,
                "line_status": line.line_status,
                "line_code": line.line_code,
                "line_code_text": line.line_code_text,
                "line_code_original_value": line.line_code_original_value,
                "line_code_confidence": line.line_code_confidence,
                "line_code_modified": False,
                "parent_row_index": line.parent_row_index,
                "needs_review": line.needs_review,
                "line_name": line.line_name,
                "line_name_original_value": line.line_name,
                "line_name_confidence": line.line_name_confidence,
                "line_name_modified": False,
            }
        )
    for start in range(0, len(lines), BATCH):
        await session.execute(insert(FinancialStatementLine).values(lines[start : start + BATCH]))
    values = [
        {
            "id": uuid.uuid4(),
            "line_id": line_ids[value.row_index],
            "column_id": column_ids[value.column_index],
            "amount": value.amount,
            "string_value": value.string_value,
            "confidence": value.confidence,
            "is_recognized": value.is_recognized,
            "is_modified": False,
            "source_text": value.source_text,
            "needs_review": value.needs_review,
        }
        for value in draft.values
    ]
    for start in range(0, len(values), BATCH):
        await session.execute(insert(FinancialStatementValue).values(values[start : start + BATCH]))
    return meta


async def add_issues(
    session: AsyncSession, document_id: uuid.UUID, drafts: Sequence[IssueDraft], *, now: datetime.datetime
) -> None:
    """Записи `document_issue` документа из черновиков пакетами по `BATCH` (без commit)."""
    issues = [
        {
            "id": uuid.uuid4(),
            "document_id": document_id,
            "kind": draft.kind,
            "code": draft.code,
            "severity": draft.severity,
            "message": draft.message,
            "line_ids": [str(line_id) for line_id in draft.line_ids] or None,
            "column_id": draft.column_id,
            "column_index": draft.column_index,
            "codes": list(draft.codes) or None,
            "field": draft.field,
            "source": draft.source,
            "relation": draft.relation,
            "expected": draft.expected,
            "actual": draft.actual,
            "created_at": now,
        }
        for draft in drafts
    ]
    for start in range(0, len(issues), BATCH):
        await session.execute(insert(DocumentIssue).values(issues[start : start + BATCH]))


async def delete_issues(  # noqa: PLR0913 - фильтры записей независимы
    session: AsyncSession,
    document_id: uuid.UUID,
    *,
    codes: Collection[str],
    kind: DocumentIssueKind | None = None,
    source: IssueSource | None = None,
    column_id: uuid.UUID | None = None,
) -> int:
    """Удалить записи `document_issue` документа с кодами `codes` (и, если заданы, видом, источником, колонкой);
    вернуть число удалённых."""
    stmt = delete(DocumentIssue).where(DocumentIssue.document_id == document_id, DocumentIssue.code.in_(codes))
    if kind is not None:
        stmt = stmt.where(DocumentIssue.kind == kind)
    if source is not None:
        stmt = stmt.where(DocumentIssue.source == source)
    if column_id is not None:
        stmt = stmt.where(DocumentIssue.column_id == column_id)
    result = await session.execute(stmt)
    return int(result.rowcount or 0)  # type: ignore[attr-defined]


async def delete_file_statements(session: AsyncSession, file_id: uuid.UUID) -> None:
    """`document_issue`, значения, строки, колонки и шапки форм всех документов файла (в порядке FK)."""
    documents = select(Document.id).where(Document.stored_file_id == file_id).scalar_subquery()
    await _delete_statements(session, documents)


async def delete_documents_statements(session: AsyncSession, document_ids: Sequence[uuid.UUID]) -> None:
    """То же для перечисленных документов (очистка форм смешанного файла: ОСВ и прочие документы файла остаются)."""
    if document_ids:
        await _delete_statements(session, list(document_ids))


async def _delete_statements(session: AsyncSession, documents: ScalarSelect[uuid.UUID] | list[uuid.UUID]) -> None:
    await session.execute(delete(DocumentIssue).where(DocumentIssue.document_id.in_(documents)))
    metas = select(FinancialStatementMeta.id).where(FinancialStatementMeta.document_id.in_(documents)).scalar_subquery()
    lines = select(FinancialStatementLine.id).where(FinancialStatementLine.meta_id.in_(metas)).scalar_subquery()
    await session.execute(delete(FinancialStatementValue).where(FinancialStatementValue.line_id.in_(lines)))
    await session.execute(delete(FinancialStatementLine).where(FinancialStatementLine.meta_id.in_(metas)))
    await session.execute(delete(FinancialStatementColumn).where(FinancialStatementColumn.meta_id.in_(metas)))
    await session.execute(delete(FinancialStatementMeta).where(FinancialStatementMeta.document_id.in_(documents)))


async def replace_issues(  # noqa: PLR0913 - фильтр замены и новые записи независимы
    session: AsyncSession,
    document_id: uuid.UUID,
    drafts: Sequence[IssueDraft],
    *,
    kinds: Collection[DocumentIssueKind],
    core_codes: Collection[str],
    now: datetime.datetime,
) -> None:
    """Пересчёт после правки: удаляет записи документа видов `kinds` (любого источника) и записи источника `core` с
    кодами `core_codes`, затем вставляет `drafts`. Прочие записи (`warning` источника `service`, `MAPPING_WARNING` и
    т. п.) сохраняются."""
    await session.execute(
        delete(DocumentIssue).where(
            DocumentIssue.document_id == document_id,
            or_(
                DocumentIssue.kind.in_(kinds),
                (DocumentIssue.source == IssueSource.core) & DocumentIssue.code.in_(core_codes),
            ),
        )
    )
    await add_issues(session, document_id, drafts, now=now)


async def file_form_periods(
    session: AsyncSession, file_id: uuid.UUID
) -> list[tuple[uuid.UUID, DocumentType, int | None, PeriodQuarter | None]]:
    """Формы ББ/ОФР файла (документы с шапкой): `(id, тип, год, квартал)` — вход `duplicate_period_forms` (Р-18)."""
    stmt = (
        select(Document.id, Document.document_type, Document.period_year, Document.period_quarter)
        .join(FinancialStatementMeta, FinancialStatementMeta.document_id == Document.id)
        .where(Document.stored_file_id == file_id)
    )
    return [(row[0], row[1], row[2], row[3]) for row in (await session.execute(stmt)).all()]


async def unresolved_lines(session: AsyncSession, meta_id: uuid.UUID) -> list[tuple[uuid.UUID, int | None, str | None]]:
    """Строки `unreadable`/`missing` с непустой ячейкой (`amount` или `string_value`) по `row_index`:
    `(line_id, row_index, line_code_text)`. Строки без значений подтверждение не блокируют."""
    has_value = (
        select(FinancialStatementValue.id)
        .where(
            FinancialStatementValue.line_id == FinancialStatementLine.id,
            or_(FinancialStatementValue.amount.is_not(None), FinancialStatementValue.string_value.is_not(None)),
        )
        .exists()
    )
    stmt = (
        select(FinancialStatementLine.id, FinancialStatementLine.row_index, FinancialStatementLine.line_code_text)
        .where(
            FinancialStatementLine.meta_id == meta_id,
            FinancialStatementLine.line_status.in_((StatementLineStatus.unreadable, StatementLineStatus.missing)),
            has_value,
        )
        .order_by(FinancialStatementLine.row_index.asc().nulls_last(), FinancialStatementLine.id)
    )
    return [(row[0], row[1], row[2]) for row in (await session.execute(stmt)).all()]


async def form_aggregate(
    session: AsyncSession,
    meta_id: uuid.UUID,
    reference_name: Callable[[str], str | None] = lambda _: None,
) -> tuple[list[FinancialStatementColumn], dict[str, CodeAggregate]]:
    """Колонки формы (по `column_index`) и агрегат строк `ok` по коду (`domain/statement_aggregate`) — вход проверок
    §7, `parent_code` и достройки по шаблону. Порядок строк формы в БД задаёт `row_index`, но проверки §7 читают коды
    по возрастанию (M-9 2b)."""
    columns = list(
        (
            await session.execute(
                select(FinancialStatementColumn)
                .where(FinancialStatementColumn.meta_id == meta_id)
                .order_by(FinancialStatementColumn.column_index)
            )
        ).scalars()
    )
    lines = (
        await session.execute(
            select(FinancialStatementLine).where(
                FinancialStatementLine.meta_id == meta_id,
                FinancialStatementLine.line_status == StatementLineStatus.ok,
                FinancialStatementLine.line_code.is_not(None),
            )
        )
    ).scalars()
    by_line: dict[uuid.UUID, dict[uuid.UUID, Decimal | None]] = {}
    rows = (
        await session.execute(
            select(FinancialStatementValue.line_id, FinancialStatementValue.column_id, FinancialStatementValue.amount)
            .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
            .where(
                FinancialStatementLine.meta_id == meta_id,
                FinancialStatementLine.line_status == StatementLineStatus.ok,
            )
        )
    ).all()
    for line_id, column_id, amount in rows:
        by_line.setdefault(line_id, {})[column_id] = amount
    agg_lines = [
        AggLine(line.id, line.row_index, line.line_status, line.line_code, by_line.get(line.id, {})) for line in lines
    ]
    return columns, aggregate_by_code(agg_lines, reference_name)


async def value_flags(session: AsyncSession, meta_id: uuid.UUID) -> list[tuple[bool, bool | None, Decimal | None]]:
    """`(is_modified, is_recognized, confidence)` всех значений формы — вход `values_ok`."""
    rows = await session.execute(
        select(
            FinancialStatementValue.is_modified,
            FinancialStatementValue.is_recognized,
            FinancialStatementValue.confidence,
        )
        .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
        .where(FinancialStatementLine.meta_id == meta_id)
    )
    return [(bool(modified), recognized, confidence) for modified, recognized, confidence in rows]


async def lock_cells(
    session: AsyncSession, meta_id: uuid.UUID, cells: Sequence[tuple[uuid.UUID, uuid.UUID]]
) -> dict[tuple[uuid.UUID, uuid.UUID], FinancialStatementValue]:
    """Значения формы `meta_id` по `(line_id, column_id)` под `FOR UPDATE` с перечитыванием; чужие ячейки (строка не
    этой формы или колонка не её) в словарь не попадают."""
    if not cells:
        return {}
    line_ids = {line_id for line_id, _ in cells}
    stmt = (
        select(FinancialStatementValue)
        .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
        .join(FinancialStatementColumn, FinancialStatementColumn.id == FinancialStatementValue.column_id)
        .where(
            FinancialStatementLine.meta_id == meta_id,
            FinancialStatementColumn.meta_id == meta_id,
            FinancialStatementValue.line_id.in_(line_ids),
        )
        .order_by(FinancialStatementValue.id)
        .with_for_update(of=FinancialStatementValue)
        .execution_options(populate_existing=True)
    )
    wanted = set(cells)
    return {
        (value.line_id, value.column_id): value
        for value in (await session.execute(stmt)).scalars()
        if (value.line_id, value.column_id) in wanted
    }


async def form_lines(session: AsyncSession, meta_id: uuid.UUID) -> list[FinancialStatementLine]:
    """Строки формы (без значений) — адресация ячеек правки по `line_id` и `line_code`."""
    stmt = select(FinancialStatementLine).where(FinancialStatementLine.meta_id == meta_id)
    return list((await session.execute(stmt)).scalars())


async def template_lines_by_code(
    session: AsyncSession, meta_id: uuid.UUID, codes: Collection[str]
) -> list[tuple[FinancialStatementLine, bool]]:
    """Строки, созданные правкой (`row_index IS NULL`, статус `ok`), с кодами `codes`, и признак «есть непустая
    сумма» (`amount IS NOT NULL` хотя бы в одной ячейке)."""
    has_amount = (
        select(FinancialStatementValue.id)
        .where(
            FinancialStatementValue.line_id == FinancialStatementLine.id, FinancialStatementValue.amount.is_not(None)
        )
        .exists()
    )
    stmt = select(FinancialStatementLine, has_amount).where(
        FinancialStatementLine.meta_id == meta_id,
        FinancialStatementLine.row_index.is_(None),
        FinancialStatementLine.line_status == StatementLineStatus.ok,
        FinancialStatementLine.line_code.in_(codes),
    )
    return [(row[0], bool(row[1])) for row in (await session.execute(stmt)).all()]


async def delete_line(session: AsyncSession, line_id: uuid.UUID) -> None:
    """Удалить строку формы вместе со значениями (порядок FK); записи `document_issue` чистит вызывающий."""
    await session.execute(delete(FinancialStatementValue).where(FinancialStatementValue.line_id == line_id))
    await session.execute(delete(FinancialStatementLine).where(FinancialStatementLine.id == line_id))
    await session.flush()


async def remove_line_from_issues(session: AsyncSession, document_id: uuid.UUID, line_id: uuid.UUID) -> None:
    """Убрать `id` удалённой строки из `line_ids` всех записей `document_issue` документа (у `line_ids` нет FK)."""
    await session.execute(
        update(DocumentIssue)
        .where(
            DocumentIssue.document_id == document_id,
            DocumentIssue.line_ids.is_not(None),
            DocumentIssue.line_ids.contains([str(line_id)]),
        )
        .values(line_ids=DocumentIssue.line_ids.op("-")(str(line_id)))
        .execution_options(synchronize_session="fetch")
    )


async def column_ids(session: AsyncSession, meta_id: uuid.UUID) -> set[uuid.UUID]:
    """Id колонок формы."""
    stmt = select(FinancialStatementColumn.id).where(FinancialStatementColumn.meta_id == meta_id)
    return set((await session.execute(stmt)).scalars())


async def add_line(session: AsyncSession, meta_id: uuid.UUID, code: str, name: str) -> FinancialStatementLine:
    """Шаблонная строка, созданная правкой аналитика: `row_index = NULL` — признак «не из источника»; статус `ok`."""
    line = FinancialStatementLine(
        id=uuid.uuid4(),
        meta_id=meta_id,
        line_code=code,
        line_code_original_value=None,
        line_code_confidence=None,
        line_code_modified=False,
        line_name=name,
        line_name_original_value=None,
        line_name_confidence=None,
        line_name_modified=False,
        row_index=None,
        source_ref=None,
        line_status=StatementLineStatus.ok,
    )
    session.add(line)
    await session.flush()
    return line


async def add_value(session: AsyncSession, line_id: uuid.UUID, column_id: uuid.UUID) -> FinancialStatementValue:
    """Пустая ячейка, созданная правкой: не распознана (`is_recognized = NULL`), исходной суммы нет."""
    value = FinancialStatementValue(
        id=uuid.uuid4(),
        line_id=line_id,
        column_id=column_id,
        amount=None,
        string_value=None,
        confidence=None,
        is_recognized=None,
        is_modified=False,
        original_amount=None,
    )
    session.add(value)
    await session.flush()
    return value


async def get_meta(
    session: AsyncSession, document_id: uuid.UUID, *, refresh: bool = False
) -> FinancialStatementMeta | None:
    """Шапка формы документа (1:1) или `None`; `refresh` — перечитать атрибуты, даже если объект уже в сессии."""
    stmt = select(FinancialStatementMeta).where(FinancialStatementMeta.document_id == document_id)
    if refresh:
        stmt = stmt.execution_options(populate_existing=True)
    return (await session.execute(stmt)).scalar_one_or_none()


async def load_form(
    session: AsyncSession, meta_id: uuid.UUID
) -> tuple[list[FinancialStatementColumn], list[tuple[FinancialStatementLine, list[FinancialStatementValue]]]]:
    """Колонки по `column_index` и строки по `line_code`, `row_index` (`NULL` последними), `id` с их значениями по
    `column_index`."""
    columns = list(
        (
            await session.execute(
                select(FinancialStatementColumn)
                .where(FinancialStatementColumn.meta_id == meta_id)
                .order_by(FinancialStatementColumn.column_index)
            )
        ).scalars()
    )
    order = {column.id: column.column_index for column in columns}
    lines = list(
        (
            await session.execute(
                select(FinancialStatementLine)
                .where(FinancialStatementLine.meta_id == meta_id)
                .order_by(
                    FinancialStatementLine.line_code,
                    FinancialStatementLine.row_index.asc().nulls_last(),
                    FinancialStatementLine.id,
                )
            )
        ).scalars()
    )
    values = (
        await session.execute(
            select(FinancialStatementValue).where(FinancialStatementValue.line_id.in_([line.id for line in lines]))
        )
    ).scalars()
    by_line: dict[uuid.UUID, list[FinancialStatementValue]] = {line.id: [] for line in lines}
    for value in values:
        by_line[value.line_id].append(value)
    for items in by_line.values():
        items.sort(key=lambda value: order[value.column_id])
    return columns, [(line, by_line[line.id]) for line in lines]


async def code_amount(
    session: AsyncSession,
    meta_id: uuid.UUID,
    code: str,
    *,
    current: bool = False,
    period_date: datetime.date | None = None,
) -> Decimal | None:
    """`SUM(amount)` строк `ok` с кодом `code` в колонке `is_current` (`current`) или с датой `period_date`; колонки
    нет или непустых сумм нет — `None`."""
    column_stmt = select(FinancialStatementColumn.id).where(FinancialStatementColumn.meta_id == meta_id)
    if current:
        column_stmt = column_stmt.where(FinancialStatementColumn.is_current.is_(True))
    else:
        column_stmt = column_stmt.where(FinancialStatementColumn.period_date == period_date)
    column_id = (await session.execute(column_stmt.order_by(FinancialStatementColumn.column_index).limit(1))).scalar()
    if column_id is None:
        return None
    stmt = (
        select(func.sum(FinancialStatementValue.amount))
        .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
        .where(
            FinancialStatementLine.meta_id == meta_id,
            FinancialStatementLine.line_code == code,
            FinancialStatementLine.line_status == StatementLineStatus.ok,
            FinancialStatementValue.column_id == column_id,
        )
    )
    return (await session.execute(stmt)).scalar_one_or_none()


async def _pair_document(session: AsyncSession, document: Document, pair_type: DocumentType) -> Document | None:
    """Форма-пара (D-5): другой тип формы того же ЮЛ, года и квартала, файл не `SUPERSEDED`; подтверждённая первой,
    затем самая поздняя по `created_at`. Документ без шапки формы не считается парой (иначе он затенял бы настоящую)."""
    stmt = (
        select(Document)
        .join(StoredFile, StoredFile.id == Document.stored_file_id)
        .join(FinancialStatementMeta, FinancialStatementMeta.document_id == Document.id)
        .where(
            Document.legal_entity_id == document.legal_entity_id,
            Document.period_year == document.period_year,
            Document.period_quarter == document.period_quarter,
            Document.document_type == pair_type,
            StoredFile.status != FileStatus.SUPERSEDED,
        )
        .order_by((Document.status == DocumentStatus.CONFIRMED).desc(), Document.created_at.desc())
        .limit(1)
    )
    return (await session.execute(stmt)).scalar_one_or_none()


async def pair_amounts(
    session: AsyncSession, document: Document
) -> tuple[tuple[Decimal | None, Decimal | None], Decimal | None]:
    """Вход межформенной проверки (D-5): `(1370 на отчётную дату, 1370 на 31.12 прошлого года)` из ББ и `2400` ОФР
    в колонке отчётного периода. Нет пары, шапки пары, строки или колонки — соответствующий `None`."""
    if document.period_year is None or document.period_quarter is None:
        return (None, None), None
    is_balance = document.document_type == DocumentType.BALANCE
    pair = await _pair_document(session, document, DocumentType.INCOME if is_balance else DocumentType.BALANCE)
    own_meta = await get_meta(session, document.id)
    pair_meta = await get_meta(session, pair.id) if pair is not None else None
    balance_meta, income_meta = (own_meta, pair_meta) if is_balance else (pair_meta, own_meta)
    balance: tuple[Decimal | None, Decimal | None] = (None, None)
    if balance_meta is not None:
        year_start = datetime.date(document.period_year - 1, 12, 31)
        balance = (
            await code_amount(session, balance_meta.id, "1370", current=True),
            await code_amount(session, balance_meta.id, "1370", period_date=year_start),
        )
    income = await code_amount(session, income_meta.id, "2400", current=True) if income_meta is not None else None
    return balance, income


async def issues(session: AsyncSession, document_id: uuid.UUID) -> list[DocumentIssue]:
    """Записи `document_issue` документа по `created_at`, `id`."""
    stmt = (
        select(DocumentIssue)
        .where(DocumentIssue.document_id == document_id)
        .order_by(DocumentIssue.created_at, DocumentIssue.id)
    )
    return list((await session.execute(stmt)).scalars())


async def issue_counts_by_document(
    session: AsyncSession, document_ids: Sequence[uuid.UUID]
) -> dict[uuid.UUID, IssueCounts]:
    """Число записей `document_issue` по `severity` одним запросом на список документов (P3-23); документа с записями
    нет в ответе (читать через `.get(id, IssueCounts())`)."""
    if not document_ids:
        return {}
    stmt = (
        select(DocumentIssue.document_id, DocumentIssue.severity, func.count(DocumentIssue.id))
        .where(DocumentIssue.document_id.in_(document_ids))
        .group_by(DocumentIssue.document_id, DocumentIssue.severity)
    )
    raw: dict[uuid.UUID, dict[IssueSeverity, int]] = {}
    for document_id, severity, count in (await session.execute(stmt)).all():
        raw.setdefault(document_id, {})[severity] = count
    return {
        document_id: IssueCounts(
            error=by_severity.get(IssueSeverity.error, 0), info=by_severity.get(IssueSeverity.info, 0)
        )
        for document_id, by_severity in raw.items()
    }