fix(userbot): a deletion that arrives before its message is applied once the message is stored
This commit is contained in:
1 parent
516b8e307f
commit
f32b772ddc
5 files changed
+56
No files matched your search
@@ -20,6 +20,9 @@ async def on_deleted_messages(client: PyroClient, messages: list[Message]) -> No
|
|||||||
box.append(message.id)
|
box.append(message.id)
|
||||||
else:
|
else:
|
||||||
channels.setdefault(chat_id, []).append(message.id)
|
channels.setdefault(chat_id, []).append(message.id)
|
||||||
|
ctx.pending_deletes.remember(None, box)
|
||||||
|
for chat_id, ids in channels.items():
|
||||||
|
ctx.pending_deletes.remember(chat_id, ids)
|
||||||
if box:
|
if box:
|
||||||
await repository.mark_deleted_box(ctx.pool, ctx.account_id, box)
|
await repository.mark_deleted_box(ctx.pool, ctx.account_id, box)
|
||||||
await notify_bg_event(ctx.pool, "delete", ctx.account_id, message_ids=box)
|
await notify_bg_event(ctx.pool, "delete", ctx.account_id, message_ids=box)
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ import asyncpg
|
|||||||
from pyrogram import Client
|
from pyrogram import Client
|
||||||
|
|
||||||
from userbot.modules.capture.identity import ChatMetaCache, PeerIdentityCache
|
from userbot.modules.capture.identity import ChatMetaCache, PeerIdentityCache
|
||||||
|
from userbot.modules.capture.pending import PendingDeletes
|
||||||
from userbot.modules.contacts import ContactCache
|
from userbot.modules.contacts import ContactCache
|
||||||
from userbot.modules.folders import FolderCache
|
from userbot.modules.folders import FolderCache
|
||||||
from userbot.modules.watches import WatchCache
|
from userbot.modules.watches import WatchCache
|
||||||
@@ -28,6 +29,7 @@ class CaptureContext:
|
|||||||
self.watches = WatchCache(pool, account_id)
|
self.watches = WatchCache(pool, account_id)
|
||||||
self.peer_identity = PeerIdentityCache()
|
self.peer_identity = PeerIdentityCache()
|
||||||
self.chat_meta = ChatMetaCache()
|
self.chat_meta = ChatMetaCache()
|
||||||
|
self.pending_deletes = PendingDeletes()
|
||||||
self.policies = PolicySet()
|
self.policies = PolicySet()
|
||||||
|
|
||||||
async def reload_policies(self) -> None:
|
async def reload_policies(self) -> None:
|
||||||
|
|||||||
@@ -61,6 +61,10 @@ async def capture_message(
|
|||||||
current = await repository.upsert_message(
|
current = await repository.upsert_message(
|
||||||
ctx.pool, ctx.account_id, message_row(message, chat_id, message.date)
|
ctx.pool, ctx.account_id, message_row(message, chat_id, message.date)
|
||||||
)
|
)
|
||||||
|
if ctx.pending_deletes.take(chat_id, message.id):
|
||||||
|
await repository.mark_deleted_late(
|
||||||
|
ctx.pool, ctx.account_id, chat_id, message.id
|
||||||
|
)
|
||||||
from userbot.modules.capture.identity import capture_identity # noqa: PLC0415
|
from userbot.modules.capture.identity import capture_identity # noqa: PLC0415
|
||||||
|
|
||||||
await capture_identity(message, ctx)
|
await capture_identity(message, ctx)
|
||||||
|
|||||||
@@ -0,0 +1,36 @@
|
|||||||
|
import time
|
||||||
|
|
||||||
|
from userbot.modules.capture.repository import CHANNEL_ID_THRESHOLD
|
||||||
|
|
||||||
|
PENDING_DELETE_SECONDS = 120.0
|
||||||
|
PENDING_DELETE_LIMIT = 10_000
|
||||||
|
|
||||||
|
DeleteKey = tuple[int | None, int]
|
||||||
|
|
||||||
|
|
||||||
|
def _key(chat_id: int | None, message_id: int) -> DeleteKey:
|
||||||
|
if chat_id is None or chat_id > CHANNEL_ID_THRESHOLD:
|
||||||
|
return None, message_id
|
||||||
|
return chat_id, message_id
|
||||||
|
|
||||||
|
|
||||||
|
class PendingDeletes:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._expires: dict[DeleteKey, float] = {}
|
||||||
|
|
||||||
|
def remember(self, chat_id: int | None, message_ids: list[int]) -> None:
|
||||||
|
now = time.monotonic()
|
||||||
|
while self._expires:
|
||||||
|
oldest = next(iter(self._expires))
|
||||||
|
fresh = self._expires[oldest] > now
|
||||||
|
if fresh and len(self._expires) < PENDING_DELETE_LIMIT:
|
||||||
|
break
|
||||||
|
del self._expires[oldest]
|
||||||
|
for message_id in message_ids:
|
||||||
|
key = _key(chat_id, message_id)
|
||||||
|
self._expires.pop(key, None)
|
||||||
|
self._expires[key] = now + PENDING_DELETE_SECONDS
|
||||||
|
|
||||||
|
def take(self, chat_id: int, message_id: int) -> bool:
|
||||||
|
expires = self._expires.pop(_key(chat_id, message_id), None)
|
||||||
|
return expires is not None and expires > time.monotonic()
|
||||||
@@ -47,6 +47,11 @@ UPDATE messages SET deleted_at = now()
|
|||||||
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
|
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
_MARK_DELETED_LATE = """
|
||||||
|
UPDATE messages SET deleted_at = now()
|
||||||
|
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND deleted_at IS NULL
|
||||||
|
"""
|
||||||
|
|
||||||
_INSERT_VERSION = """
|
_INSERT_VERSION = """
|
||||||
INSERT INTO message_versions
|
INSERT INTO message_versions
|
||||||
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
||||||
@@ -310,6 +315,12 @@ async def add_version(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> b
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
async def mark_deleted_late(
|
||||||
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||||
|
) -> None:
|
||||||
|
await pool.execute(_MARK_DELETED_LATE, account_id, chat_id, message_id)
|
||||||
|
|
||||||
|
|
||||||
async def current_media(
|
async def current_media(
|
||||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||||
) -> asyncpg.Record | None:
|
) -> asyncpg.Record | None:
|
||||||
|
|||||||
Reference in new issue
Block a user