Switch TelegramFile to SQLAlchemy core

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()

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