Switch from SQLAlchemy to asyncpg/aiosqlite

mautrix_telegram/__main__.py

@@ -13,39 +13,27 @@
 #
 # 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, Any
-import sys
+from __future__ import annotations
+
+from typing import Any
 
 from telethon import __version__ as __telethon_version__
-from alchemysession import AlchemySessionContainer
 
 from mautrix.types import UserID, RoomID
 from mautrix.bridge import Bridge
-from mautrix.util.db import Base
-from mautrix.bridge.state_store.sqlalchemy import SQLBridgeStateStore
 
 from .web.provisioning import ProvisioningAPI
 from .web.public import PublicBridgeWebsite
-from .abstract_user import init as init_abstract_user
-from .bot import Bot, init as init_bot
+from .abstract_user import AbstractUser
+from .bot import Bot
 from .config import Config
-from .context import Context
-from .db import init as init_db
-from .formatter import init as init_formatter
+from .db import init as init_db, upgrade_table
 from .matrix import MatrixHandler
-from .portal import Portal, init as init_portal
-from .puppet import Puppet, init as init_puppet
-from .user import User, init as init_user
+from .portal import Portal
+from .puppet import Puppet
+from .user import User
 from .version import version, linkified_version
 
-import sqlalchemy as sql
-from sqlalchemy.engine.base import Engine
-
-try:
-    import prometheus_client as prometheus
-except ImportError:
-    prometheus = None
-
 
 class TelegramBridge(Bridge):
     module = "mautrix_telegram"
@@ -57,55 +45,46 @@ class TelegramBridge(Bridge):
     markdown_version = linkified_version
     config_class = Config
     matrix_class = MatrixHandler
-    state_store_class = SQLBridgeStateStore
+    upgrade_table = upgrade_table
 
-    db: 'Engine'
     config: Config
-    session_container: AlchemySessionContainer
-    bot: Bot
+    bot: Bot | None
+    public_website: PublicBridgeWebsite | None
+    provisioning_api: ProvisioningAPI | None
 
     def prepare_db(self) -> None:
-        if not sql:
-            raise RuntimeError("SQLAlchemy is not installed")
-        self.db = sql.create_engine(self.config["appservice.database"],
-                                    **self.config["appservice.database_opts"])
-        Base.metadata.bind = self.db
-        if not self.db.has_table("alembic_version"):
-            self.log.critical("alembic_version table not found. "
-                              "Did you forget to `alembic upgrade head`?")
-            sys.exit(10)
-
+        super().prepare_db()
         init_db(self.db)
-        self.session_container = AlchemySessionContainer(
-            engine=self.db, table_base=Base, session=False,
-            table_prefix="telethon_", manage_tables=False)
 
-    def make_state_store(self) -> None:
-        self.state_store = self.state_store_class(self.get_puppet, self.get_double_puppet)
-
-    def _prepare_website(self, context: Context) -> None:
+    def _prepare_website(self) -> None:
         if self.config["appservice.public.enabled"]:
-            public_website = PublicBridgeWebsite(self.loop)
-            self.az.app.add_subapp(self.config["appservice.public.prefix"], public_website.app)
-            context.public_website = public_website
+            self.public_website = PublicBridgeWebsite(self.loop)
+            self.az.app.add_subapp(
+                self.config["appservice.public.prefix"], self.public_website.app
+            )
+        else:
+            self.public_website = None
 
         if self.config["appservice.provisioning.enabled"]:
-            provisioning_api = ProvisioningAPI(context)
-            self.az.app.add_subapp(self.config["appservice.provisioning.prefix"],
-                                   provisioning_api.app)
-            context.provisioning_api = provisioning_api
+            self.provisioning_api = ProvisioningAPI(self)
+            self.az.app.add_subapp(
+                self.config["appservice.provisioning.prefix"], self.provisioning_api.app
+            )
+        else:
+            self.provisioning_api = None
 
     def prepare_bridge(self) -> None:
-        self.bot = init_bot(self.config)
-        context = Context(self.az, self.config, self.loop, self.session_container, self, self.bot)
-        self._prepare_website(context)
-        self.matrix = context.mx = MatrixHandler(context)
-
-        init_abstract_user(context)
-        init_formatter(context)
-        init_portal(context)
-        self.add_startup_actions(init_puppet(context))
-        self.add_startup_actions(init_user(context))
+        self._prepare_website()
+        AbstractUser.init_cls(self)
+        bot_token: str = self.config["telegram.bot_token"]
+        if bot_token and not bot_token.lower().startswith("disable"):
+            self.bot = AbstractUser.relaybot = Bot(bot_token)
+        else:
+            self.bot = AbstractUser.relaybot = None
+        self.matrix = MatrixHandler(self)
+        Portal.init_cls(self)
+        self.add_startup_actions(Puppet.init_cls(self))
+        self.add_startup_actions(User.init_cls(self))
         if self.bot:
             self.add_startup_actions(self.bot.start())
         if self.config["bridge.resend_bridge_info"]:
@@ -115,7 +94,7 @@ class TelegramBridge(Bridge):
         self.config["bridge.resend_bridge_info"] = False
         self.config.save()
         self.log.info("Re-sending bridge info state event to all portals")
-        for portal in Portal.all():
+        async for portal in Portal.all():
             await portal.update_bridge_info()
         self.log.info("Finished re-sending bridge info state events")
 
@@ -124,19 +103,19 @@ class TelegramBridge(Bridge):
             puppet.stop()
         self.shutdown_actions = (user.stop() for user in User.by_tgid.values())
 
-    async def get_user(self, user_id: UserID, create: bool = True) -> User:
-        user = User.get_by_mxid(user_id, create=create)
+    async def get_user(self, user_id: UserID, create: bool = True) -> User | None:
+        user = await User.get_by_mxid(user_id, create=create)
         if user:
             await user.ensure_started()
         return user
 
-    async def get_portal(self, room_id: RoomID) -> Portal:
-        return Portal.get_by_mxid(room_id)
+    async def get_portal(self, room_id: RoomID) -> Portal | None:
+        return await Portal.get_by_mxid(room_id)
 
-    async def get_puppet(self, user_id: UserID, create: bool = False) -> Puppet:
+    async def get_puppet(self, user_id: UserID, create: bool = False) -> Puppet | None:
         return await Puppet.get_by_mxid(user_id, create=create)
 
-    async def get_double_puppet(self, user_id: UserID) -> Puppet:
+    async def get_double_puppet(self, user_id: UserID) -> Puppet | None:
         return await Puppet.get_by_custom_mxid(user_id)
 
     def is_bridge_ghost(self, user_id: UserID) -> bool:
@@ -145,7 +124,7 @@ class TelegramBridge(Bridge):
     async def count_logged_in_users(self) -> int:
         return len([user for user in User.by_tgid.values() if user.tgid])
 
-    async def manhole_global_namespace(self, user_id: UserID) -> Dict[str, Any]:
+    async def manhole_global_namespace(self, user_id: UserID) -> dict[str, Any]:
         return {
             **await super().manhole_global_namespace(user_id),
             "User": User,