Switch from SQLAlchemy to asyncpg/aiosqlite

mautrix_telegram/db/__init__.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,19 +13,23 @@
 #
 # 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 sqlalchemy.engine.base import Engine
+from mautrix.util.async_db import Database
 
-from mautrix.client.state_store.sqlalchemy import UserProfile, RoomState
+from .upgrade import upgrade_table
 
 from .bot_chat import BotChat
 from .message import Message
 from .portal import Portal
 from .puppet import Puppet
 from .telegram_file import TelegramFile
-from .user import User, UserPortal, Contact
+from .user import User
+from .telethon_session import PgSession
 
 
-def init(db_engine: Engine) -> None:
-    for table in (Portal, Message, User, Contact, UserPortal, Puppet, TelegramFile, UserProfile,
-                  RoomState, BotChat):
-        table.bind(db_engine)
+def init(db: Database) -> None:
+    for table in (Portal, Message, User, Puppet, TelegramFile, BotChat, PgSession):
+        table.db = db
+
+
+__all__ = ["upgrade_table", "init", "Portal", "Message", "User", "Puppet", "TelegramFile",
+           "BotChat", "PgSession"]