88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
from pyrogram import Client
|
|
from pyrogram.errors import BadRequest, Forbidden
|
|
from pyrogram.types import User
|
|
|
|
from userbot.modules.avatars import note_avatar
|
|
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.parse import snapshot_from_high_level
|
|
from userbot.modules.profiles.repository import write_profile
|
|
from userbot.modules.profiles.snapshots import save_group, save_private
|
|
from userbot.modules.raw_json import raw_json
|
|
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 not isinstance(user, User):
|
|
continue
|
|
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
|
await write_profile(
|
|
ctx.pool,
|
|
ctx.account_id,
|
|
user.id,
|
|
fields,
|
|
raw_json(user),
|
|
history=ctx.resolve_peer(
|
|
user.id, is_bot=bool(user.is_bot)
|
|
).profile_history,
|
|
)
|
|
if photo_file_id and photo_unique_id:
|
|
await note_avatar(
|
|
ctx.pool,
|
|
ctx.account_id,
|
|
user.id,
|
|
"peer",
|
|
photo_unique_id,
|
|
photo_file_id,
|
|
)
|
|
|
|
|
|
@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})
|