From 493955c59f699c712dd068665d4a2d37ec62fd1e Mon Sep 17 00:00:00 2001 From: Toni Spets Date: Mon, 11 Dec 2023 17:42:10 +0200 Subject: [PATCH] WIP: Pump disappearing messages from db --- mautrix_telegram/db/disappearing_message.py | 50 +++++++++++++++---- mautrix_telegram/db/upgrade/__init__.py | 1 + .../db/upgrade/v00_latest_revision.py | 4 +- .../v19_disappearing_message_unqueue.py | 26 ++++++++++ mautrix_telegram/portal.py | 40 ++++++++++++++- mautrix_telegram/version.py | 7 ++- 6 files changed, 115 insertions(+), 13 deletions(-) create mode 100644 mautrix_telegram/db/upgrade/v19_disappearing_message_unqueue.py diff --git a/mautrix_telegram/db/disappearing_message.py b/mautrix_telegram/db/disappearing_message.py index 6c217173..6487bec3 100644 --- a/mautrix_telegram/db/disappearing_message.py +++ b/mautrix_telegram/db/disappearing_message.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, ClassVar import asyncpg +import time from mautrix.bridge import AbstractDisappearingMessage from mautrix.types import EventID, RoomID @@ -27,6 +28,7 @@ class DisappearingMessage(AbstractDisappearingMessage): + unqueued_ts: int | None = None db: ClassVar[Database] = fake_db async def insert(self) -> None: @@ -50,6 +52,40 @@ async def delete(self) -> None: def _from_row(cls, row: asyncpg.Record) -> DisappearingMessage: return cls(**row) + """ + Get all scheduled messages that will expire in given seconds that haven't yet been unqueued. + + This will also stamp them in the database for being unqueued so every time this method is called + there should be a unique set of events. If seconds is None then all events will be returned + regardless of being requested before. + + The first call on startup should be with None and subsequent with the previous value. + """ + @classmethod + async def unqueue_expiring(cls, seconds: int | None = None) -> list[DisappearingMessage]: + unqueued_ts = int(time.time() * 1000) + + rows = None + if seconds is None: + q = """ + SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message + WHERE expiration_ts <= $1 + """ + rows = await cls.db.fetch(q, unqueued_ts) + else: + q = """ + SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message + WHERE expiration_ts <= $1 AND (unqueued_ts IS NULL OR unqueued_ts < $2) + """ + rows = await cls.db.fetch(q, unqueued_ts + (seconds * 1000), unqueued_ts) + + msgs = [cls._from_row(r) for r in rows] + for msg in msgs: + msg.unqueued_ts = unqueued_ts + await msg.update() + + return msgs + @classmethod async def get(cls, room_id: RoomID, event_id: EventID) -> DisappearingMessage | None: q = """ @@ -63,16 +99,10 @@ async def get(cls, room_id: RoomID, event_id: EventID) -> DisappearingMessage | @classmethod async def get_all_scheduled(cls) -> list[DisappearingMessage]: - q = """ - SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message - WHERE expiration_ts IS NOT NULL - """ - return [cls._from_row(r) for r in await cls.db.fetch(q)] + # Stubbed because we pump with unqueue_expiring + return [] @classmethod async def get_unscheduled_for_room(cls, room_id: RoomID) -> list[DisappearingMessage]: - q = """ - SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message - WHERE room_id = $1 AND expiration_ts IS NULL - """ - return [cls._from_row(r) for r in await cls.db.fetch(q, room_id)] + # Stubbed because we pump with unqueue_expiring + return [] diff --git a/mautrix_telegram/db/upgrade/__init__.py b/mautrix_telegram/db/upgrade/__init__.py index 34cedc6e..a8f0b9e8 100644 --- a/mautrix_telegram/db/upgrade/__init__.py +++ b/mautrix_telegram/db/upgrade/__init__.py @@ -21,4 +21,5 @@ v16_backfill_type, v17_message_find_recent, v18_puppet_contact_info_set, + v19_disappearing_message_unqueue, ) diff --git a/mautrix_telegram/db/upgrade/v00_latest_revision.py b/mautrix_telegram/db/upgrade/v00_latest_revision.py index 66c1989a..c65e78dd 100644 --- a/mautrix_telegram/db/upgrade/v00_latest_revision.py +++ b/mautrix_telegram/db/upgrade/v00_latest_revision.py @@ -15,7 +15,7 @@ # along with this program. If not, see . from mautrix.util.async_db import Connection, Scheme -latest_version = 18 +latest_version = 19 async def create_latest_tables(conn: Connection, scheme: Scheme) -> int: @@ -92,10 +92,12 @@ async def create_latest_tables(conn: Connection, scheme: Scheme) -> int: event_id TEXT, expiration_seconds BIGINT, expiration_ts BIGINT, + unqueued_ts BIGINT, PRIMARY KEY (room_id, event_id) )""" ) + await conn.execute("CREATE INDEX disappearing_message_expiration_ts ON disappearing_message(expiration_ts)") await conn.execute( """CREATE TABLE puppet ( id BIGINT PRIMARY KEY, diff --git a/mautrix_telegram/db/upgrade/v19_disappearing_message_unqueue.py b/mautrix_telegram/db/upgrade/v19_disappearing_message_unqueue.py new file mode 100644 index 00000000..89c8780d --- /dev/null +++ b/mautrix_telegram/db/upgrade/v19_disappearing_message_unqueue.py @@ -0,0 +1,26 @@ +# mautrix-telegram - A Matrix-Telegram puppeting bridge +# Copyright (C) 2022 Tulir Asokan +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +from mautrix.util.async_db import Connection + +from . import upgrade_table + + +@upgrade_table.register(description="Add index on disappearing_message expiration_ts, unqueued_ts column") +async def upgrade_v19(conn: Connection) -> None: + await conn.execute( + "ALTER TABLE disappearing_message ADD COLUMN unqueued_ts BIGINT" + ) + await conn.execute("CREATE INDEX disappearing_message_expiration_ts ON disappearing_message(expiration_ts)") diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index 2ab062db..97a45216 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -322,6 +322,8 @@ class Portal(DBPortal, BasePortal): _msg_conv: putil.TelegramMessageConverter + _disappearing_event: asyncio.Event + def __init__( self, tgid: TelegramID, @@ -468,6 +470,42 @@ def set_dm_room_metadata(self) -> bool: or (self.encrypted and self.private_chat_portal_meta != "never") ) + @classmethod + async def _disappearing_message_loop(cls, seconds: int | None = None) -> None: + try: + seconds = None + while True: + print("fetching disappearing") + cls._disappearing_event.clear() + msgs = await cls.disappearing_msg_class.unqueue_expiring(seconds) + print(f"got {len(msgs)} rows") + for msg in msgs: + print("handling disappear thing") + portal = await cls.bridge.get_portal(msg.room_id) + if portal and portal.mxid: + background_task.create(portal._disappear_event(msg)) + else: + await msg.delete() + + try: + await asyncio.wait_for(cls._disappearing_event.wait(), 10) + except TimeoutError: + pass + + seconds = 10 + except RuntimeError: + return + + @classmethod + async def restart_scheduled_disappearing(cls) -> None: + cls._disappearing_event = asyncio.Event() + background_task.create(cls._disappearing_message_loop()) + + @classmethod + async def notify_disappearing_message_loop(cls) -> None: + print("notifying disappear loop") + cls._disappearing_event.set() + @classmethod def init_cls(cls, bridge: "TelegramBridge") -> None: BasePortal.bridge = bridge @@ -3531,7 +3569,7 @@ async def _mark_disappearing( ) await dm.insert() if expires_at: - background_task.create(self._disappear_event(dm)) + Portal.notify_disappearing_message_loop() async def _create_room_on_action( self, source: au.AbstractUser, action: TypeMessageAction diff --git a/mautrix_telegram/version.py b/mautrix_telegram/version.py index 0b226803..ccb8cd1d 100644 --- a/mautrix_telegram/version.py +++ b/mautrix_telegram/version.py @@ -1 +1,6 @@ -from .get_version import git_revision, git_tag, linkified_version, version +# Generated in setup.py + +git_tag = None +git_revision = "e3a067c2" +version = "0.11.3+dev.e3a067c2" +linkified_version = "0.11.3+dev.[e3a067c2](https://github.com/mautrix/telegram/commit/e3a067c27aa3d9dd5e82db307218cc66c8356ddd)"