Aurora-Core/app/services/artwork/download.py

"""Downloads one Artwork row's original_url and stores it locally
(section 10; the "Artwork processing" background worker from section 1's
architecture diagram).

Deliberately narrow: only ever called for artwork already marked
`is_primary` (see `app/services/ingestion/runner.py:get_primary_artwork_ids`
and `app/workers/tasks.py`, which cascade a download job per primary image
after each ingestion succeeds) — TMDb returns dozens of backdrop/logo
alternates per title, and downloading every one of them automatically would
be a lot of unnecessary storage and bandwidth for images Aurora's own
clients are unlikely to ever show. Downloading on request for anything
beyond the primary image is a reasonable follow-up if that's ever needed.
"""

from __future__ import annotations

import hashlib
import uuid

import httpx
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import get_settings
from app.core.logging import get_logger
from app.models.artwork import Artwork
from app.services.artwork.storage import ArtworkStorage

logger = get_logger(__name__)

_EXTENSION_BY_CONTENT_TYPE = {
    "image/jpeg": "jpg",
    "image/jpg": "jpg",
    "image/png": "png",
    "image/webp": "webp",
    "image/gif": "gif",
    "image/svg+xml": "svg",
}


class ArtworkDownloadError(Exception):
    """Raised for any download-side failure (bad content-type, too large,
    network error, ...). Callers (the arq task) let this propagate so the
    job gets retried like any other ingestion failure — see
    app/services/ingestion/runner.py."""


async def download_artwork(
    session: AsyncSession,
    storage: ArtworkStorage,
    artwork_id: uuid.UUID,
    *,
    client: httpx.AsyncClient | None = None,
    force: bool = False,
) -> bool:
    """Fetch `artwork.original_url`, store it, and set `local_path` +
    `checksum`. Returns False (no-op) if the row is gone or already has a
    local copy (unless `force=True`); True if it downloaded something.
    """

    artwork = await session.get(Artwork, artwork_id)
    if artwork is None:
        return False
    if artwork.local_path and not force:
        return False

    settings = get_settings()
    owns_client = client is None
    http_client = client or httpx.AsyncClient(timeout=20.0, follow_redirects=True)
    try:
        content, content_type = await _fetch(http_client, artwork.original_url, settings.ARTWORK_DOWNLOAD_MAX_BYTES)
    finally:
        if owns_client:
            await http_client.aclose()

    checksum = hashlib.sha256(content).hexdigest()
    extension = _EXTENSION_BY_CONTENT_TYPE.get(content_type, "jpg")
    relative_path = f"{artwork.artwork_type.value}/{checksum[:2]}/{checksum}.{extension}"

    local_url = await storage.save(relative_path, content)

    artwork.local_path = local_url
    artwork.checksum = checksum
    await session.flush()
    logger.info(
        "artwork_downloaded",
        artwork_id=str(artwork_id),
        checksum=checksum,
        bytes=len(content),
    )
    return True


async def _fetch(client: httpx.AsyncClient, url: str, max_bytes: int) -> tuple[bytes, str]:
    try:
        async with client.stream("GET", url) as response:
            response.raise_for_status()
            content_type = response.headers.get("content-type", "").split(";")[0].strip().lower()
            if not content_type.startswith("image/"):
                raise ArtworkDownloadError(f"expected an image, got content-type={content_type!r}")

            body = bytearray()
            async for chunk in response.aiter_bytes():
                body.extend(chunk)
                if len(body) > max_bytes:
                    raise ArtworkDownloadError(
                        f"artwork at {url} exceeds ARTWORK_DOWNLOAD_MAX_BYTES ({max_bytes} bytes)"
                    )
            return bytes(body), content_type
    except httpx.HTTPError as exc:
        raise ArtworkDownloadError(f"failed to download {url}: {exc}") from exc