persistence/document_uploads.py

"""Owned persistence operations for queued asynchronous document uploads."""

from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import datetime
from typing import Optional

from sqlalchemy import case, select
from sqlalchemy.exc import IntegrityError

from backend.app.storage.base import document_pdf_object_key, validate_object_key
from persistence.document_ingestion import (
    DOCUMENT_STATUS_FAILED,
    DOCUMENT_STATUS_PROCESSING,
    DOCUMENT_STATUS_QUEUED,
    DOCUMENT_STATUS_READY,
)
from persistence.ownership import require_existing_user, require_owned_document


DOCUMENT_LIFECYCLE_STATUSES = {
    DOCUMENT_STATUS_QUEUED,
    DOCUMENT_STATUS_PROCESSING,
    DOCUMENT_STATUS_READY,
    DOCUMENT_STATUS_FAILED,
}
_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")


class UploadRequestConflictError(ValueError):
    """Raised when an owned upload request UUID is reused for different bytes."""


class DocumentUploadStateError(RuntimeError):
    """Raised when persisted document state cannot be safely exposed or recovered."""


@dataclass(frozen=True)
class StoredDocumentStatus:
    document_id: str
    user_id: str
    filename: Optional[str]
    source_hash: Optional[str]
    status: str
    error_message: Optional[str]
    created_at: datetime
    updated_at: datetime


@dataclass(frozen=True)
class QueuedDocumentWriteResult:
    document: StoredDocumentStatus
    created: bool


def _safe_error_message(value: object) -> Optional[str]:
    if value is None:
        return None
    normalized = " ".join(str(value).split()).strip()
    return normalized[:2000] or None


def _project_document(document: object) -> StoredDocumentStatus:
    status = str(getattr(document, "status"))
    if status not in DOCUMENT_LIFECYCLE_STATUSES:
        raise DocumentUploadStateError("Document has an unsupported lifecycle status.")
    raw_source_hash = getattr(document, "source_hash", None)
    source_hash = str(raw_source_hash) if raw_source_hash is not None else None
    return StoredDocumentStatus(
        document_id=str(getattr(document, "id")),
        user_id=str(getattr(document, "user_id")),
        filename=getattr(document, "filename", None),
        source_hash=source_hash,
        status=status,
        error_message=_safe_error_message(getattr(document, "error_message", None)),
        created_at=getattr(document, "created_at"),
        updated_at=getattr(document, "updated_at"),
    )


def _matching_upload_request(
    session: object,
    *,
    user_id: str,
    upload_request_id: str,
):
    from persistence.models import Document

    return session.execute(
        select(Document).where(
            Document.user_id == user_id,
            Document.upload_request_id == upload_request_id,
        )
    ).scalar_one_or_none()


def recover_owned_upload_request(
    *,
    user_id: str,
    upload_request_id: str,
    source_hash: str,
) -> Optional[StoredDocumentStatus]:
    """Recover an owned upload before any new object is written."""
    from persistence.db import SessionLocal

    with SessionLocal() as session:
        document = _matching_upload_request(
            session,
            user_id=user_id,
            upload_request_id=upload_request_id,
        )
        if document is None:
            return None
        if document.source_hash is None:
            raise DocumentUploadStateError("Document source hash is unavailable.")
        if document.source_hash != source_hash:
            raise UploadRequestConflictError(
                "upload_request_id was already used with a different document."
            )
        return _project_document(document)


def find_owned_usable_duplicate(*, user_id: str, source_hash: str) -> Optional[StoredDocumentStatus]:
    """Prefer the newest active match, then the newest ready match for this owner."""
    from persistence.db import SessionLocal
    from persistence.models import Document

    active_statuses = (DOCUMENT_STATUS_QUEUED, DOCUMENT_STATUS_PROCESSING)
    with SessionLocal() as session:
        document = session.execute(
            select(Document)
            .where(
                Document.user_id == user_id,
                Document.source_hash == source_hash,
                Document.status.in_((*active_statuses, DOCUMENT_STATUS_READY)),
            )
            .order_by(
                case((Document.status.in_(active_statuses), 0), else_=1),
                Document.created_at.desc(),
                Document.id.desc(),
            )
            .limit(1)
        ).scalar_one_or_none()
        return _project_document(document) if document is not None else None


def create_or_recover_queued_upload(
    *,
    user_id: str,
    document_id: str,
    upload_request_id: Optional[str],
    filename: str,
    source_hash: str,
    object_key: str,
    content_type: str,
    byte_size: int,
    embedding_model: str,
    embedding_format: str,
    retrieval_mode: str,
    reranker_model: str,
    k_initial: int,
    k_final: int,
) -> QueuedDocumentWriteResult:
    """Create one queued document or recover the winner of an idempotency race."""
    from persistence.db import SessionLocal
    from persistence.models import Document

    validated_key = validate_object_key(object_key)
    if validated_key != document_pdf_object_key(document_id):
        raise ValueError("Document object key does not match the generated document ID.")
    if _SHA256_PATTERN.fullmatch(source_hash) is None:
        raise ValueError("source_hash must be a lowercase SHA-256 digest.")
    if byte_size <= 0:
        raise ValueError("byte_size must be positive.")

    with SessionLocal() as session:
        try:
            require_existing_user(session, user_id)
            document = Document(
                id=document_id,
                user_id=user_id,
                source_type="upload",
                filename=filename,
                source_hash=source_hash,
                cache_key=f"async-upload:{document_id}",
                object_key=validated_key,
                content_type=content_type,
                byte_size=byte_size,
                upload_request_id=upload_request_id,
                status=DOCUMENT_STATUS_QUEUED,
                error_message=None,
                embedding_model=embedding_model,
                embedding_format=embedding_format,
                retrieval_mode=retrieval_mode,
                reranker_model=reranker_model,
                k_initial=k_initial,
                k_final=k_final,
            )
            session.add(document)
            session.flush()
            session.refresh(document)
            projected = _project_document(document)
            session.commit()
            return QueuedDocumentWriteResult(document=projected, created=True)
        except IntegrityError:
            session.rollback()
            if upload_request_id is None:
                raise
            winner = _matching_upload_request(
                session,
                user_id=user_id,
                upload_request_id=upload_request_id,
            )
            if winner is None:
                raise
            if winner.source_hash != source_hash:
                raise UploadRequestConflictError(
                    "upload_request_id was already used with a different document."
                )
            return QueuedDocumentWriteResult(
                document=_project_document(winner),
                created=False,
            )
        except Exception:
            session.rollback()
            raise


def get_owned_document_status(*, user_id: str, document_id: str) -> StoredDocumentStatus:
    """Return one document only after ownership is enforced by the query."""
    from persistence.db import SessionLocal

    with SessionLocal() as session:
        document = require_owned_document(
            session,
            user_id=user_id,
            document_id=document_id,
        )
        return _project_document(document)


def should_enqueue_owned_document(*, user_id: str, document_id: str) -> StoredDocumentStatus:
    """Lock and return current state before an idempotent re-publication decision."""
    from persistence.db import SessionLocal
    from persistence.models import Document

    with SessionLocal() as session:
        document = session.execute(
            select(Document)
            .where(Document.id == document_id, Document.user_id == user_id)
            .with_for_update()
        ).scalar_one_or_none()
        if document is None:
            from persistence.ownership import OwnedResourceNotFoundError

            raise OwnedResourceNotFoundError("Document not found.")
        return _project_document(document)