Switch from SQLAlchemy to asyncpg/aiosqlite

mautrix_telegram/util/__init__.py

@@ -2,3 +2,5 @@ from .file_transfer import transfer_file_to_matrix, convert_image
 from .parallel_file_transfer import parallel_transfer_to_telegram
 from .recursive_dict import recursive_del, recursive_set, recursive_get
 from .color_log import ColorFormatter
+from .send_lock import PortalSendLock
+from .deduplication import PortalDedup

mautrix_telegram/util/deduplication.py

@@ -0,0 +1,126 @@
+# mautrix-telegram - A Matrix-Telegram puppeting bridge
+# Copyright (C) 2021 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 <https://www.gnu.org/licenses/>.
+from typing import Optional, Deque, Dict, Tuple, TYPE_CHECKING
+from collections import deque
+import hashlib
+
+from telethon.tl.patched import Message, MessageService
+from telethon.tl.types import (MessageMediaContact, MessageMediaDocument, MessageMediaGeo,
+                               MessageMediaPhoto, TypeMessage, TypeUpdates, UpdateNewMessage,
+                               UpdateNewChannelMessage)
+
+from mautrix.types import EventID
+
+from ..types import TelegramID
+
+if TYPE_CHECKING:
+    from ..portal import Portal
+
+DedupMXID = Tuple[EventID, TelegramID]
+
+
+class PortalDedup:
+    pre_db_check: bool = False
+    cache_queue_length: int = 20
+
+    _dedup: Deque[str]
+    _dedup_mxid: Dict[str, DedupMXID]
+    _dedup_action: Deque[str]
+    _portal: 'Portal'
+
+    def __init__(self, portal: 'Portal') -> None:
+        self._dedup = deque()
+        self._dedup_mxid = {}
+        self._dedup_action = deque()
+        self._portal = portal
+
+    @property
+    def _always_force_hash(self) -> bool:
+        return self._portal.peer_type == 'chat'
+
+    @staticmethod
+    def _hash_event(event: TypeMessage) -> str:
+        # Non-channel messages are unique per-user (wtf telegram), so we have no other choice than
+        # to deduplicate based on a hash of the message content.
+
+        # The timestamp is only accurate to the second, so we can't rely solely on that either.
+        if isinstance(event, MessageService):
+            hash_content = [event.date.timestamp(), event.from_id, event.action]
+        else:
+            hash_content = [event.date.timestamp(), event.message.strip()]
+            if event.fwd_from:
+                hash_content += [event.fwd_from.from_id]
+            elif isinstance(event, Message) and event.media:
+                try:
+                    hash_content += {
+                        MessageMediaContact: lambda media: [media.user_id],
+                        MessageMediaDocument: lambda media: [media.document.id],
+                        MessageMediaPhoto: lambda media: [media.photo.id if media.photo else 0],
+                        MessageMediaGeo: lambda media: [media.geo.long, media.geo.lat],
+                    }[type(event.media)](event.media)
+                except KeyError:
+                    pass
+        return hashlib.md5("-"
+                           .join(str(a) for a in hash_content)
+                           .encode("utf-8")
+                           ).hexdigest()
+
+    def check_action(self, event: TypeMessage) -> bool:
+        evt_hash = self._hash_event(event) if self._always_force_hash else event.id
+        if evt_hash in self._dedup_action:
+            return True
+
+        self._dedup_action.append(evt_hash)
+
+        if len(self._dedup_action) > self.cache_queue_length:
+            self._dedup_action.popleft()
+        return False
+
+    def update(self, event: TypeMessage, mxid: DedupMXID = None,
+               expected_mxid: Optional[DedupMXID] = None, force_hash: bool = False
+               ) -> Optional[DedupMXID]:
+        evt_hash = self._hash_event(event) if self._always_force_hash or force_hash else event.id
+        try:
+            found_mxid = self._dedup_mxid[evt_hash]
+        except KeyError:
+            return EventID("None"), TelegramID(0)
+
+        if found_mxid != expected_mxid:
+            return found_mxid
+        self._dedup_mxid[evt_hash] = mxid
+        return None
+
+    def check(self, event: TypeMessage, mxid: DedupMXID = None, force_hash: bool = False
+              ) -> Optional[DedupMXID]:
+        evt_hash = (self._hash_event(event)
+                    if self._always_force_hash or force_hash
+                    else event.id)
+        if evt_hash in self._dedup:
+            return self._dedup_mxid[evt_hash]
+
+        self._dedup_mxid[evt_hash] = mxid
+        self._dedup.append(evt_hash)
+
+        if len(self._dedup) > self.cache_queue_length:
+            del self._dedup_mxid[self._dedup.popleft()]
+        return None
+
+    def register_outgoing_actions(self, response: TypeUpdates) -> None:
+        for update in response.updates:
+            check_dedup = (isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage))
+                           and isinstance(update.message, MessageService))
+            if check_dedup:
+                self.check(update.message)

