diff --git a/backend/src/api/routers/backfill.py b/backend/src/api/routers/backfill.py index 87304ac..5cf8306 100644 --- a/backend/src/api/routers/backfill.py +++ b/backend/src/api/routers/backfill.py @@ -19,6 +19,11 @@ class BackfillRequest(BaseModel): full: bool = False +class StoriesBackfillRequest(BaseModel): + account_id: int + peer_id: int + + class FetchMediaRequest(BaseModel): account_id: int chat_id: int @@ -76,6 +81,16 @@ async def enqueue_backfill( return EnqueueResponse(job_id=job_id) +@router.post("/stories/backfill", status_code=201) +async def enqueue_stories_backfill( + pool: FromDishka[asyncpg.Pool], body: StoriesBackfillRequest +) -> EnqueueResponse: + job_id = await enqueue( + pool, body.account_id, "backfill_stories", {"peer_id": body.peer_id} + ) + return EnqueueResponse(job_id=job_id) + + @router.post("/media/fetch", status_code=201) async def enqueue_fetch_media( pool: FromDishka[asyncpg.Pool], body: FetchMediaRequest diff --git a/backend/src/userbot/handlers/stories.py b/backend/src/userbot/handlers/stories.py index 774673c..c0d9ccf 100644 --- a/backend/src/userbot/handlers/stories.py +++ b/backend/src/userbot/handlers/stories.py @@ -1,52 +1,14 @@ -from io import BytesIO - from pyrogram.types import Story from userbot import PyroClient -from userbot.modules.stories import repository - - -def _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 +from userbot.modules.stories.service import save_story @PyroClient.on_story() async def on_story(client: PyroClient, story: Story) -> None: - ctx = client.capture - if ctx is None: + if client.capture is None: return - media_kind = story.media.name.lower() if story.media else None - storage_key: str | None = None - file_size: int | None = None - downloaded = False - if not story.deleted and story.media is not None: - buffer = await client.download_media(story, in_memory=True) - if isinstance(buffer, BytesIO): - data = buffer.getvalue() - storage_key = ctx.storage.put(data) - file_size = len(data) - downloaded = True - await repository.upsert_story( - ctx.pool, - ctx.account_id, - _peer_id(story), - story.id, - story.date, - story.expire_date, - story.caption, - media_kind, - storage_key, - file_size, - story.views, - str(story.raw), - pinned=bool(story.pinned), - deleted=bool(story.deleted), - downloaded=downloaded, - ) + await save_story(client, client.capture, story) handlers = on_story.handlers diff --git a/backend/src/userbot/modules/jobs/handlers/__init__.py b/backend/src/userbot/modules/jobs/handlers/__init__.py index 0fd736f..05b1dc4 100644 --- a/backend/src/userbot/modules/jobs/handlers/__init__.py +++ b/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", diff --git a/backend/src/userbot/modules/jobs/handlers/backfill_stories.py b/backend/src/userbot/modules/jobs/handlers/backfill_stories.py new file mode 100644 index 0000000..f292565 --- /dev/null +++ b/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}) diff --git a/backend/src/userbot/modules/stories/repository.py b/backend/src/userbot/modules/stories/repository.py index 7007207..06a72de 100644 --- a/backend/src/userbot/modules/stories/repository.py +++ b/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, diff --git a/backend/src/userbot/modules/stories/service.py b/backend/src/userbot/modules/stories/service.py new file mode 100644 index 0000000..5dff4c0 --- /dev/null +++ b/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, + ) diff --git a/frontend/src/lib/api/endpoints.ts b/frontend/src/lib/api/endpoints.ts index 8580586..3e1f727 100644 --- a/frontend/src/lib/api/endpoints.ts +++ b/frontend/src/lib/api/endpoints.ts @@ -315,7 +315,7 @@ export function getMessageAt(chatId: number, date: string): Promise { } export function getStories( - peerId: number, + peerId: number | null, page: Page = {} ): Promise { return request("/stories", { @@ -367,6 +367,15 @@ export function enqueueBackfill( }); } +export function enqueueStoriesBackfill( + peerId: number +): Promise<{ job_id: number }> { + return request<{ job_id: number }>("/stories/backfill", { + method: "POST", + body: { account_id: accounts.selectedId, peer_id: peerId }, + }); +} + export function discoverPeers( query: string, remote = false diff --git a/frontend/src/lib/components/jobs/JobList.svelte b/frontend/src/lib/components/jobs/JobList.svelte index 3f60e1b..c35f589 100644 --- a/frontend/src/lib/components/jobs/JobList.svelte +++ b/frontend/src/lib/components/jobs/JobList.svelte @@ -16,6 +16,7 @@ const KIND_LABELS: Record = { backfill: "Бэкфилл", + backfill_stories: "Бэкфилл сторис", fetch_media: "Докачка медиа", fetch_avatar: "Аватар", fetch_custom_emoji: "Кастом-эмодзи", @@ -88,12 +89,12 @@ } function processed(job: JobView): number | null { - const value = job.progress.processed; + const value = job.progress.processed ?? job.progress.saved; return typeof value === "number" ? value : null; } function chatId(job: JobView): number | null { - const value = job.params.chat_id; + const value = job.params.chat_id ?? job.params.peer_id; return typeof value === "number" ? value : null; } diff --git a/frontend/src/lib/components/profile/ProfilePanel.svelte b/frontend/src/lib/components/profile/ProfilePanel.svelte index 76f5171..11eaf6f 100644 --- a/frontend/src/lib/components/profile/ProfilePanel.svelte +++ b/frontend/src/lib/components/profile/ProfilePanel.svelte @@ -6,19 +6,21 @@ import ProfileInfo from "$lib/components/profile/ProfileInfo.svelte"; import SharedLinks from "$lib/components/profile/SharedLinks.svelte"; import SharedMedia from "$lib/components/profile/SharedMedia.svelte"; + import StoriesArchive from "$lib/components/stories/StoriesArchive.svelte"; import Avatar from "$lib/components/ui/Avatar.svelte"; import EmptyState from "$lib/components/ui/EmptyState.svelte"; import Icon from "$lib/components/ui/Icon.svelte"; import { peerName } from "$lib/format/peer"; import { chats } from "$lib/stores/chats.svelte"; - type Tab = "info" | "media" | "files" | "links" | "calendar"; + type Tab = "info" | "media" | "files" | "links" | "stories" | "calendar"; const TABS: { id: Tab; icon: string; label: string }[] = [ { id: "info", icon: "info", label: "Инфо" }, { id: "media", icon: "photo", label: "Медиа" }, { id: "files", icon: "document", label: "Файлы" }, { id: "links", icon: "link", label: "Ссылки" }, + { id: "stories", icon: "play-story", label: "Сторис" }, { id: "calendar", icon: "calendar", label: "Календарь" }, ]; @@ -114,6 +116,8 @@ {:else if tab === "links"} + {:else if tab === "stories"} + {:else if tab === "calendar"} {/if} diff --git a/frontend/src/lib/components/stories/AllStoriesArchive.svelte b/frontend/src/lib/components/stories/AllStoriesArchive.svelte index 52af996..3ff42dd 100644 --- a/frontend/src/lib/components/stories/AllStoriesArchive.svelte +++ b/frontend/src/lib/components/stories/AllStoriesArchive.svelte @@ -1,16 +1,23 @@ {#if peerId === null} -{:else if items.length === 0} - {#if loading} -
- {:else} - - {/if} {:else} -
- {#each items as item, index (item.story_id)} - - {/each} +
+ +
- {#if !done} - - {/if} + {#if items.length === 0} + {#if loading} +
+ {:else} + + {/if} + {:else} +
+ {#each items as item, index (item.story_id)} + + {/each} +
+ + {#if !done} + + {/if} - {#if peerId !== null}