Wrap database-changing statements in db.begin()

mautrix_telegram/db/base.py

@@ -44,14 +44,16 @@ class BaseBase:
         pass
 
     def update(self, **values) -> None:
-        self.db.execute(self.t.update()
-                        .where(self._edit_identity)
-                        .values(**values))
+        with self.db.begin() as conn:
+            conn.execute(self.t.update()
+                         .where(self._edit_identity)
+                         .values(**values))
         for key, value in values.items():
             setattr(self, key, value)
 
     def delete(self) -> None:
-        self.db.execute(self.t.delete().where(self._edit_identity))
+        with self.db.begin() as conn:
+            conn.execute(self.t.delete().where(self._edit_identity))
 
 
 Base = declarative_base(cls=BaseBase)

mautrix_telegram/db/bot_chat.py

@@ -30,7 +30,8 @@ class BotChat(Base):
 
     @classmethod
     def delete(cls, id: TelegramID) -> None:
-        cls.db.execute(cls.t.delete().where(cls.c.id == id))
+        with cls.db.begin() as conn:
+            conn.execute(cls.t.delete().where(cls.c.id == id))
 
     @classmethod
     def all(cls) -> Iterable['BotChat']:
@@ -40,4 +41,5 @@ class BotChat(Base):
             yield cls(id=id, type=type)
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(id=self.id, type=self.type))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(id=self.id, type=self.type))

mautrix_telegram/db/message.py

@@ -68,20 +68,23 @@ class Message(Base):
 
     @classmethod
     def update_by_tgid(cls, s_tgid: TelegramID, s_tg_space: TelegramID, **values) -> None:
-        cls.db.execute(cls.t.update()
-                       .where(and_(cls.c.tgid == s_tgid, cls.c.tg_space == s_tg_space))
-                       .values(**values))
+        with cls.db.begin() as conn:
+            conn.execute(cls.t.update()
+                         .where(and_(cls.c.tgid == s_tgid, cls.c.tg_space == s_tg_space))
+                         .values(**values))
 
     @classmethod
     def update_by_mxid(cls, s_mxid: MatrixEventID, s_mx_room: MatrixRoomID, **values) -> None:
-        cls.db.execute(cls.t.update()
-                       .where(and_(cls.c.mxid == s_mxid, cls.c.mx_room == s_mx_room))
-                       .values(**values))
+        with cls.db.begin() as conn:
+            conn.execute(cls.t.update()
+                         .where(and_(cls.c.mxid == s_mxid, cls.c.mx_room == s_mx_room))
+                         .values(**values))
 
     @property
     def _edit_identity(self):
         return and_(self.c.tgid == self.tgid, self.c.tg_space == self.tg_space)
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(mxid=self.mxid, mx_room=self.mx_room, tgid=self.tgid,
-                                               tg_space=self.tg_space))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(mxid=self.mxid, mx_room=self.mx_room,
+                                                tgid=self.tgid, tg_space=self.tg_space))

mautrix_telegram/db/portal.py

@@ -74,7 +74,8 @@ class Portal(Base):
         return and_(self.c.tgid == self.tgid, self.c.tg_receiver == self.tg_receiver)
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(
-            tgid=self.tgid, tg_receiver=self.tg_receiver, peer_type=self.peer_type,
-            megagroup=self.megagroup, mxid=self.mxid, config=self.config, username=self.username,
-            title=self.title, about=self.about, photo_id=self.photo_id))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(
+                tgid=self.tgid, tg_receiver=self.tg_receiver, peer_type=self.peer_type,
+                megagroup=self.megagroup, mxid=self.mxid, config=self.config,
+                username=self.username, title=self.title, about=self.about, photo_id=self.photo_id))

mautrix_telegram/db/puppet.py

@@ -79,8 +79,9 @@ class Puppet(Base):
         return self.c.id == self.id
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(
-            id=self.id, custom_mxid=self.custom_mxid, access_token=self.access_token,
-            displayname=self.displayname, displayname_source=self.displayname_source,
-            username=self.username, photo_id=self.photo_id, is_bot=self.is_bot,
-            matrix_registered=self.matrix_registered))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(
+                id=self.id, custom_mxid=self.custom_mxid, access_token=self.access_token,
+                displayname=self.displayname, displayname_source=self.displayname_source,
+                username=self.username, photo_id=self.photo_id, is_bot=self.is_bot,
+                matrix_registered=self.matrix_registered))

mautrix_telegram/db/room_state.py

@@ -47,14 +47,16 @@ class RoomState(Base):
             return None
 
     def update(self) -> None:
