From b2e183e363b9e21e53fa7ed83473d00a88fdfe0a Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Thu, 20 Dec 2018 00:07:04 +0200 Subject: [PATCH] Switch TelegramFile to SQLAlchemy core --- mautrix_telegram/db.py | 23 ++++++++++++++- mautrix_telegram/util/file_transfer.py | 39 ++++++++++++++------------ 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/mautrix_telegram/db.py b/mautrix_telegram/db.py index 4fad2d2e..a4f36ea0 100644 --- a/mautrix_telegram/db.py +++ b/mautrix_telegram/db.py @@ -278,7 +278,9 @@ class BotChat(Base): class TelegramFile(Base): - query = None # type: Query + db = None # type: Engine + t = None # type: Table + c = None # type: ImmutableColumnCollection __tablename__ = "telegram_file" id = Column(String, primary_key=True) @@ -292,6 +294,25 @@ class TelegramFile(Base): thumbnail_id = Column("thumbnail", String, ForeignKey("telegram_file.id"), nullable=True) thumbnail = relationship("TelegramFile", uselist=False) + @classmethod + def get(cls, id: str) -> Optional['TelegramFile']: + rows = cls.db.execute(cls.t.select().where(cls.c.id == id)) + try: + id, mxc, mime, conv, ts, s, w, h, thumb_id = next(rows) + thumb = None + if thumb_id: + thumb = cls.get(thumb_id) + return TelegramFile(id=id, mxc=mxc, mime_type=mime, was_converted=conv, timestamp=ts, + size=s, width=w, height=h, thumbnail_id=thumb_id, thumbnail=thumb) + except StopIteration: + return None + + def insert(self) -> None: + self.db.execute(self.t.insert().values( + id=self.id, mxc=self.mxc, mime_type=self.mime_type, timestamp=self.timestamp, + size=self.size, width=self.width, height=self.height, + thumbnail_id=self.thumbnail.id if self.thumbnail else self.thumbnail_id)) + def init(db_session, db_engine) -> None: query = db_session.query_property() diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py index b60fc71b..415a3309 100644 --- a/mautrix_telegram/util/file_transfer.py +++ b/mautrix_telegram/util/file_transfer.py @@ -23,7 +23,6 @@ import asyncio import magic from sqlalchemy import orm from sqlalchemy.exc import IntegrityError, InvalidRequestError -from sqlalchemy.orm.exc import FlushError from telethon.tl.types import (Document, FileLocation, InputFileLocation, InputDocumentFileLocation, PhotoSize, PhotoCachedSize) @@ -117,6 +116,10 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In if not loc_id: return None + db_file = DBTelegramFile.get(loc_id) + if db_file: + return db_file + video_ext = mimetypes.guess_extension(mime) if VideoFileClip and video_ext: try: @@ -131,22 +134,29 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In content_uri = await intent.upload_file(file, mime_type) - return DBTelegramFile(id=loc_id, mxc=content_uri, mime_type=mime_type, - was_converted=False, timestamp=int(time.time()), size=len(file), - width=width, height=height) + db_file = DBTelegramFile(id=loc_id, mxc=content_uri, mime_type=mime_type, + was_converted=False, timestamp=int(time.time()), size=len(file), + width=width, height=height) + try: + db_file.insert() + except (IntegrityError, InvalidRequestError) 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.") + return db_file transfer_locks = {} # type: Dict[str, asyncio.Lock] -async def transfer_file_to_matrix(db: orm.Session, client: MautrixTelegramClient, intent: IntentAPI, +async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI, location: TypeLocation, thumbnail: Optional[TypeLocation] = None, is_sticker: bool = False) -> Optional[DBTelegramFile]: location_id = _location_to_id(location) if not location_id: return None - db_file = DBTelegramFile.query.get(location_id) + db_file = DBTelegramFile.get(location_id) if db_file: return db_file @@ -156,15 +166,15 @@ async def transfer_file_to_matrix(db: orm.Session, client: MautrixTelegramClient lock = asyncio.Lock() transfer_locks[location_id] = lock async with lock: - return await _unlocked_transfer_file_to_matrix(db, client, intent, location_id, location, + return await _unlocked_transfer_file_to_matrix(client, intent, location_id, location, thumbnail, is_sticker) -async def _unlocked_transfer_file_to_matrix(db: orm.Session, client: MautrixTelegramClient, - intent: IntentAPI, loc_id: str, location: TypeLocation, +async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI, + loc_id: str, location: TypeLocation, thumbnail: Optional[TypeLocation], is_sticker: bool) -> Optional[DBTelegramFile]: - db_file = DBTelegramFile.query.get(loc_id) + db_file = DBTelegramFile.get(loc_id) if db_file: return db_file @@ -201,16 +211,9 @@ async def _unlocked_transfer_file_to_matrix(db: orm.Session, client: MautrixTele mime_type) try: - db.add(db_file) - db.commit() - except FlushError 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.") + db_file.insert() except (IntegrityError, InvalidRequestError) as e: - db.rollback() 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.") - return db_file