beaver/beavergram
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
hh
2026-09-01 21:38:22 +02:00
parent 9c12628304
commit 1fa4f18a9b
35 changed files with 473 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/userbot/modules/capture/repository.py
+1 -1
View File
@@ -122,7 +122,7 @@ async def max_message_id(
pool: asyncpg.Pool, account_id: int, chat_id: int
) -> int | None:
return await pool.fetchval(
"SELECT max(message_id) FROM messages WHERE account_id = $1 AND chat_id = $2",
"SELECT last_message_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2",
account_id,
chat_id,
)
backend/src/userbot/modules/client.py
+1
View File
@@ -31,6 +31,7 @@ class PyroClient(Client):
lang_code="en",
system_lang_code="en-US",
client_platform=enums.ClientPlatform.DESKTOP,
max_concurrent_transmissions=4,
)
self.capture: CaptureContext | None = None
backend/src/userbot/modules/download.py
+1 -1
View File
@@ -50,7 +50,7 @@ async def _reset_media_sessions(client: Client) -> None:
async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
try:
buffer = await client.download_media(target, in_memory=True)
except TimeoutError:
except (TimeoutError, OSError):
await _reset_media_sessions(client)
buffer = await client.download_media(target, in_memory=True)
return buffer.getvalue() if isinstance(buffer, BytesIO) else None
backend/src/userbot/modules/jobs/consumer.py
+3
View File
@@ -69,6 +69,9 @@ class JobConsumer:
requeued = await repository.requeue_running(self.pool, self.account_id)
if requeued:
logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
pruned = await repository.prune_finished(self.pool, self.account_id)
if pruned:
logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]")
conn = await asyncpg.connect(dsn=env.db.connection_url)
await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set())
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
backend/src/userbot/modules/jobs/handlers/backfill.py
+15 -5
View File
@@ -1,5 +1,5 @@
from pyrogram import Client
from pyrogram.errors import PeerIdInvalid
from pyrogram.errors import FloodPremiumWait, FloodWait, PeerIdInvalid
from pyrogram.types import Message
from userbot.modules.capture import capture_message
@@ -11,6 +11,7 @@ from userbot.modules.jobs.registry import register
from userbot.modules.stt import repository as stt_repo
from userbot.modules.stt import should_transcribe_on_backfill
from userbot.modules.stt.gate import safe_transcribe
from utils.logging import logger
from utils.policy.models import CaptureToggles
SAVE_EVERY = 100
@@ -61,21 +62,30 @@ async def backfill(ctx: JobContext) -> None:
min_id = await resolve_min_id(ctx, chat_id)
await ctx.save_cursor({"max_id": max_id, "min_id": min_id})
processed = ctx.job.progress.get("processed", 0)
errors = ctx.job.progress.get("errors", 0)
self_id = client.me.id if client.me else None
try:
async for message in client.get_chat_history(
chat_id, max_id=max_id, min_id=min_id
):
await capture_message(client, message, capture, toggles)
await maybe_transcribe(client, capture, chat_id, message, self_id)
try:
await capture_message(client, message, capture, toggles)
await maybe_transcribe(client, capture, chat_id, message, self_id)
except (FloodWait, FloodPremiumWait, PeerIdInvalid):
raise
except Exception:
errors += 1
logger.exception(f"Backfill capture failed for {chat_id}/{message.id}")
processed += 1
if processed % SAVE_EVERY == 0:
next_max = message.id - 1
await ctx.save_cursor({"max_id": next_max, "min_id": min_id})
await ctx.report_progress({"processed": processed, "max_id": next_max})
await ctx.report_progress(
{"processed": processed, "errors": errors, "max_id": next_max}
)
if await ctx.is_canceled():
return
except PeerIdInvalid:
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
return
await ctx.report_progress({"processed": processed, "done": True})
await ctx.report_progress({"processed": processed, "errors": errors, "done": True})
backend/src/userbot/modules/jobs/handlers/enrich_chat.py
+19 -7
View File
@@ -65,16 +65,28 @@ async def _enrich_members(client: Client, ctx: CaptureContext, chat_id: int) ->
return
async def _fetch_users(client: Client, ids: list[int]) -> list[User]:
try:
users = await client.get_users(ids)
except BadRequest:
users = []
for sender_id in ids:
try:
users.append(await client.get_users(sender_id))
except BadRequest:
continue
if isinstance(users, User):
return [users]
return [user for user in users if isinstance(user, User)]
async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None:
rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id)
ids = [row["sender_id"] for row in rows]
for sender_id in ids:
try:
user = await client.get_users(sender_id)
except BadRequest:
continue
if isinstance(user, User):
await _save_user(ctx, user)
if not ids:
return
for user in await _fetch_users(client, ids):
await _save_user(ctx, user)
@register("enrich_chat")
backend/src/userbot/modules/jobs/handlers/fetch_media.py
+6 -2
View File
@@ -18,6 +18,10 @@ async def fetch_media(ctx: JobContext) -> None:
if isinstance(message, list):
message = message[0] if message else None
if message is None or message.empty:
return
msg = f"message {chat_id}/{message_id} is unavailable"
raise RuntimeError(msg)
toggles = CaptureToggles(media=True, self_destruct_media=True)
await capture_media(client, message, capture, chat_id, message_id, toggles)
ok = await capture_media(client, message, capture, chat_id, message_id, toggles)
if not ok:
msg = f"media download failed for {chat_id}/{message_id}"
raise RuntimeError(msg)
backend/src/userbot/modules/jobs/repository.py
+10
View File
@@ -48,6 +48,16 @@ async def requeue_running(pool: asyncpg.Pool, account_id: int) -> int:
return int(result.split()[-1])
async def prune_finished(pool: asyncpg.Pool, account_id: int) -> int:
result = await pool.execute(
"DELETE FROM jobs WHERE account_id = $1 "
"AND status IN ('done', 'canceled', 'failed') "
"AND finished_at < now() - interval '30 days'",
account_id,
)
return int(result.split()[-1])
async def save_cursor(pool: asyncpg.Pool, job_id: int, cursor: dict[str, Any]) -> None:
await pool.execute(
"UPDATE jobs SET cursor = $2::jsonb, updated_at = now() WHERE id = $1",
backend/src/userbot/modules/media/downloader.py
+62 -9
View File
@@ -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
backend/src/userbot/runner.py
+2 -14
View File
@@ -49,18 +49,6 @@ async def _cancel(task: asyncio.Task) -> None:
await task
async def _enqueue_once(pool: asyncpg.Pool, account_id: int, kind: str) -> None:
existing = await pool.fetchval(
"SELECT 1 FROM jobs WHERE account_id = $1 AND kind = $2 "
"AND status IN ('pending', 'running') LIMIT 1",
account_id,
kind,
)
if existing is None:
await enqueue(pool, account_id, kind, {})
logger.info(f"[green]Queued {kind}.[/]")
class AccountRegistry:
def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
self._pool = pool
@@ -121,8 +109,8 @@ class AccountRegistry:
client, asyncio.create_task(consumer.run()), device_model
)
logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
await _enqueue_once(self._pool, account_id, "sync_dialogs")
await _enqueue_once(self._pool, account_id, "sync_contacts")
await enqueue(self._pool, account_id, "sync_dialogs", {})
await enqueue(self._pool, account_id, "sync_contacts", {})
async def _stop(self, session_name: str) -> None:
account = self._running.pop(session_name, None)