feat(userbot,api,frontend): backfill stories and unbreak all-stories list

backend/src/userbot/modules/jobs/handlers/__init__.py

@@ -1,5 +1,6 @@
 from userbot.modules.jobs.handlers import (
     backfill,
+    backfill_stories,
     enrich_chat,
     fetch_avatar,
     fetch_custom_emoji,
@@ -12,6 +13,7 @@ from userbot.modules.jobs.handlers import (
 
 __all__ = [
     "backfill",
+    "backfill_stories",
     "enrich_chat",
     "fetch_avatar",
     "fetch_custom_emoji",

backend/src/userbot/modules/jobs/handlers/backfill_stories.py

@@ -0,0 +1,70 @@
+from collections.abc import AsyncIterator, Callable
+
+from pyrogram import Client
+from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError
+from pyrogram.types import Story
+
+from userbot.modules.capture.context import CaptureContext
+from userbot.modules.jobs.context import JobContext
+from userbot.modules.jobs.registry import register
+from userbot.modules.stories.service import save_story
+
+SAVE_EVERY = 10
+
+StorySource = Callable[[], AsyncIterator[Story]]
+
+
+def _sources(client: Client, peer_id: int, *, own: bool) -> dict[str, StorySource]:
+    sources: dict[str, StorySource] = {
+        "active": lambda: client.get_chat_stories(peer_id),
+        "pinned": lambda: client.get_pinned_stories(peer_id),
+    }
+    if own:
+        sources["archived"] = lambda: client.get_archived_stories(peer_id)
+    return sources
+
+
+async def _drain(
+    ctx: JobContext, capture: CaptureContext, name: str, source: StorySource
+) -> int:
+    client = ctx.client
+    if client is None:
+        return 0
+    saved = 0
+    async for story in source():
+        try:
+            await save_story(client, capture, story)
+        except (FloodWait, FloodPremiumWait):
+            raise
+        except RPCError:
+            continue
+        saved += 1
+        if saved % SAVE_EVERY == 0:
+            await ctx.report_progress({"saved": saved, "source": name})
+            if await ctx.is_canceled():
+                break
+    return saved
+
+
+@register("backfill_stories")
+async def backfill_stories(ctx: JobContext) -> None:
+    client = ctx.client
+    if client is None:
+        return
+    capture = getattr(client, "capture", None)
+    if capture is None:
+        return
+    peer_id = ctx.job.params["peer_id"]
+    own = client.me is not None and client.me.id == peer_id
+    saved = 0
+    errors: dict[str, str] = {}
+    for name, source in _sources(client, peer_id, own=own).items():
+        try:
+            saved += await _drain(ctx, capture, name, source)
+        except (FloodWait, FloodPremiumWait):
+            raise
+        except RPCError as exc:
+            errors[name] = type(exc).__name__
+        if await ctx.is_canceled():
+            break
+    await ctx.report_progress({"saved": saved, "done": True, "errors": errors})

backend/src/userbot/modules/stories/repository.py

@@ -22,6 +22,20 @@ ON CONFLICT (account_id, peer_id, story_id) DO UPDATE SET
 """
 
 
+async def is_downloaded(
+    pool: asyncpg.Pool, account_id: int, peer_id: int, story_id: int
+) -> bool:
+    return bool(
+        await pool.fetchval(
+            "SELECT downloaded FROM stories "
+            "WHERE account_id = $1 AND peer_id = $2 AND story_id = $3",
+            account_id,
+            peer_id,
+            story_id,
+        )
+    )
+
+
 async def upsert_story(  # noqa: PLR0913
     pool: asyncpg.Pool,
     account_id: int,

backend/src/userbot/modules/stories/service.py

@@ -0,0 +1,49 @@
+from io import BytesIO
+
+from pyrogram import Client
+from pyrogram.types import Story
+
+from userbot.modules.capture.context import CaptureContext
+from userbot.modules.stories import repository
+
+
+def story_peer_id(story: Story) -> int:
+    if story.chat is not None:
+        return story.chat.id or 0
+    if story.from_user is not None:
+        return story.from_user.id or 0
+    return 0
+
+
+async def save_story(client: Client, capture: CaptureContext, story: Story) -> None:
+    peer_id = story_peer_id(story)
+    storage_key: str | None = None
+    file_size: int | None = None
+    downloaded = False
+    stored = await repository.is_downloaded(
+        capture.pool, capture.account_id, peer_id, story.id
+    )
+    if not (stored or story.deleted or story.media is None):
+        buffer = await client.download_media(story, in_memory=True)
+        if isinstance(buffer, BytesIO):
+            data = buffer.getvalue()
+            storage_key = capture.storage.put(data)
+            file_size = len(data)
+            downloaded = True
+    await repository.upsert_story(
+        capture.pool,
+        capture.account_id,
+        peer_id,
+        story.id,
+        story.date,
+        story.expire_date,
+        story.caption,
+        story.media.name.lower() if story.media else None,
+        storage_key,
+        file_size,
+        story.views,
+        str(story.raw),
+        pinned=bool(story.pinned),
+        deleted=bool(story.deleted),
+        downloaded=downloaded,
+    )