datasette-media/datasette_media/__init__.py

import asyncio
from datasette import hookimpl
from datasette.utils.asgi import Response, asgi_send_file
from concurrent import futures
import httpx
from mimetypes import guess_type
from PIL import Image
import io
from . import utils
 
transform_executor = None
RESERVED_MEDIA_TYPES = ("transform_threads", "enable_transform")
 
PNG_1x1 = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x00\x00\x00\x00:~\x9bU\x00\x00\x00\nIDATx\x9cc\xfa\x0f\x00\x01\x05\x01\x02\xcf\xa0.\xcd\x00\x00\x00\x00IEND\xaeB`\x82"
 
 
@hookimpl
def register_routes():
    return [
        (r"/-/media/(?P<media_type>[^/]+)/(?P<key>.+)", serve_media),
    ]
 
 
async def serve_media(datasette, request, send):
    global transform_executor
    plugin_config = datasette.plugin_config("datasette-media") or {}
    pool_size = plugin_config.get("transform_threads") or 4
    if transform_executor is None:
        transform_executor = futures.ThreadPoolExecutor(max_workers=pool_size)
 
    media_type = request.url_vars["media_type"]
    key = request.url_vars["key"]
 
    config = plugin_config.get(media_type)
    if media_type in RESERVED_MEDIA_TYPES or config is None:
        return Response.html("<h1>Invalid media type</h1>", status=404)
    sql = config.get("sql")
    if sql is None:
        return Response.html("<h1>Missing SQL from configuration</h1>", status=404)
    database = datasette.get_database(config.get("database"))
    results = await database.execute(sql, {"key": key})
    row = results.first()
    if row is None:
        return Response.html("<h1>404 - no results</h1>", status=404)
 
    # We need filepath or content
    content = None
    content_type = None
    content_url = None
    content_filename = None
    filepath = None
 
    row_keys = row.keys()
    if (
        "filepath" not in row_keys
        and "content" not in row_keys
        and "content_url" not in row_keys
    ):
        return Response.html(
            "<h1>404 - SQL must return 'filepath' or 'content' or 'content_url'</h1>",
            status=404,
        )
    if "content" in row_keys:
        content = row["content"]
    elif "content_url" in row_keys:
        content_url = row["content_url"]
    else:
        filepath = row["filepath"]
 
    if "content_filename" in row_keys:
        content_filename = row["content_filename"]
 
    # Images are special cases, triggered by a few different conditions
    should_transform = utils.should_transform(row, config, request)
    if should_transform:
        if content is None and content_url:
            async with httpx.AsyncClient() as client:
                response = await client.get(row["content_url"])
                content = response.content
                content_type = response.headers["content-type"]
        image_bytes = content or open(filepath, "rb").read()
        image = await asyncio.get_event_loop().run_in_executor(
            transform_executor,
            lambda: utils.transform_image(image_bytes, **should_transform),
        )
        response = utils.ImageResponse(image, format=should_transform.get("format"))
        if content_filename:
            response.headers[
                "content-disposition"
            ] = 'attachment; filename="{}"'.format(content_filename)
        return response
    else:
        # content_url is proxied as a special case
        if content_url:
            client = httpx.AsyncClient()
            async with client.stream("GET", content_url) as response:
                content_type = response.headers["content-type"]
                content_length = response.headers.get("content-length")
                headers = [(b"content-type", content_type.encode("utf-8"))]
                if content_length:
                    headers.append(
                        (b"content-length", str(content_length).encode("utf-8"))
                    )
                if content_filename:
                    headers.append(
                        (
                            b"content-disposition",
                            'attachment; filename="{}"'.format(content_filename).encode(
                                "utf-8"
                            ),
                        )
                    )
 
                await send(
                    {"type": "http.response.start", "status": 200, "headers": headers}
                )
                async for chunk in response.aiter_bytes():
                    await send(
                        {
                            "type": "http.response.body",
                            "body": chunk,
                            "more_body": True,
                        }
                    )
                await send({"type": "http.response.body", "body": b""})
                return
 
        if "content_type" in row_keys:
            content_type = row["content_type"]
 
        # Non-image files are returned directly
        if filepath:
            await asgi_send_file(
                send,
                filepath,
                filename=content_filename,
                content_type=content_type or guess_type(filepath)[0],
            )
        else:
            status_code = 200
            if not content:
                status_code = 404
                content = PNG_1x1
                content_type = "application/png"
                content_filename = None
            return Response(
                content,
                content_type=content_type or "application/octet-stream",
                headers={
                    "content-disposition": 'attachment; filename="{}"'.format(
                        content_filename
                    )
                }
                if content_filename
                else None,
                status=status_code,
            )