Coverage for slidge/db/avatar.py: 88%
226 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-10 04:45 +0000
1import asyncio
2import hashlib
3import io
4import logging
5from collections.abc import AsyncIterator
6from concurrent.futures import ThreadPoolExecutor
7from contextlib import asynccontextmanager
8from http import HTTPStatus
9from pathlib import Path
10from typing import Literal
12import aiohttp
13from multidict import CIMultiDictProxy
14from PIL.Image import Image
15from PIL.Image import open as open_image
16from sqlalchemy import select
18from ..core import config
19from ..core.attachment_upload import AttachmentUploader
20from ..util.lock import NamedLockMixin
21from ..util.types import AnySession, AvatarMetadata
22from ..util.types import Avatar as AvatarType
23from .models import Avatar
24from .store import AvatarStore
26_AVATAR_FETCH_CONCURRENCY = 8
27_AVATAR_DOWNLOAD_TIMEOUT = 30
30class CachedAvatar:
31 def __init__(self, stored: Avatar, root_dir: Path) -> None:
32 self.stored = stored
33 self._root = root_dir
35 @property
36 def pk(self) -> int | None:
37 return self.stored.id
39 @property
40 def hash(self) -> str | None:
41 """
42 SHA1 of avatar in PNG format, stored in self._root, meant to be
43 served in-band. `None` if the avatar is only available through HTTP.
44 """
45 return self.stored.hash
47 @property
48 def height(self) -> int:
49 return self.stored.height
51 @property
52 def width(self) -> int:
53 return self.stored.width
55 @property
56 def etag(self) -> str | None:
57 return self.stored.etag
59 @property
60 def last_modified(self) -> str | None:
61 return self.stored.last_modified
63 @property
64 def data(self) -> bytes:
65 return self.path.read_bytes()
67 @property
68 def path(self) -> Path:
69 assert self.hash
70 return (self._root / self.hash).with_suffix(".png")
73class NotModified(Exception):
74 pass
77class AvatarDownloadError(Exception):
78 def __init__(self, url: str, status: int | None = None, message: str = "") -> None:
79 self.url = url
80 self.status = status
81 self.message = message
82 if status is None:
83 super().__init__(f"{message} ({url})" if message else url)
84 else:
85 super().__init__(f"{status} {message} ({url})")
88class AvatarCache(NamedLockMixin):
89 dir: Path
90 http: aiohttp.ClientSession
91 store: AvatarStore
93 def __init__(self) -> None:
94 self._thread_pool = ThreadPoolExecutor(config.AVATAR_RESAMPLING_THREADS)
95 self._download_semaphore = asyncio.BoundedSemaphore(_AVATAR_FETCH_CONCURRENCY)
96 super().__init__()
98 def from_stored(self, stored: Avatar) -> CachedAvatar:
99 return CachedAvatar(stored, self.dir)
101 def set_dir(self, path: Path) -> None:
102 self.dir = path
103 self.dir.mkdir(exist_ok=True)
104 for f in path.glob("*"):
105 if f.suffix != ".png":
106 # FIXME: remove this before 1.0.0!
107 # slidge v0.5.0 wrote useless non-PNG files in here, this
108 # cleans them up
109 f.unlink()
110 log.debug("Checking avatar files")
111 with self.store.session(expire_on_commit=False) as orm:
112 for stored in orm.query(Avatar).all():
113 avatar = CachedAvatar(stored, path)
114 if avatar.hash is None or (avatar.hash and avatar.path.exists()):
115 continue
116 log.warning(
117 "Removing avatar %s from store because %s does not exist",
118 avatar.hash,
119 avatar.path,
120 )
121 orm.delete(stored)
122 orm.commit()
124 def close(self) -> None:
125 self._thread_pool.shutdown(cancel_futures=True)
127 def __get_http_headers(
128 self, cached: CachedAvatar | Avatar | None = None
129 ) -> dict[str, str]:
130 headers = {}
131 if (
132 cached
133 and cached.hash
134 and (self.dir / cached.hash).with_suffix(".png").exists()
135 ):
136 if last_modified := cached.last_modified:
137 headers["If-Modified-Since"] = last_modified
138 if etag := cached.etag:
139 headers["If-None-Match"] = etag
140 return headers
142 @asynccontextmanager
143 async def _fetch(
144 self,
145 url: str,
146 headers: dict[str, str],
147 method: Literal["GET", "HEAD"],
148 ) -> AsyncIterator[aiohttp.ClientResponse]:
149 """Common bits of logic for HTTP requests fetching avatar (meta)data."""
150 async with self._download_semaphore:
151 try:
152 async with self.http.request(
153 method,
154 url,
155 headers=headers,
156 timeout=aiohttp.ClientTimeout(total=_AVATAR_DOWNLOAD_TIMEOUT),
157 ) as response:
158 yield response
159 except aiohttp.ClientResponseError as e:
160 raise AvatarDownloadError(url, e.status, e.message) from e
161 except (aiohttp.ClientError, TimeoutError) as e:
162 raise AvatarDownloadError(url, message=str(e)) from e
164 async def __download_if_modified(
165 self,
166 url: str,
167 headers: dict[str, str],
168 ) -> tuple[CIMultiDictProxy[str], bytes]:
169 """
170 Download avatar only if it has been modified compared to what we have
171 in cache.
173 :return: HTTP response headers, data
174 :raise: NotModified if fetching was not necessary
175 """
176 async with self._fetch(url, headers, "GET") as response:
177 if response.status == HTTPStatus.NOT_MODIFIED:
178 log.debug("Using avatar cache for %s", url)
179 raise NotModified
180 response.raise_for_status()
181 data = await response.read()
182 return response.headers, data
184 async def __is_modified(self, url: str, headers: dict[str, str]) -> bool:
185 async with self._fetch(url, headers, "HEAD") as response:
186 response.raise_for_status()
187 if response.status == HTTPStatus.NOT_MODIFIED:
188 return False
189 cached_last_modified = headers.get("If-Modified-Since")
190 if not cached_last_modified:
191 return True
192 response_last_modified = response.headers.get("last-modified")
193 if not response_last_modified:
194 return True
195 return cached_last_modified != response_last_modified
197 async def url_modified(self, url: str) -> bool:
198 with self.store.session() as orm:
199 cached = orm.query(Avatar).filter_by(url=url).one_or_none()
200 if cached is None:
201 return True
202 headers = self.__get_http_headers(cached)
203 return await self.__is_modified(url, headers)
205 @staticmethod
206 async def __open_image(avatar: AvatarType) -> Image:
207 if avatar.data is not None:
208 return open_image(io.BytesIO(avatar.data))
209 elif avatar.path is not None:
210 return open_image(avatar.path)
211 raise TypeError("Avatar must be bytes or a Path", avatar)
213 async def get(
214 self,
215 avatar: AvatarType,
216 session: AnySession | None = None,
217 convert: bool = True,
218 ) -> CachedAvatar:
219 if avatar.unique_id is not None:
220 with self.store.session() as orm:
221 stored = (
222 orm.query(Avatar)
223 .filter_by(legacy_id=str(avatar.unique_id))
224 .one_or_none()
225 )
226 if stored is not None:
227 return self.from_stored(stored)
229 if avatar.url is not None:
230 return await self.__fetch_url_if_not_cached(
231 avatar, session=session, convert=convert
232 )
234 return await self.__process(
235 avatar, await self.__open_image(avatar), session=session, convert=convert
236 )
238 async def __fetch_url_if_not_cached(
239 self,
240 avatar: AvatarType,
241 session: AnySession | None = None,
242 convert: bool = True,
243 ) -> CachedAvatar:
244 assert avatar.url is not None
245 async with self.lock(avatar.unique_id or avatar.url):
246 with self.store.session() as orm:
247 if avatar.unique_id is None:
248 stored = orm.query(Avatar).filter_by(url=avatar.url).one_or_none()
249 else:
250 stored = (
251 orm.query(Avatar)
252 .filter_by(legacy_id=str(avatar.unique_id))
253 .one_or_none()
254 )
255 if stored is not None:
256 return self.from_stored(stored)
258 try:
259 response_headers, data = await self.__download_if_modified(
260 avatar.url, self.__get_http_headers(stored)
261 )
262 except NotModified:
263 assert stored is not None
264 return self.from_stored(stored)
266 return await self.__process(
267 avatar,
268 open_image(io.BytesIO(data)),
269 response_headers,
270 session,
271 convert=convert,
272 img_bytes=data,
273 )
275 async def __process(
276 self,
277 avatar: AvatarType,
278 img: Image,
279 response_headers: CIMultiDictProxy[str] | None = None,
280 session: AnySession | None = None,
281 convert: bool = True,
282 img_bytes: bytes | None = None,
283 ) -> CachedAvatar:
284 if convert:
285 too_big = (size := config.AVATAR_SIZE) and any(x > size for x in img.size)
286 if too_big:
287 await asyncio.get_event_loop().run_in_executor(
288 self._thread_pool, img.thumbnail, (size, size)
289 )
290 img_bytes = _get_png_bytes(img)
291 log.debug("Resampled image to %s", img.size)
292 else:
293 img_bytes = (
294 _get_any_bytes(img_bytes, avatar)
295 if img.format == "PNG"
296 else _get_png_bytes(img)
297 )
298 else:
299 img_bytes = _get_any_bytes(img_bytes, avatar)
301 hash_ = hashlib.sha1(img_bytes).hexdigest()
302 http_metadata = await self.__upload(session, avatar, img, hash_, len(img_bytes))
303 if convert:
304 # convert means that the avatar must be converted in order to be
305 # served in band.
306 # currently, it is only False for space avatars (that have no protocol
307 # for in-band serving).
308 file_path = (self.dir / hash_).with_suffix(".png")
309 if file_path.exists():
310 log.warning("Overwriting %s", file_path)
311 with file_path.open("wb") as file:
312 file.write(img_bytes)
313 with self.store.session(expire_on_commit=False) as orm:
314 stored = orm.execute(
315 select(Avatar).where(Avatar.hash == hash_)
316 ).scalar()
318 if stored is not None:
319 if (
320 avatar.unique_id is not None
321 and str(avatar.unique_id) != stored.legacy_id
322 ):
323 log.warning(
324 "Updating the 'unique' hash of an avatar, was '%s', is now '%s'",
325 stored.legacy_id,
326 avatar.unique_id,
327 )
328 stored.legacy_id = str(avatar.unique_id)
329 stored.set_http_metadata(http_metadata)
330 orm.add(stored)
331 orm.commit()
333 return self.from_stored(stored)
335 stored = Avatar(
336 hash=hash_ if convert else None,
337 height=img.height,
338 width=img.width,
339 url=avatar.url,
340 legacy_id=avatar.unique_id,
341 )
342 stored.set_http_metadata(http_metadata)
344 if response_headers:
345 stored.etag = response_headers.get("etag")
346 stored.last_modified = response_headers.get("last-modified")
348 with self.store.session(expire_on_commit=False) as orm:
349 if avatar.url is not None:
350 existing = orm.execute(
351 select(Avatar).filter_by(url=avatar.url)
352 ).scalar_one_or_none()
353 if existing is not None:
354 orm.delete(existing)
355 orm.commit()
356 orm.add(stored)
357 orm.commit()
358 return self.from_stored(stored)
360 async def __upload(
361 self,
362 session: AnySession | None,
363 avatar: AvatarType,
364 img: Image,
365 hash_: str,
366 size: int,
367 ) -> AvatarMetadata | None:
368 if not session:
369 return None
370 if config.USE_ATTACHMENT_ORIGINAL_URLS and avatar.url is not None:
371 assert img.format
372 return AvatarMetadata(
373 url=avatar.url,
374 width=img.width,
375 height=img.height,
376 bytes=size,
377 type=img.format.lower(),
378 id=hash_,
379 )
381 return await AttachmentUploader(session.xmpp, session).upload_avatar(
382 avatar, img, hash_
383 )
386def _get_png_bytes(img: Image) -> bytes:
387 with io.BytesIO() as f:
388 img.save(f, format="PNG")
389 return f.getvalue()
392def _get_any_bytes(img_bytes: bytes | None, avatar: AvatarType) -> bytes:
393 r = img_bytes or avatar.data
394 if r is None:
395 if avatar.path is None:
396 raise RuntimeError("NEVER")
397 r = avatar.path.read_bytes()
398 return r
401avatar_cache = AvatarCache()
402log = logging.getLogger(__name__)
403_download_lock = asyncio.Lock()
405__all__ = ("AvatarType", "CachedAvatar", "avatar_cache")