"""Shared api upload/enqueue service (used by Telegram and B2B adapters).""" from __future__ import annotations import uuid from typing import TYPE_CHECKING from fastapi import HTTPException from src.contract_check.api.schemas import DocumentUploadResponse from src.contract_check.core.billing.errors import NoCredits from src.contract_check.core.billing.quota import ( release_document_slot, reserve_document_slot, ) from src.contract_check.core.config import get_settings from src.contract_check.core.credits import adjust_credits, reserve_credit from src.contract_check.core.db.repositories import ( DocumentRepository, JobRepository, UserRepository, ) from src.contract_check.core.db.repositories.credits import CreditsRepository from src.contract_check.core.extraction.formats import SUPPORTED_SUFFIXES from src.contract_check.core.logging import get_logger, new_correlation_id from src.contract_check.core.metrics import credits_reserved, documents_uploaded from src.contract_check.core.mq.messages import DocumentUploaded from src.contract_check.core.mq.topology import RK_EXTRACT from src.contract_check.core.s3 import original_key if TYPE_CHECKING: from fastapi import UploadFile from sqlalchemy.ext.asyncio import AsyncSession from src.contract_check.core.mq.publisher import Publisher from src.contract_check.core.s3.port import Storage log = get_logger(__name__) def _content_type_from_suffix(suffix: str) -> str: return { ".pdf": "application/pdf", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", ".rtf": "application/rtf", ".txt": "text/plain", ".csv": "text/csv", ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".tif": "image/tiff", ".tiff": "image/tiff", }.get(suffix, "application/octet-stream") async def upload_and_enqueue( session: AsyncSession, storage: Storage, publisher: Publisher, user_id: uuid.UUID, file: UploadFile, ) -> DocumentUploadResponse: """Validate, store in MinIO, create document/job rows, reserve credit, publish. Returns 202 payload: {document_id, correlation_id, credits_left}. Raises HTTPException on validation, credit, storage, DB, or MQ failure. """ if file.filename is None: raise HTTPException(status_code=400, detail="filename is required") suffix = file.filename.lower().split(".")[-1] if "." in file.filename else "" suffix = f".{suffix}" if suffix else "" if suffix not in SUPPORTED_SUFFIXES: raise HTTPException( status_code=400, detail=f"unsupported format {suffix!r}; supported: {sorted(SUPPORTED_SUFFIXES)}", ) settings = get_settings() # Reject Billing Holds before reading the body or writing storage. if await UserRepository(session).get_billing_hold(user_id): raise HTTPException(status_code=402, detail="billing hold") credits_repo = CreditsRepository(session) doc_repo = DocumentRepository(session) job_repo = JobRepository(session) document_id = uuid.uuid4() correlation_id = new_correlation_id() s3_key = original_key(str(user_id), str(document_id), suffix) content_type = file.content_type or _content_type_from_suffix(suffix) data = await _read_upload_with_limit(file, settings.max_upload_bytes) if len(data) == 0: raise HTTPException(status_code=400, detail="empty file") try: await storage.put(s3_key, data, content_type=content_type) except Exception as exc: log.error("s3_upload_failed", document_id=str(document_id), error=str(exc)) raise HTTPException(status_code=500, detail="failed to store document") from exc try: await doc_repo.create( document_id=document_id, user_id=user_id, s3_key=s3_key, filename=file.filename or "document", mime=content_type, bytes_=len(data), status="queued", ) await job_repo.create( document_id=document_id, correlation_id=uuid.UUID(correlation_id), queue="extract", status="pending", ) except Exception as exc: log.error("db_enqueue_failed", document_id=str(document_id), error=str(exc)) # Best-effort cleanup so a DB failure does not orphan the stored object. try: await storage.delete(s3_key) except Exception as cleanup_exc: # noqa: BLE001 log.warning( "upload_cleanup_failed", document_id=str(document_id), s3_key=s3_key, error=str(cleanup_exc), ) raise HTTPException(status_code=500, detail="failed to enqueue document") from exc source = "credits" if settings.plans_enabled: try: source = await reserve_document_slot(session, user_id, document_id) except NoCredits: raise HTTPException(status_code=402, detail="no credits available") from None else: if not await reserve_credit(session, user_id, document_id=document_id): raise HTTPException(status_code=402, detail="no credits available") await session.commit() credits_reserved.inc() try: msg = DocumentUploaded( correlation_id=uuid.UUID(correlation_id), document_id=document_id, user_id=user_id, s3_key=s3_key, filename=file.filename, mime=content_type, ) await publisher.publish(msg, routing_key=RK_EXTRACT) documents_uploaded.inc() except Exception as exc: log.error("mq_publish_failed", document_id=str(document_id), error=str(exc)) if source == "quota": await release_document_slot(session, document_id) else: await adjust_credits( session, user_id, 1, kind="refund_auto", document_id=document_id, ) await doc_repo.mark_failed(document_id, stage="publish_failed") await session.commit() raise HTTPException(status_code=500, detail="failed to publish job") from exc credits_left = await credits_repo.get_balance(user_id) return DocumentUploadResponse( document_id=str(document_id), correlation_id=str(correlation_id), credits_left=credits_left, ) async def _read_upload_with_limit(file: UploadFile, max_bytes: int) -> bytes: """Read ``file`` in chunks, raising 413 if ``max_bytes`` is exceeded. The file is never fully buffered in memory beyond ``max_bytes + 1``. """ chunks: list[bytes] = [] total = 0 chunk_size = 64 * 1024 while True: chunk = await file.read(chunk_size) if not chunk: break total += len(chunk) if total > max_bytes: raise HTTPException( status_code=413, detail=f"file exceeds maximum size of {max_bytes} bytes", ) chunks.append(chunk) return b"".join(chunks)