diff --git a/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py b/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py index 78bc91d..bc7ed0d 100644 --- a/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py +++ b/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py @@ -1,9 +1,42 @@ from io import BytesIO +from pyrogram import Client +from pyrogram.errors import FileIdInvalid, FileReferenceExpired, FileReferenceInvalid +from pyrogram.types import Photo + from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.registry import register +STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid) + + +async def _fresh_file_id(client: Client, owner_id: int, unique_id: str) -> str | None: + chat = await client.get_chat(owner_id) + photo = chat.photo + if photo and photo.big_photo_unique_id == unique_id: + return photo.big_file_id + photos = client.get_chat_photos(owner_id) + if photos is None: + return None + async for item in photos: + if isinstance(item, Photo) and item.file_unique_id == unique_id: + return item.file_id + return None + + +async def _download( + client: Client, owner_id: int, unique_id: str, file_id: str +) -> bytes | None: + try: + buffer = await client.download_media(file_id, in_memory=True) + except STALE_FILE_ID: + fresh = await _fresh_file_id(client, owner_id, unique_id) + if fresh is None: + return None + buffer = await client.download_media(fresh, in_memory=True) + return buffer.getvalue() if isinstance(buffer, BytesIO) else None + @register("fetch_avatar") async def fetch_avatar(ctx: JobContext) -> None: @@ -24,10 +57,9 @@ async def fetch_avatar(ctx: JobContext) -> None: file_id, downloaded = found if downloaded or file_id is None: return - buffer = await client.download_media(file_id, in_memory=True) - if not isinstance(buffer, BytesIO): + data = await _download(client, owner_id, unique_id, file_id) + if data is None: return - data = buffer.getvalue() storage_key = capture.storage.put(data) await mark_avatar_downloaded( ctx.pool,