Files
kite/backend/api/documents/upload.py
T

280 lines
8.6 KiB
Python

"""Document upload endpoints."""
from __future__ import annotations
import uuid
from pathlib import Path
from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession
from config import settings
from db.models import CloudConnection, Document, User
from deps.auth import get_regular_user
from deps.db import get_db
from deps.utils import get_client_ip
from services.audit import write_audit_log
from services.rate_limiting import account_limiter
from storage import get_storage_backend
from storage.cloud_backend_factory import build_cloud_backend
from storage.cloud_utils import decrypt_credentials
from storage.exceptions import CloudConnectionError
from tasks.document_tasks import extract_and_classify
from api.documents.shared import UploadUrlRequest, _CLOUD_PROVIDERS
router = APIRouter()
def _new_minio_document(user_id: uuid.UUID, filename: str, content_type: str) -> Document:
doc_id = uuid.uuid4()
suffix = Path(filename).suffix.lower()
return Document(
id=doc_id,
user_id=user_id,
filename=filename,
content_type=content_type,
size_bytes=0,
storage_backend="minio",
status="pending",
object_key=f"{user_id}/{doc_id}/{uuid.uuid4()}{suffix}",
)
async def _create_presigned_upload(
session: AsyncSession,
user_id: uuid.UUID,
filename: str,
content_type: str,
) -> dict:
doc = _new_minio_document(user_id, filename, content_type)
session.add(doc)
await session.commit()
upload_url = await get_storage_backend().generate_presigned_put_url(
doc.object_key, expires_minutes=15
)
return {"upload_url": upload_url, "document_id": str(doc.id)}
async def _get_active_cloud_connection(
session: AsyncSession,
user_id: uuid.UUID,
provider: str,
) -> CloudConnection:
result = await session.execute(
select(CloudConnection).where(
CloudConnection.user_id == user_id,
CloudConnection.provider == provider,
CloudConnection.status == "ACTIVE",
)
)
conn = result.scalar_one_or_none()
if conn is None:
raise HTTPException(
status_code=404,
detail=f"No active {provider} connection found. Please connect in Settings.",
)
return conn
def _decrypt_cloud_credentials(conn: CloudConnection, user_id: uuid.UUID) -> dict:
return decrypt_credentials(settings.cloud_creds_key.encode(), str(user_id), conn.credentials_enc)
async def _record_upload(
session: AsyncSession,
request: Request,
current_user: User,
doc: Document,
size_bytes: int,
storage_backend: str,
) -> None:
await write_audit_log(
session,
event_type="document.uploaded",
user_id=current_user.id,
actor_id=current_user.id,
resource_id=doc.id,
ip_address=get_client_ip(request) if request else None,
metadata_={"size_bytes": size_bytes, "storage_backend": storage_backend},
)
@router.post("/upload-url")
@account_limiter.limit("100/minute")
async def request_upload_url(
request: Request,
body: UploadUrlRequest,
session: AsyncSession = Depends(get_db),
current_user: User = Depends(get_regular_user),
):
"""Create a pending Document row and return a presigned PUT URL."""
request.state.current_user = current_user
return await _create_presigned_upload(
session,
current_user.id,
body.filename,
body.content_type,
)
@router.post("/upload")
@account_limiter.limit("100/minute")
async def upload_document(
request: Request,
file: UploadFile = File(...),
target_backend: str = Form("minio"),
cloud_folder_path: str = Form(None),
session: AsyncSession = Depends(get_db),
current_user: User = Depends(get_regular_user),
):
"""Direct multipart upload endpoint supporting cloud backends."""
request.state.current_user = current_user
if target_backend == "minio":
return await _create_presigned_upload(
session,
current_user.id,
file.filename or "upload",
file.content_type or "application/octet-stream",
)
if target_backend not in _CLOUD_PROVIDERS:
raise HTTPException(
status_code=422,
detail=f"Invalid target_backend '{target_backend}'. Valid values: minio, {', '.join(sorted(_CLOUD_PROVIDERS))}",
)
conn = await _get_active_cloud_connection(session, current_user.id, target_backend)
credentials = _decrypt_cloud_credentials(conn, current_user.id)
file_bytes = await file.read()
filename = file.filename or "upload"
content_type = file.content_type or "application/octet-stream"
extension = Path(filename).suffix.lower()
doc_id = uuid.uuid4()
cloud_backend = build_cloud_backend(target_backend, credentials)
try:
object_key = await cloud_backend.put_object(
str(current_user.id),
str(doc_id),
file_bytes,
extension,
content_type,
cloud_folder=cloud_folder_path or None,
original_filename=filename if cloud_folder_path else None,
)
except CloudConnectionError as exc:
raise HTTPException(
status_code=503,
detail="Cloud connection requires re-authentication. Please reconnect in Settings.",
) from exc
if cloud_folder_path:
from services.cloud_cache import invalidate_provider_cache # lazy import
invalidate_provider_cache(str(current_user.id), target_backend)
doc = Document(
id=doc_id,
user_id=current_user.id,
filename=filename,
content_type=content_type,
size_bytes=len(file_bytes),
storage_backend=target_backend,
status="uploaded",
object_key=object_key,
)
session.add(doc)
await _record_upload(session, request, current_user, doc, len(file_bytes), target_backend)
await session.commit()
extract_and_classify.delay(str(doc.id))
return {"document_id": str(doc.id), "storage_backend": target_backend}
@router.post("/{doc_id}/confirm")
@account_limiter.limit("100/minute")
async def confirm_upload(
doc_id: str,
request: Request,
session: AsyncSession = Depends(get_db),
current_user: User = Depends(get_regular_user),
):
"""Confirm a presigned PUT upload and enforce quota."""
request.state.current_user = current_user
try:
uid = uuid.UUID(doc_id)
except ValueError:
raise HTTPException(status_code=404, detail="Document not found")
doc = await session.get(Document, uid)
if doc is None or doc.user_id != current_user.id:
raise HTTPException(status_code=404, detail="Document not found")
try:
size = await get_storage_backend().stat_object(doc.object_key)
except Exception as exc:
code = getattr(exc, "code", "")
if code == "NoSuchKey":
raise HTTPException(
status_code=422,
detail="Upload not found — presigned URL may have expired",
)
raise HTTPException(status_code=502, detail=f"Storage error: {exc}")
doc.size_bytes = size
await session.flush()
result = await session.execute(
text(
"UPDATE quotas "
"SET used_bytes = used_bytes + :delta "
"WHERE user_id = :uid "
" AND (used_bytes + :delta) <= limit_bytes "
"RETURNING used_bytes, limit_bytes"
),
{"delta": size, "uid": doc.user_id.hex},
)
row = result.fetchone()
if row is None:
quota_result = await session.execute(
text("SELECT used_bytes, limit_bytes FROM quotas WHERE user_id = :uid"),
{"uid": doc.user_id.hex},
)
q = quota_result.fetchone()
await session.delete(doc)
try:
await get_storage_backend().delete_object(doc.object_key)
except Exception:
pass
await session.commit()
raise HTTPException(
status_code=413,
detail={
"used_bytes": q.used_bytes if q else 0,
"limit_bytes": q.limit_bytes if q else 0,
"rejected_bytes": size,
},
)
used_bytes = row.used_bytes
doc.status = "uploaded"
await _record_upload(session, request, current_user, doc, size, "minio")
await session.commit()
extract_and_classify.delay(str(doc.id))
return {
"id": str(doc.id),
"size_bytes": size,
"used_bytes": used_bytes,
"status": "uploaded",
}