102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
"""Screenshot upload and extraction API."""
|
|
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import get_settings
|
|
from app.models.database import get_db
|
|
from app.models import Case, Screenshot, Transaction
|
|
from app.schemas import ScreenshotResponse, ScreenshotListResponse, TransactionListResponse
|
|
from app.services.extractor import extract_and_save
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _allowed(filename: str) -> bool:
|
|
ext = (Path(filename).suffix or "").lstrip(".").lower()
|
|
return ext in get_settings().allowed_extensions
|
|
|
|
|
|
@router.get("/{case_id}/screenshots", response_model=ScreenshotListResponse)
|
|
async def list_screenshots(case_id: int, db: AsyncSession = Depends(get_db)):
|
|
r = await db.execute(select(Case).where(Case.id == case_id))
|
|
if not r.scalar_one_or_none():
|
|
raise HTTPException(status_code=404, detail="Case not found")
|
|
r = await db.execute(select(Screenshot).where(Screenshot.case_id == case_id).order_by(Screenshot.created_at))
|
|
screenshots = r.scalars().all()
|
|
return ScreenshotListResponse(items=[ScreenshotResponse.model_validate(s) for s in screenshots])
|
|
|
|
|
|
@router.post("/{case_id}/screenshots", response_model=ScreenshotListResponse)
|
|
async def upload_screenshots(
|
|
case_id: int,
|
|
files: list[UploadFile] = File(...),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
r = await db.execute(select(Case).where(Case.id == case_id))
|
|
case = r.scalar_one_or_none()
|
|
if not case:
|
|
raise HTTPException(status_code=404, detail="Case not found")
|
|
settings = get_settings()
|
|
upload_dir = settings.upload_dir.resolve()
|
|
case_dir = upload_dir / str(case_id)
|
|
case_dir.mkdir(parents=True, exist_ok=True)
|
|
created: list[Screenshot] = []
|
|
for f in files:
|
|
if not f.filename or not _allowed(f.filename):
|
|
continue
|
|
stem = uuid.uuid4().hex[:12]
|
|
suffix = Path(f.filename).suffix
|
|
path = case_dir / f"{stem}{suffix}"
|
|
content = await f.read()
|
|
path.write_bytes(content)
|
|
rel_path = str(path.relative_to(upload_dir))
|
|
screenshot = Screenshot(
|
|
case_id=case_id,
|
|
filename=f.filename,
|
|
file_path=rel_path,
|
|
status="pending",
|
|
)
|
|
db.add(screenshot)
|
|
created.append(screenshot)
|
|
await db.commit()
|
|
for s in created:
|
|
await db.refresh(s)
|
|
return ScreenshotListResponse(items=[ScreenshotResponse.model_validate(s) for s in created])
|
|
|
|
|
|
@router.post("/{case_id}/screenshots/{screenshot_id}/extract", response_model=TransactionListResponse)
|
|
async def extract_transactions(
|
|
case_id: int,
|
|
screenshot_id: int,
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
r = await db.execute(select(Screenshot).where(Screenshot.id == screenshot_id, Screenshot.case_id == case_id))
|
|
screenshot = r.scalar_one_or_none()
|
|
if not screenshot:
|
|
raise HTTPException(status_code=404, detail="Screenshot not found")
|
|
settings = get_settings()
|
|
full_path = settings.upload_dir.resolve() / screenshot.file_path
|
|
if not full_path.exists():
|
|
raise HTTPException(status_code=404, detail="File not found on disk")
|
|
image_bytes = full_path.read_bytes()
|
|
try:
|
|
transactions = await extract_and_save(case_id, screenshot_id, image_bytes)
|
|
except Exception as e:
|
|
r = await db.execute(select(Screenshot).where(Screenshot.id == screenshot_id))
|
|
sc = r.scalar_one_or_none()
|
|
if sc:
|
|
sc.status = "failed"
|
|
await db.commit()
|
|
raise HTTPException(status_code=502, detail=f"Extraction failed: {e!s}")
|
|
r = await db.execute(select(Screenshot).where(Screenshot.id == screenshot_id))
|
|
sc = r.scalar_one_or_none()
|
|
if sc:
|
|
sc.status = "extracted"
|
|
await db.commit()
|
|
return TransactionListResponse(items=transactions)
|