64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
from pyrogram import Client
|
|
from pyrogram.errors import BadRequest, Forbidden
|
|
from pyrogram.types import User
|
|
|
|
from userbot.modules.capture.context import CaptureContext
|
|
from userbot.modules.jobs.context import JobContext
|
|
from userbot.modules.jobs.registry import register
|
|
from userbot.modules.profiles.snapshots import save_group, save_private, save_user
|
|
from utils.logging import logger
|
|
|
|
SAVE_EVERY = 100
|
|
USERS_BATCH = 200
|
|
|
|
_UPSERT_DIALOG = """
|
|
INSERT INTO dialogs (account_id, chat_id) VALUES ($1, $2)
|
|
ON CONFLICT (account_id, chat_id) DO UPDATE SET updated_at = now()
|
|
"""
|
|
|
|
|
|
async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) -> None:
|
|
for start in range(0, len(ids), USERS_BATCH):
|
|
batch = ids[start : start + USERS_BATCH]
|
|
try:
|
|
result = await client.get_users(batch)
|
|
except (BadRequest, Forbidden):
|
|
continue
|
|
users = result if isinstance(result, list) else [result]
|
|
for user in users:
|
|
if isinstance(user, User):
|
|
await save_user(ctx, user)
|
|
|
|
|
|
@register("sync_dialogs")
|
|
async def sync_dialogs(ctx: JobContext) -> None:
|
|
client = ctx.client
|
|
if client is None:
|
|
return
|
|
capture = getattr(client, "capture", None)
|
|
if capture is None:
|
|
return
|
|
processed = ctx.job.progress.get("processed", 0)
|
|
nameless: list[int] = []
|
|
async for dialog in client.get_dialogs():
|
|
chat = dialog.chat
|
|
if chat is None or chat.id is None:
|
|
continue
|
|
chat_id = chat.id
|
|
try:
|
|
if chat_id > 0:
|
|
if not await save_private(capture, chat):
|
|
nameless.append(chat_id)
|
|
else:
|
|
await save_group(capture, chat)
|
|
except (BadRequest, Forbidden):
|
|
pass
|
|
except Exception:
|
|
logger.exception(f"Failed to snapshot dialog {chat_id}")
|
|
await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id)
|
|
processed += 1
|
|
if processed % SAVE_EVERY == 0:
|
|
await ctx.report_progress({"processed": processed})
|
|
await _enrich_users(client, capture, nameless)
|
|
await ctx.report_progress({"processed": processed, "done": True})
|