-
-
Notifications
You must be signed in to change notification settings - Fork 1
/
__init__.py
160 lines (146 loc) · 5.55 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
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")
@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 content:
return Response(
content,
content_type=content_type or "application/octet-stream",
headers={
"content-disposition": 'attachment; filename="{}"'.format(
content_filename
)
}
if content_filename
else None,
)
else:
print(
"""
asgi_send_file(
send={},
filepath={},
filename={},
content_type={} or guess_type(filepath)[0],
)
""".format(
send, filepath, content_filename, content_type
)
)
await asgi_send_file(
send,
filepath,
filename=content_filename,
content_type=content_type or guess_type(filepath)[0],
)