mautrix_telegram/util/file_transfer.py

@@ -21,7 +21,8 @@ import asyncio
 import tempfile
 
 import magic
-from sqlalchemy.exc import IntegrityError, InvalidRequestError
+from asyncpg import UniqueViolationError
+from sqlite3 import IntegrityError
 
 from telethon.tl.types import (Document, InputFileLocation, InputDocumentFileLocation,
                                TypePhotoSize, PhotoSize, PhotoCachedSize, InputPhotoFileLocation,
@@ -123,7 +124,7 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In
     if custom_data:
         loc_id += "-mau_custom_thumbnail"
 
-    db_file = DBTelegramFile.get(loc_id)
+    db_file = await DBTelegramFile.get(loc_id)
     if db_file:
         return db_file
 
@@ -154,8 +155,8 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In
                              was_converted=False, timestamp=int(time.time()), size=len(file),
                              width=width, height=height, decryption_info=decryption_info)
     try:
-        db_file.insert()
-    except (IntegrityError, InvalidRequestError) as e:
+        await db_file.insert()
+    except (UniqueViolationError, IntegrityError) as e:
         log.exception(f"{e.__class__.__name__} while saving transferred file thumbnail data. "
                       "This was probably caused by two simultaneous transfers of the same file, "
                       "and might (but probably won't) cause problems with thumbnails or something.")
@@ -176,7 +177,7 @@ async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentA
     if not location_id:
         return None
 
-    db_file = DBTelegramFile.get(location_id)
+    db_file = await DBTelegramFile.get(location_id)
     if db_file:
         return db_file
 
@@ -197,7 +198,7 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten
                                             tgs_convert: Optional[dict], filename: Optional[str],
                                             encrypt: bool, parallel_id: Optional[int]
                                             ) -> Optional[DBTelegramFile]:
-    db_file = DBTelegramFile.get(loc_id)
+    db_file = await DBTelegramFile.get(loc_id)
     if db_file:
         return db_file
 
@@ -263,8 +264,8 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten
             width=converted_anim.width, height=converted_anim.height)
 
     try:
-        db_file.insert()
-    except (IntegrityError, InvalidRequestError) as e:
+        await db_file.insert()
+    except (UniqueViolationError, IntegrityError) as e:
         log.exception(f"{e.__class__.__name__} while saving transferred file data. "
                       "This was probably caused by two simultaneous transfers of the same file, "
                       "and should not cause any problems.")

mautrix_telegram/util/send_lock.py

@@ -0,0 +1,44 @@
+# mautrix-telegram - A Matrix-Telegram puppeting bridge
+# Copyright (C) 2021 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 <https://www.gnu.org/licenses/>.
+from typing import Dict
+from asyncio import Lock
+
+from ..types import TelegramID
+
+
+class FakeLock:
+    async def __aenter__(self) -> None:
+        pass
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        pass
+
+
+class PortalSendLock:
+    _send_locks: Dict[int, Lock]
+    _noop_lock: Lock = FakeLock()
+
+    def __init__(self) -> None:
+        self._send_locks = {}
+
+    def __call__(self, user_id: TelegramID, required: bool = True) -> Lock:
+        if user_id is None and required:
+            raise ValueError("Required send lock for none id")
+        try:
+            return self._send_locks[user_id]
+        except KeyError:
+            return (self._send_locks.setdefault(user_id, Lock())
+                    if required else self._noop_lock)