diff --git a/backend/src/userbot/handlers/edits.py b/backend/src/userbot/handlers/edits.py index 396b940..833b63d 100644 --- a/backend/src/userbot/handlers/edits.py +++ b/backend/src/userbot/handlers/edits.py @@ -2,10 +2,9 @@ from pyrogram.types import Message from userbot.modules.capture import repository from userbot.modules.capture.chat_meta import meta_from_chat -from userbot.modules.capture.message import sender_id +from userbot.modules.capture.message import message_row from userbot.modules.client import PyroClient -from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl -from userbot.modules.raw_json import raw_json +from userbot.modules.media import capture_media from utils.events import notify_bg_event @@ -27,18 +26,7 @@ async def on_edited_message(client: PyroClient, message: Message) -> None: if not toggles.track_edits_deletes: return changed = await repository.add_version( - ctx.pool, - ctx.account_id, - chat_id, - message.id, - message.date, - sender_id(message), - message.text or message.caption, - raw_json(message), - message.edit_date, - media_unique_id(message), - has_media=message.media is not None, - is_self_destruct=self_destruct_ttl(message) is not None, + ctx.pool, ctx.account_id, message_row(message, chat_id, message.date) ) if not changed: return diff --git a/backend/src/userbot/modules/capture/message.py b/backend/src/userbot/modules/capture/message.py index ad1a14a..34bdfab 100644 --- a/backend/src/userbot/modules/capture/message.py +++ b/backend/src/userbot/modules/capture/message.py @@ -1,11 +1,14 @@ +from datetime import datetime + from pyrogram import Client from pyrogram.types import Message from userbot.modules.capture import repository from userbot.modules.capture.context import CaptureContext +from userbot.modules.capture.repository import MessageRow from userbot.modules.links import extract_links from userbot.modules.links import repository as links_repository -from userbot.modules.media import capture_media, self_destruct_ttl +from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl from userbot.modules.raw_json import raw_json from utils.policy.models import CaptureToggles @@ -34,27 +37,35 @@ def callbacks(message: Message) -> list[tuple[int, str | None, bytes | None]]: return buttons +def message_row(message: Message, chat_id: int, date: datetime) -> MessageRow: + return MessageRow( + chat_id=chat_id, + message_id=message.id, + date=date, + sender_id=sender_id(message), + text=message.text or message.caption, + raw=raw_json(message), + has_media=message.media is not None, + is_self_destruct=self_destruct_ttl(message) is not None, + edit_date=message.edit_date, + media_unique_id=media_unique_id(message), + ) + + async def capture_message( client: Client, message: Message, ctx: CaptureContext, toggles: CaptureToggles ) -> None: if message.empty or message.chat is None or message.date is None: return chat_id = message.chat.id or 0 - await repository.upsert_message( - ctx.pool, - ctx.account_id, - chat_id, - message.id, - message.date, - sender_id(message), - message.text or message.caption, - raw_json(message), - has_media=message.media is not None, - is_self_destruct=self_destruct_ttl(message) is not None, + current = await repository.upsert_message( + ctx.pool, ctx.account_id, message_row(message, chat_id, message.date) ) from userbot.modules.capture.identity import capture_identity # noqa: PLC0415 await capture_identity(message, ctx) + if not current: + return await capture_media(client, message, ctx, chat_id, message.id, toggles) buttons = callbacks(message) if buttons: diff --git a/backend/src/userbot/modules/capture/repository.py b/backend/src/userbot/modules/capture/repository.py index e87c327..181320f 100644 --- a/backend/src/userbot/modules/capture/repository.py +++ b/backend/src/userbot/modules/capture/repository.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from datetime import datetime import asyncpg @@ -21,13 +22,19 @@ _TOUCH_EDITED = """ INSERT INTO messages (account_id, chat_id, message_id, date, sender_id, text, raw, has_media, is_self_destruct, edited_at) -VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9, now()) +VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9, CASE WHEN $10 THEN now() END) ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET text = EXCLUDED.text, raw = EXCLUDED.raw, has_media = EXCLUDED.has_media, is_self_destruct = EXCLUDED.is_self_destruct, - edited_at = now() + edited_at = coalesce(EXCLUDED.edited_at, messages.edited_at) +""" + +_LOCK_MESSAGE = """ +SELECT pg_advisory_xact_lock( + hashtextextended(concat_ws(':', 'message', $1::int, $2::bigint, $3::bigint), 0) +) """ _FIND_DELETED = """ @@ -61,9 +68,32 @@ WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3 ON CONFLICT DO NOTHING """ +_INSERT_PAST_VERSION = """ +INSERT INTO message_versions + (account_id, chat_id, message_id, observed_at, edit_date, text, raw) +SELECT $1::int, $2::bigint, $3::bigint, + coalesce( + (SELECT min(v.observed_at) - interval '1 microsecond' + FROM message_versions v + WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3 + AND v.edit_date > $4::timestamptz), + $4::timestamptz, + $5::timestamptz + ), + $4::timestamptz, $6::text, $7::jsonb +WHERE NOT EXISTS ( + SELECT 1 FROM message_versions v + WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3 + AND v.edit_date IS NOT DISTINCT FROM $4::timestamptz + AND v.text IS NOT DISTINCT FROM $6::text +) +ON CONFLICT DO NOTHING +""" + _CURRENT_CONTENT = """ SELECT m.text AS text, + m.raw->>'edit_date' AS edit_date, (SELECT d.unique_id FROM media d WHERE d.account_id = m.account_id AND d.chat_id = m.chat_id AND d.message_id = m.message_id) AS media_unique_id @@ -121,33 +151,108 @@ WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL) """ -async def upsert_message( # noqa: PLR0913 - pool: asyncpg.Pool, - account_id: int, - chat_id: int, - message_id: int, - date: datetime, - sender_id: int | None, - text: str | None, - raw: str, - *, - has_media: bool, - is_self_destruct: bool, -) -> None: - await pool.execute( - _UPSERT_MESSAGE, - account_id, - chat_id, - message_id, - date, - sender_id, - text, - raw, - has_media, - is_self_destruct, +@dataclass(frozen=True, slots=True) +class MessageRow: + chat_id: int + message_id: int + date: datetime + sender_id: int | None + text: str | None + raw: str + has_media: bool + is_self_destruct: bool + edit_date: datetime | None + media_unique_id: str | None + + @property + def columns(self) -> tuple: + return ( + self.chat_id, + self.message_id, + self.date, + self.sender_id, + self.text, + self.raw, + self.has_media, + self.is_self_destruct, + ) + + +async def _lock_current( + conn: asyncpg.Connection, account_id: int, row: MessageRow +) -> asyncpg.Record | None: + await conn.execute(_LOCK_MESSAGE, account_id, row.chat_id, row.message_id) + return await conn.fetchrow( + _CURRENT_CONTENT, account_id, row.chat_id, row.message_id ) +def _is_stale(current: asyncpg.Record, row: MessageRow) -> bool: + stored = current["edit_date"] + if stored is None: + return False + if row.edit_date is None: + return True + try: + return row.edit_date < datetime.fromisoformat(stored) + except (TypeError, ValueError): + return False + + +async def _record_past( + conn: asyncpg.Connection, account_id: int, row: MessageRow +) -> None: + await conn.execute( + _INSERT_PAST_VERSION, + account_id, + row.chat_id, + row.message_id, + row.edit_date, + row.date, + row.text, + row.raw, + ) + + +async def _record_edit( + conn: asyncpg.Connection, account_id: int, row: MessageRow, *, mark_edited: bool +) -> None: + await conn.execute(_SNAPSHOT_ORIGINAL, account_id, row.chat_id, row.message_id) + await conn.execute(_TOUCH_EDITED, account_id, *row.columns, mark_edited) + await conn.execute( + _INSERT_VERSION, + account_id, + row.chat_id, + row.message_id, + row.edit_date, + row.text, + row.raw, + ) + + +async def upsert_message(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool: + async with pool.acquire() as conn, conn.transaction(): + current = await _lock_current(conn, account_id, row) + if current is None: + await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns) + return True + stored_media = current["media_unique_id"] + changed = current["text"] != row.text or ( + stored_media is not None and stored_media != row.media_unique_id + ) + if _is_stale(current, row): + if changed: + await _record_past(conn, account_id, row) + return False + if changed: + await _record_edit( + conn, account_id, row, mark_edited=row.edit_date is not None + ) + else: + await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns) + return True + + async def max_message_id( pool: asyncpg.Pool, account_id: int, chat_id: int ) -> int | None: @@ -188,45 +293,20 @@ async def mark_deleted_channel( await _mark_deleted(pool, account_id, rows) -async def add_version( # noqa: PLR0913 - pool: asyncpg.Pool, - account_id: int, - chat_id: int, - message_id: int, - date: datetime, - sender_id: int | None, - text: str | None, - raw: str, - edit_date: datetime | None, - media_unique_id: str | None, - *, - has_media: bool, - is_self_destruct: bool, -) -> bool: +async def add_version(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool: async with pool.acquire() as conn, conn.transaction(): - current = await conn.fetchrow(_CURRENT_CONTENT, account_id, chat_id, message_id) - text_changed = current is None or current["text"] != text - media_changed = ( - current is not None and current["media_unique_id"] != media_unique_id - ) - if not (text_changed or media_changed): - return False - await conn.execute(_SNAPSHOT_ORIGINAL, account_id, chat_id, message_id) - await conn.execute( - _TOUCH_EDITED, - account_id, - chat_id, - message_id, - date, - sender_id, - text, - raw, - has_media, - is_self_destruct, - ) - await conn.execute( - _INSERT_VERSION, account_id, chat_id, message_id, edit_date, text, raw - ) + current = await _lock_current(conn, account_id, row) + if current is not None: + changed = ( + current["text"] != row.text + or current["media_unique_id"] != row.media_unique_id + ) + if not changed: + return False + if _is_stale(current, row): + await _record_past(conn, account_id, row) + return False + await _record_edit(conn, account_id, row, mark_edited=True) return True