From f74a67dd793a1799449846315659cc11a467c6cf Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 6 May 2022 16:18:45 +0300 Subject: [PATCH] Store instance database engine in database --- maubot/db/__init__.py | 4 +- maubot/db/instance.py | 40 +++++++++++--- maubot/db/upgrade/__init__.py | 2 +- .../upgrade/v02_instance_database_engine.py | 25 +++++++++ maubot/instance.py | 52 ++++++++++++++++--- maubot/loader/meta.py | 6 ++- .../frontend/src/pages/dashboard/Instance.js | 3 +- 7 files changed, 113 insertions(+), 19 deletions(-) create mode 100644 maubot/db/upgrade/v02_instance_database_engine.py diff --git a/maubot/db/__init__.py b/maubot/db/__init__.py index d6aeb09..68833ce 100644 --- a/maubot/db/__init__.py +++ b/maubot/db/__init__.py @@ -1,7 +1,7 @@ from mautrix.util.async_db import Database from .client import Client -from .instance import Instance +from .instance import DatabaseEngine, Instance from .upgrade import upgrade_table @@ -10,4 +10,4 @@ def init(db: Database) -> None: table.db = db -__all__ = ["upgrade_table", "init", "Client", "Instance"] +__all__ = ["upgrade_table", "init", "Client", "Instance", "DatabaseEngine"] diff --git a/maubot/db/instance.py b/maubot/db/instance.py index dff7064..5bb3f6a 100644 --- a/maubot/db/instance.py +++ b/maubot/db/instance.py @@ -16,6 +16,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, ClassVar +from enum import Enum from asyncpg import Record from attr import dataclass @@ -26,6 +27,11 @@ from mautrix.util.async_db import Database fake_db = Database.create("") if TYPE_CHECKING else None +class DatabaseEngine(Enum): + SQLITE = "sqlite" + POSTGRES = "postgres" + + @dataclass class Instance: db: ClassVar[Database] = fake_db @@ -35,21 +41,31 @@ class Instance: enabled: bool primary_user: UserID config_str: str + database_engine: DatabaseEngine | None + + @property + def database_engine_str(self) -> str | None: + return self.database_engine.value if self.database_engine else None @classmethod def _from_row(cls, row: Record | None) -> Instance | None: if row is None: return None - return cls(**row) + data = {**row} + db_engine = data.pop("database_engine", None) + return cls(**data, database_engine=DatabaseEngine(db_engine) if db_engine else None) + + _columns = "id, type, enabled, primary_user, config, database_engine" @classmethod async def all(cls) -> list[Instance]: - rows = await cls.db.fetch("SELECT id, type, enabled, primary_user, config FROM instance") + q = f"SELECT {cls._columns} FROM instance" + rows = await cls.db.fetch(q) return [cls._from_row(row) for row in rows] @classmethod async def get(cls, id: str) -> Instance | None: - q = "SELECT id, type, enabled, primary_user, config FROM instance WHERE id=$1" + q = f"SELECT {cls._columns} FROM instance WHERE id=$1" return cls._from_row(await cls.db.fetchrow(q, id)) async def update_id(self, new_id: str) -> None: @@ -58,17 +74,27 @@ class Instance: @property def _values(self): - return self.id, self.type, self.enabled, self.primary_user, self.config_str + return ( + self.id, + self.type, + self.enabled, + self.primary_user, + self.config_str, + self.database_engine_str, + ) async def insert(self) -> None: q = ( - "INSERT INTO instance (id, type, enabled, primary_user, config) " - "VALUES ($1, $2, $3, $4, $5)" + "INSERT INTO instance (id, type, enabled, primary_user, config, database_engine) " + "VALUES ($1, $2, $3, $4, $5, $6)" ) await self.db.execute(q, *self._values) async def update(self) -> None: - q = "UPDATE instance SET type=$2, enabled=$3, primary_user=$4, config=$5 WHERE id=$1" + q = """ + UPDATE instance SET type=$2, enabled=$3, primary_user=$4, config=$5, database_engine=$6 + WHERE id=$1 + """ await self.db.execute(q, *self._values) async def delete(self) -> None: diff --git a/maubot/db/upgrade/__init__.py b/maubot/db/upgrade/__init__.py index 146e713..ed96422 100644 --- a/maubot/db/upgrade/__init__.py +++ b/maubot/db/upgrade/__init__.py @@ -2,4 +2,4 @@ from mautrix.util.async_db import UpgradeTable upgrade_table = UpgradeTable() -from . import v01_initial_revision +from . import v01_initial_revision, v02_instance_database_engine diff --git a/maubot/db/upgrade/v02_instance_database_engine.py b/maubot/db/upgrade/v02_instance_database_engine.py new file mode 100644 index 0000000..7d2d7e7 --- /dev/null +++ b/maubot/db/upgrade/v02_instance_database_engine.py @@ -0,0 +1,25 @@ +# maubot - A plugin-based Matrix bot system. +# Copyright (C) 2022 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 +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +from __future__ import annotations + +from mautrix.util.async_db import Connection + +from . import upgrade_table + + +@upgrade_table.register(description="Store instance database engine") +async def upgrade_v2(conn: Connection) -> None: + await conn.execute("ALTER TABLE instance ADD COLUMN database_engine TEXT") diff --git a/maubot/instance.py b/maubot/instance.py index b9b1c23..4d797e4 100644 --- a/maubot/instance.py +++ b/maubot/instance.py @@ -34,7 +34,7 @@ from mautrix.util.config import BaseProxyConfig, RecursiveDict from mautrix.util.logging import TraceLogger from .client import Client -from .db import Instance as DBInstance +from .db import DatabaseEngine, Instance as DBInstance from .lib.plugin_db import ProxyPostgresDatabase from .loader import DatabaseType, PluginLoader, ZippedPluginLoader from .plugin_base import Plugin @@ -71,10 +71,21 @@ class PluginInstance(DBInstance): started: bool def __init__( - self, id: str, type: str, enabled: bool, primary_user: UserID, config: str = "" + self, + id: str, + type: str, + enabled: bool, + primary_user: UserID, + config: str = "", + database_engine: DatabaseEngine | None = None, ) -> None: super().__init__( - id=id, type=type, enabled=bool(enabled), primary_user=primary_user, config_str=config + id=id, + type=type, + enabled=bool(enabled), + primary_user=primary_user, + config_str=config, + database_engine=database_engine, ) def __hash__(self) -> int: @@ -111,6 +122,8 @@ class PluginInstance(DBInstance): "database": ( self.inst_db is not None and self.maubot.config["api_features.instance_database"] ), + "database_interface": self.loader.meta.database_type_str if self.loader else "unknown", + "database_engine": self.database_engine_str, } def _introspect_sqlalchemy(self) -> dict: @@ -269,12 +282,27 @@ class PluginInstance(DBInstance): self, upgrade_table: UpgradeTable | None = None, actually_start: bool = True ) -> None: if self.loader.meta.database_type == DatabaseType.SQLALCHEMY: + if self.database_engine is None: + await self.update_db_engine(DatabaseEngine.SQLITE) + elif self.database_engine == DatabaseEngine.POSTGRES: + raise RuntimeError( + "Instance database engine is marked as Postgres, but plugin uses legacy " + "database interface, which doesn't support postgres." + ) self.inst_db = sql.create_engine(f"sqlite:///{self._sqlite_db_path}") elif self.loader.meta.database_type == DatabaseType.ASYNCPG: + if self.database_engine is None: + if os.path.exists(self._sqlite_db_path) or not self.maubot.plugin_postgres_db: + await self.update_db_engine(DatabaseEngine.SQLITE) + else: + await self.update_db_engine(DatabaseEngine.POSTGRES) instance_db_log = db_log.getChild(self.id) - # TODO should there be a way to choose between SQLite and Postgres - # for individual instances? Maybe checking the existence of the SQLite file. - if self.maubot.plugin_postgres_db: + if self.database_engine == DatabaseEngine.POSTGRES: + if not self.maubot.plugin_postgres_db: + raise RuntimeError( + "Instance database engine is marked as Postgres, but this maubot isn't " + "configured to support Postgres for plugin databases" + ) self.inst_db = ProxyPostgresDatabase( pool=self.maubot.plugin_postgres_db, instance_id=self.id, @@ -334,7 +362,12 @@ class PluginInstance(DBInstance): self.log.debug("Disabling webapp after plugin meta reload") self.disable_webapp() if self.loader.meta.database: - await self.start_database(cls.get_db_upgrade_table()) + try: + await self.start_database(cls.get_db_upgrade_table()) + except Exception: + self.log.exception("Failed to start instance database") + await self.update_enabled(False) + return config_class = cls.get_config_class() if config_class: try: @@ -455,6 +488,11 @@ class PluginInstance(DBInstance): self.enabled = enabled await self.update() + async def update_db_engine(self, db_engine: DatabaseEngine | None) -> None: + if db_engine is not None and db_engine != self.database_engine: + self.database_engine = db_engine + await self.update() + @classmethod @async_getter_lock async def get( diff --git a/maubot/loader/meta.py b/maubot/loader/meta.py index f16937b..d368e24 100644 --- a/maubot/loader/meta.py +++ b/maubot/loader/meta.py @@ -13,7 +13,7 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -from typing import List +from typing import List, Optional from attr import dataclass from packaging.version import InvalidVersion, Version @@ -63,3 +63,7 @@ class PluginMeta(SerializableAttrs): extra_files: List[str] = [] dependencies: List[str] = [] soft_dependencies: List[str] = [] + + @property + def database_type_str(self) -> Optional[str]: + return self.database_type.value if self.database else None diff --git a/maubot/management/frontend/src/pages/dashboard/Instance.js b/maubot/management/frontend/src/pages/dashboard/Instance.js index c487fb0..049c5e5 100644 --- a/maubot/management/frontend/src/pages/dashboard/Instance.js +++ b/maubot/management/frontend/src/pages/dashboard/Instance.js @@ -43,7 +43,7 @@ class Instance extends BaseMainView { } get entryKeys() { - return ["id", "primary_user", "enabled", "started", "type", "config"] + return ["id", "primary_user", "enabled", "started", "type", "config", "database_engine"] } get initialState() { @@ -54,6 +54,7 @@ class Instance extends BaseMainView { started: true, type: "", config: "", + database_engine: "", saving: false, deleting: false,