-        self.db.execute(self.t.update()
-                        .where(self.c.room_id == self.room_id)
-                        .values(power_levels=self._power_levels_text))
+        with self.db.begin() as conn:
+            conn.execute(self.t.update()
+                         .where(self.c.room_id == self.room_id)
+                         .values(power_levels=self._power_levels_text))
 
     @property
     def _edit_identity(self):
         return self.c.room_id == self.room_id
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(room_id=self.room_id,
-                                               power_levels=self._power_levels_text))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(room_id=self.room_id,
+                                                power_levels=self._power_levels_text))

mautrix_telegram/db/telegram_file.py

@@ -49,7 +49,9 @@ class TelegramFile(Base):
             return None
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(
-            id=self.id, mxc=self.mxc, mime_type=self.mime_type, was_converted=self.was_converted,
-            timestamp=self.timestamp, size=self.size, width=self.width, height=self.height,
-            thumbnail=self.thumbnail.id if self.thumbnail else self.thumbnail_id))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(
+                id=self.id, mxc=self.mxc, mime_type=self.mime_type,
+                was_converted=self.was_converted, timestamp=self.timestamp, size=self.size,
+                width=self.width, height=self.height,
+                thumbnail=self.thumbnail.id if self.thumbnail else self.thumbnail_id))

mautrix_telegram/db/user.py

@@ -65,9 +65,10 @@ class User(Base):
         return self.c.mxid == self.mxid
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(
-            mxid=self.mxid, tgid=self.tgid, tg_username=self.tg_username, tg_phone=self.tg_phone,
-            saved_contacts=self.saved_contacts))
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(
+                mxid=self.mxid, tgid=self.tgid, tg_username=self.tg_username,
+                tg_phone=self.tg_phone, saved_contacts=self.saved_contacts))
 
     @property
     def contacts(self) -> Iterable[TelegramID]:
@@ -78,10 +79,11 @@ class User(Base):
 
     @contacts.setter
     def contacts(self, puppets: Iterable[TelegramID]) -> None:
-        self.db.execute(Contact.t.delete().where(Contact.c.user == self.tgid))
-        insert_puppets = [{"user": self.tgid, "contact": tgid} for tgid in puppets]
-        if insert_puppets:
-            self.db.execute(Contact.t.insert(), insert_puppets)
+        with self.db.begin() as conn:
+            conn.execute(Contact.t.delete().where(Contact.c.user == self.tgid))
+            insert_puppets = [{"user": self.tgid, "contact": tgid} for tgid in puppets]
+            if insert_puppets:
+                conn.execute(Contact.t.insert(), insert_puppets)
 
     @property
     def portals(self) -> Iterable[Tuple[TelegramID, TelegramID]]:
@@ -92,14 +94,15 @@ class User(Base):
 
     @portals.setter
     def portals(self, portals: Iterable[Tuple[TelegramID, TelegramID]]) -> None:
-        self.db.execute(UserPortal.t.delete().where(UserPortal.c.user == self.tgid))
-        insert_portals = [{
-            "user": self.tgid,
-            "portal": tgid,
-            "portal_receiver": tg_receiver
-        } for tgid, tg_receiver in portals]
-        if insert_portals:
-            self.db.execute(UserPortal.t.insert(), insert_portals)
+        with self.db.begin() as conn:
+            conn.execute(UserPortal.t.delete().where(UserPortal.c.user == self.tgid))
+            insert_portals = [{
+                "user": self.tgid,
+                "portal": tgid,
+                "portal_receiver": tg_receiver
+            } for tgid, tg_receiver in portals]
+            if insert_portals:
+                conn.execute(UserPortal.t.insert(), insert_portals)
 
     def delete(self) -> None:
         super().delete()

mautrix_telegram/db/user_profile.py

@@ -50,7 +50,8 @@ class UserProfile(Base):
 
     @classmethod
     def delete_all(cls, room_id: MatrixRoomID) -> None:
-        cls.db.execute(cls.t.delete().where(cls.c.room_id == room_id))
+        with cls.db.begin() as conn:
+            conn.execute(cls.t.delete().where(cls.c.room_id == room_id))
 
     def update(self) -> None:
         super().update(membership=self.membership, displayname=self.displayname,
@@ -61,8 +62,8 @@ class UserProfile(Base):
         return and_(self.c.room_id == self.room_id, self.c.user_id == self.user_id)
 
     def insert(self) -> None:
-        self.db.execute(self.t.insert().values(room_id=self.room_id, user_id=self.user_id,
-                                               membership=self.membership,
-                                               displayname=self.displayname,
-                                               avatar_url=self.avatar_url))
-
+        with self.db.begin() as conn:
+            conn.execute(self.t.insert().values(room_id=self.room_id, user_id=self.user_id,
+                                                membership=self.membership,
+                                                displayname=self.displayname,
+                                                avatar_url=self.avatar_url))