Switch from SQLAlchemy to asyncpg/aiosqlite

mautrix_telegram/db/puppet.py

@@ -1,5 +1,5 @@
 # mautrix-telegram - A Matrix-Telegram puppeting bridge
-# Copyright (C) 2019 Tulir Asokan
+# 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
@@ -13,51 +13,106 @@
 #
 # 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, Iterable
+from __future__ import annotations
 
-from sqlalchemy import Column, Integer, BigInteger, String, Text, Boolean
-from sqlalchemy.sql import expression, func
+from typing import ClassVar, TYPE_CHECKING
+
+from asyncpg import Record
+from attr import dataclass
+from yarl import URL
 
 from mautrix.types import UserID, SyncToken
-from mautrix.util.db import Base
+from mautrix.util.async_db import Database
 
 from ..types import TelegramID
 
+fake_db = Database.create("") if TYPE_CHECKING else None
 
-class Puppet(Base):
-    __tablename__ = "puppet"
 
-    id: TelegramID = Column(BigInteger, primary_key=True)
-    custom_mxid: UserID = Column(String, nullable=True)
-    access_token: str = Column(String, nullable=True)
-    next_batch: SyncToken = Column(String, nullable=True)
-    base_url: str = Column(Text, nullable=True)
-    displayname: str = Column(String, nullable=True)
-    displayname_source: TelegramID = Column(BigInteger, nullable=True)
-    displayname_contact: bool = Column(Boolean, nullable=False, server_default=expression.true())
-    displayname_quality: int = Column(Integer, nullable=False, server_default="0")
-    username: str = Column(String, nullable=True)
-    photo_id: str = Column(String, nullable=True)
-    is_bot: bool = Column(Boolean, nullable=True)
-    matrix_registered: bool = Column(Boolean, nullable=False, server_default=expression.false())
-    disable_updates: bool = Column(Boolean, nullable=False, server_default=expression.false())
+@dataclass
+class Puppet:
+    db: ClassVar[Database] = fake_db
+
+    id: TelegramID
+
+    is_registered: bool
+
+    displayname: str | None
+    displayname_source: TelegramID | None
+    displayname_contact: bool
+    displayname_quality: int
+    disable_updates: bool
+    username: str | None
+    photo_id: str | None
+    is_bot: bool | None
+
+    custom_mxid: UserID | None
+    access_token: str | None
+    next_batch: SyncToken | None
+    base_url: URL | None
 
     @classmethod
-    def all_with_custom_mxid(cls) -> Iterable['Puppet']:
-        yield from cls._select_all(cls.c.custom_mxid != None)
+    def _from_row(cls, row: Record | None) -> Puppet | None:
+        if row is None:
+            return None
+        data = {**row}
+        base_url = data.pop("base_url", None)
+        return cls(**data, base_url=URL(base_url) if base_url else None)
+
+    columns: ClassVar[str] = (
+        "id, is_registered, displayname, displayname_source, displayname_contact, "
+        "displayname_quality, disable_updates, username, photo_id, is_bot, "
+        "custom_mxid, access_token, next_batch, base_url"
+    )
 
     @classmethod
-    def get_by_tgid(cls, tgid: TelegramID) -> Optional['Puppet']:
-        return cls._select_one_or_none(cls.c.id == tgid)
+    async def all_with_custom_mxid(cls) -> list[Puppet]:
+        q = f"SELECT {cls.columns} FROM puppet WHERE custom_mxid<>''"
+        return [cls._from_row(row) for row in await cls.db.fetch(q)]
 
     @classmethod
-    def get_by_custom_mxid(cls, mxid: UserID) -> Optional['Puppet']:
-        return cls._select_one_or_none(cls.c.custom_mxid == mxid)
+    async def get_by_tgid(cls, tgid: TelegramID) -> Puppet | None:
+        q = f"SELECT {cls.columns} FROM puppet WHERE id=$1"
+        return cls._from_row(await cls.db.fetchrow(q, tgid))
 
     @classmethod
-    def get_by_username(cls, username: str) -> Optional['Puppet']:
-        return cls._select_one_or_none(func.lower(cls.c.username) == username)
+    async def get_by_custom_mxid(cls, mxid: UserID) -> Puppet | None:
+        q = f"SELECT {cls.columns} FROM puppet WHERE custom_mxid=$1"
+        return cls._from_row(await cls.db.fetchrow(q, mxid))
 
     @classmethod
-    def get_by_displayname(cls, displayname: str) -> Optional['Puppet']:
-        return cls._select_one_or_none(cls.c.displayname == displayname)
+    async def find_by_username(cls, username: str) -> Puppet | None:
+        q = f"SELECT {cls.columns} FROM puppet WHERE lower(username)=$1"
+        return cls._from_row(await cls.db.fetchrow(q, username.lower()))
+
+    @classmethod
+    async def find_by_displayname(cls, displayname: str) -> Puppet | None:
+        q = f"SELECT {cls.columns} FROM puppet WHERE displayname=$1"
+        return cls._from_row(await cls.db.fetchrow(q, displayname))
+
+    @property
+    def _values(self):
+        return (self.id, self.is_registered, self.displayname, self.displayname_source,
+                self.displayname_contact, self.displayname_quality, self.disable_updates,
+                self.username, self.photo_id, self.is_bot, self.custom_mxid, self.access_token,
+                self.next_batch, str(self.base_url) if self.base_url else None)
+
+    async def save(self) -> None:
+        q = (
+            "UPDATE puppet "
+            "SET is_registered=$2, displayname=$3, displayname_source=$4, displayname_contact=$5,"
+            "    displayname_quality=$6, disable_updates=$7, username=$8, photo_id=$9, is_bot=$10,"
+            "    custom_mxid=$11, access_token=$12, next_batch=$13, base_url=$14 "
+            "WHERE id=$1"
+        )
+        await self.db.execute(q, *self._values)
+
+    async def insert(self) -> None:
+        q = (
+            "INSERT INTO puppet ("
+            "    id, is_registered, displayname, displayname_source, displayname_contact,"
+            "    displayname_quality, disable_updates, username, photo_id, is_bot,"
+            "    custom_mxid, access_token, next_batch, base_url"
+            ") VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)"
+        )
+        await self.db.execute(q, *self._values)