perf(*): bound list queries, batch fetches, harden media downloads

backend/src/userbot/modules/media/downloader.py

@@ -1,14 +1,23 @@
 from typing import Any
 
 from pyrogram import Client
+from pyrogram.errors import (
+    FileIdInvalid,
+    FileReferenceExpired,
+    FileReferenceInvalid,
+    FloodPremiumWait,
+    FloodWait,
+)
 from pyrogram.types import Message
 
 from userbot.modules.capture import repository
 from userbot.modules.capture.context import CaptureContext
-from userbot.modules.download import download_bytes
+from userbot.modules.download import Downloadable, download_bytes
 from utils.logging import logger
 from utils.policy.models import CaptureToggles
 
+STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
+
 _MEDIA_ATTRS = (
     "photo",
     "video",
@@ -48,6 +57,48 @@ def media_unique_id(message: Message) -> str | None:
     return getattr(obj, "file_unique_id", None) if obj is not None else None
 
 
+def _download_target(message: Message) -> Downloadable | None:
+    kind, obj = media_object(message)
+    if obj is None:
+        return None
+    return message if getattr(message, kind or "", None) is obj else obj
+
+
+async def _fresh_target(
+    client: Client, chat_id: int, message_id: int
+) -> Downloadable | None:
+    message = await client.get_messages(chat_id, message_id)
+    if not isinstance(message, Message) or message.empty:
+        return None
+    return _download_target(message)
+
+
+async def _download(
+    client: Client, target: Downloadable | None, chat_id: int, message_id: int
+) -> bytes | None:
+    if target is None:
+        return None
+    try:
+        return await download_bytes(client, target)
+    except TimeoutError:
+        logger.warning(
+            f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
+        )
+        return None
+    except STALE_FILE_ID:
+        fresh = await _fresh_target(client, chat_id, message_id)
+        if fresh is None:
+            return None
+        try:
+            return await download_bytes(client, fresh)
+        except (TimeoutError, *STALE_FILE_ID):
+            logger.warning(
+                f"[yellow]Media download failed after refetch "
+                f"for {chat_id}/{message_id}.[/]"
+            )
+            return None
+
+
 async def capture_media(  # noqa: PLR0913
     client: Client,
     message: Message,
@@ -55,10 +106,10 @@ async def capture_media(  # noqa: PLR0913
     chat_id: int,
     message_id: int,
     toggles: CaptureToggles,
-) -> None:
+) -> bool:
     kind, obj = media_object(message)
     if obj is None:
-        return
+        return True
     unique_id = getattr(obj, "file_unique_id", None)
     ttl = getattr(obj, "ttl_seconds", None)
     want = toggles.self_destruct_media if ttl else toggles.media
@@ -66,6 +117,7 @@ async def capture_media(  # noqa: PLR0913
     mime = getattr(obj, "mime_type", None)
     storage_key: str | None = None
     downloaded = False
+    flood: FloodWait | FloodPremiumWait | None = None
     if want:
         existing = await repository.current_media(
             ctx.pool, ctx.account_id, chat_id, message_id
@@ -80,13 +132,11 @@ async def capture_media(  # noqa: PLR0913
             file_size = existing["file_size"]
             downloaded = True
         else:
-            target = message if getattr(message, kind or "", None) is obj else obj
+            target = _download_target(message)
             try:
-                data = await download_bytes(client, target)
-            except TimeoutError:
-                logger.warning(
-                    f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
-                )
+                data = await _download(client, target, chat_id, message_id)
+            except (FloodWait, FloodPremiumWait) as exc:
+                flood = exc
                 data = None
             if data is not None:
                 storage_key = ctx.storage.put(data)
@@ -105,3 +155,6 @@ async def capture_media(  # noqa: PLR0913
         unique_id,
         downloaded=downloaded,
     )
+    if flood is not None:
+        raise flood
+    return downloaded or not want