Switch to asyncpg

gitlab_matrix/bot.py

@@ -17,23 +17,25 @@
 from typing import Type
 
 from mautrix.util.config import BaseProxyConfig
+from mautrix.util.async_db import UpgradeTable
 from maubot import Plugin
 
-from .db import Database
+from .db import DBManager
 from .util import Config
 from .webhook import GitlabWebhook
 from .commands import GitlabCommands
+from .migrations import upgrade_table
 
 
 class GitlabBot(Plugin):
-    db: Database
+    db: DBManager
     webhook: GitlabWebhook
     commands: GitlabCommands
 
     async def start(self) -> None:
         self.config.load_and_update()
 
-        self.db = Database(self.database)
+        self.db = DBManager(self.database)
         self.webhook = await GitlabWebhook(self).start()
         self.commands = GitlabCommands(self)
 
@@ -46,3 +48,7 @@ async def stop(self) -> None:
     @classmethod
     def get_config_class(cls) -> Type[BaseProxyConfig]:
         return Config
+
+    @classmethod
+    def get_db_upgrade_table(cls) -> UpgradeTable:
+        return upgrade_table

gitlab_matrix/commands/alias.py

@@ -29,18 +29,20 @@ async def alias(self) -> None:
     @command.argument("url", "server URL")
     @command.argument("alias", "server alias")
     async def alias_add(self, evt: MessageEvent, url: str, alias: str) -> None:
-        if url not in self.bot.db.get_servers(evt.sender):
+        servers = await self.bot.db.get_servers(evt.sender)
+        if url not in servers:
             await evt.reply("You can't add an alias to a GitLab server you are not logged in to.")
             return
-        if self.bot.db.has_alias(evt.sender, alias):
+        has_alias = await self.bot.db.has_alias(evt.sender, alias)
+        if has_alias:
             await evt.reply("Alias already in use.")
             return
-        self.bot.db.add_alias(evt.sender, url, alias)
+        await self.bot.db.add_alias(evt.sender, url, alias)
         await evt.reply(f"Added alias {alias} to server {url}")
 
     @alias.subcommand("list", aliases=("l", "ls"), help="Show your Gitlab server aliases.")
     async def alias_list(self, evt: MessageEvent) -> None:
-        aliases = self.bot.db.get_aliases(evt.sender)
+        aliases = await self.bot.db.get_aliases(evt.sender)
         if not aliases:
             await evt.reply("You don't have any aliases.")
             return
@@ -52,5 +54,5 @@ async def alias_list(self, evt: MessageEvent) -> None:
                       help="Remove a alias to a Gitlab server.")
     @command.argument("alias", "server alias")
     async def alias_rm(self, evt: MessageEvent, alias: str) -> None:
-        self.bot.db.rm_alias(evt.sender, alias)
+        await self.bot.db.rm_alias(evt.sender, alias)
         await evt.reply(f"Removed alias {alias}.")

gitlab_matrix/commands/room.py

@@ -48,5 +48,5 @@ async def default_repo(self, evt: MessageEvent, repo: str, gl: Gl) -> None:
                 await evt.reply(f"Couldn't find {repo} on {gl.url}")
                 return
             raise
-        self.bot.db.set_default_repo(evt.room_id, gl.url, repo)
+        await self.bot.db.set_default_repo(evt.room_id, gl.url, repo)
         await evt.reply(f"Changed the default repo to {repo} on {gl.url}")

gitlab_matrix/commands/server.py

@@ -32,12 +32,12 @@ async def server(self) -> None:
     @server.subcommand("default", aliases=("d",), help="Change your default GitLab server.")
     @command.argument("url", "server URL")
     async def server_default(self, evt: MessageEvent, url: str) -> None:
-        self.bot.db.change_default(evt.sender, url)
+        await self.bot.db.change_default(evt.sender, url)
         await evt.reply(f"Changed the default server to {url}")
 
     @server.subcommand("list", aliases=("ls",), help="Show your GitLab servers.")
     async def server_list(self, evt: MessageEvent) -> None:
-        servers = self.bot.db.get_servers(evt.sender)
+        servers = await self.bot.db.get_servers(evt.sender)
         if not servers:
             await evt.reply("You are not logged in to any server.")
             return
@@ -60,13 +60,13 @@ async def server_login(self, evt: MessageEvent, url: str, token: str) -> None:
                              exc_info=True)
             await evt.reply(f"GitLab login failed: {e}")
             return
-        self.bot.db.add_login(evt.sender, url, token)
+        await self.bot.db.add_login(evt.sender, url, token)
         await evt.reply(f"Successfully logged into GitLab at {url} as {gl.user.name}")
 
     @server.subcommand("logout", help="Remove the access token from the bot's database.")
     @command.argument("url", "server URL")
     async def server_logout(self, evt: MessageEvent, url: str) -> None:
-        self.bot.db.rm_login(evt.sender, url)
+        await self.bot.db.rm_login(evt.sender, url)
         await evt.reply(f"Removed {url} from the database.")
 
     @Command.gitlab.subcommand("ping", aliases=("p",), help="Ping the bot.")

gitlab_matrix/commands/webhook.py

@@ -34,7 +34,7 @@ async def webhook(self) -> None:
     @with_gitlab_session
     async def webhook_add(self, evt: MessageEvent, repo: str, gl: Gl) -> None:
         token = secrets.token_urlsafe(64)
-        self.bot.db.add_webhook_room(token, evt.room_id)
+        await self.bot.db.add_webhook_room(token, evt.room_id)
         project = gl.projects.get(repo)
         hook = project.hooks.create({
             "url": f"{self.bot.webapp_url}/webhooks",