app/api/media.py

from __future__ import annotations

from collections.abc import AsyncIterator
from pathlib import Path

import aiofiles
from fastapi import APIRouter, Request
from fastapi.responses import StreamingResponse

from app.core.exceptions import NotFoundError
from app.core.response import SuccessResponse
from app.security.assets import CanonicalAssetNotFoundError
from app.services.media_service import MediaProcessor, Operation

router = APIRouter(tags=["media"])


def get_processor(request: Request) -> MediaProcessor:
    return request.app.state.container.processor


async def execute_operation(request: Request, name: str, operation: Operation) -> SuccessResponse:
    request.state.operation = name
    processor = get_processor(request)
    resolved = await processor.resolver.resolve(request)
    return await processor.run(resolved, name, operation)


async def stream_file(path: Path) -> AsyncIterator[bytes]:
    async with aiofiles.open(path, "rb") as media:
        while chunk := await media.read(1024 * 1024):
            yield chunk


@router.get("/v1/media/{request_id}/{filename}", name="download_media")
async def download_media(request: Request, request_id: str, filename: str) -> StreamingResponse:
    request.state.operation = "media.download"
    container = request.app.state.container
    # The filesystem locator is not an authorization token. In authenticated
    # deployments an output must have been issued by the pipeline to the
    # caller's authoritative workspace before it can be downloaded.
    if container.settings.auth_enabled:
        context = getattr(request.state, "auth", None)
        if context is None or not context.workspace_id:
            raise NotFoundError("Output file not found")
        try:
            asset = await container.assets.get_owned(
                workspace_id=context.workspace_id,
                request_id=request_id,
                filename=filename,
            )
            path = container.cleanup.resolve_download(request_id, asset.filename)
            await container.assets.verify_file(asset, path)
        except CanonicalAssetNotFoundError as exc:
            raise NotFoundError("Output file not found") from exc
    else:
        path = container.cleanup.resolve_download(request_id, filename)
    media_type = container.validator.infer_mime(path)
    headers = {
        "Content-Disposition": f'attachment; filename="{path.name}"',
        "Content-Length": str(path.stat().st_size),
        "X-Request-ID": request_id,
    }
    return StreamingResponse(stream_file(path), media_type=media_type, headers=headers)