Coverage for slidge/db/avatar.py: 89%
225 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-10 12:26 +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, avatar: AvatarType, session: AnySession | None = None
215 ) -> CachedAvatar:
216 if avatar.unique_id is not None:
217 with self.store.session() as orm:
218 stored = (
219 orm.query(Avatar)
220 .filter_by(legacy_id=str(avatar.unique_id))
221 .one_or_none()
222 )
223 if stored is not None:
224 return self.from_stored(stored)
226 if avatar.url is not None:
227 return await self.__fetch_url_if_not_cached(avatar, session=session)
229 return await self.__process(
230 avatar, await self.__open_image(avatar), session=session
231 )
233 async def __fetch_url_if_not_cached(
234 self,
235 avatar: AvatarType,
236 session: AnySession | None = None,
237 convert: bool = True,
238 ) -> CachedAvatar:
239 assert avatar.url is not None
240 async with self.lock(avatar.unique_id or avatar.url):
241 with self.store.session() as orm:
242 if avatar.unique_id is None:
243 stored = orm.query(Avatar).filter_by(url=avatar.url).one_or_none()
244 else:
245 stored = (
246 orm.query(Avatar)
247 .filter_by(legacy_id=str(avatar.unique_id))
248 .one_or_none()
249 )
250 if stored is not None:
251 return self.from_stored(stored)
253 try:
254 response_headers, data = await self.__download_if_modified(
255 avatar.url, self.__get_http_headers(stored)
256 )
257 except NotModified:
258 assert stored is not None
259 return self.from_stored(stored)
261 return await self.__process(
262 avatar,
263 open_image(io.BytesIO(data)),
264 response_headers,
265 session,
266 img_bytes=data,
267 )
269 async def __process(
270 self,
271 avatar: AvatarType,
272 img: Image,
273 response_headers: CIMultiDictProxy[str] | None = None,
274 session: AnySession | None = None,
275 img_bytes: bytes | None = None,
276 ) -> CachedAvatar:
277 too_big = (size := config.AVATAR_SIZE) and any(x > size for x in img.size)
278 if too_big:
279 await asyncio.get_event_loop().run_in_executor(
280 self._thread_pool, img.thumbnail, (size, size)
281 )
282 img_bytes = _get_png_bytes(img)
283 log.debug("Resampled image to %s", img.size)
284 else:
285 img_bytes = (
286 _get_any_bytes(img_bytes, avatar)
287 if img.format == "PNG"
288 else _get_png_bytes(img)
289 )
291 hash_ = hashlib.sha1(img_bytes).hexdigest()
292 http_metadata = await self.__upload(session, avatar, img, hash_, len(img_bytes))
293 file_path = (self.dir / hash_).with_suffix(".png")
294 if file_path.exists():
295 log.warning("Overwriting %s", file_path)
296 with file_path.open("wb") as file:
297 file.write(img_bytes)
298 with self.store.session(expire_on_commit=False) as orm:
299 stored = orm.execute(select(Avatar).where(Avatar.hash == hash_)).scalar()
301 if stored is not None:
302 if (
303 avatar.unique_id is not None
304 and str(avatar.unique_id) != stored.legacy_id
305 ):
306 log.warning(
307 "Updating the 'unique' hash of an avatar, was '%s', is now '%s'",
308 stored.legacy_id,
309 avatar.unique_id,
310 )
311 stored.legacy_id = str(avatar.unique_id)
312 stored.set_http_metadata(http_metadata)
313 orm.add(stored)
314 orm.commit()
316 return self.from_stored(stored)
318 stored = Avatar(
319 hash=hash_,
320 height=img.height,
321 width=img.width,
322 url=avatar.url,
323 legacy_id=avatar.unique_id,
324 )
325 stored.set_http_metadata(http_metadata)
327 if response_headers:
328 stored.etag = response_headers.get("etag")
329 stored.last_modified = response_headers.get("last-modified")
331 with self.store.session(expire_on_commit=False) as orm:
332 if avatar.url is not None:
333 existing = orm.execute(
334 select(Avatar).filter_by(url=avatar.url)
335 ).scalar_one_or_none()
336 if existing is not None:
337 orm.delete(existing)
338 orm.commit()
339 orm.add(stored)
340 orm.commit()
341 return self.from_stored(stored)
343 async def __upload(
344 self,
345 session: AnySession | None,
346 avatar: AvatarType,
347 img: Image,
348 hash_: str,
349 size: int,
350 ) -> AvatarMetadata | None:
351 if not session:
352 return None
353 if config.USE_ATTACHMENT_ORIGINAL_URLS and avatar.url is not None:
354 assert img.format
355 return AvatarMetadata(
356 url=avatar.url,
357 width=img.width,
358 height=img.height,
359 bytes=size,
360 type=img.format.lower(),
361 id=hash_,
362 )
364 return await AttachmentUploader(session.xmpp, session).upload_avatar(
365 avatar, img, hash_
366 )
369def _get_png_bytes(img: Image) -> bytes:
370 if img.mode == "CMYK":
371 # PNG does not support CMYK.
372 img = img.convert("RGB")
373 with io.BytesIO() as f:
374 img.save(f, format="PNG")
375 return f.getvalue()
378def _get_any_bytes(img_bytes: bytes | None, avatar: AvatarType) -> bytes:
379 r = img_bytes or avatar.data
380 if r is None:
381 if avatar.path is None:
382 raise RuntimeError("NEVER")
383 r = avatar.path.read_bytes()
384 return r
387avatar_cache = AvatarCache()
388log = logging.getLogger(__name__)
389_download_lock = asyncio.Lock()
391__all__ = ("AvatarType", "CachedAvatar", "avatar_cache")