Update standalone mode to asyncpg/aiosqlite

maubot/standalone/database.py

@@ -15,19 +15,41 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from __future__ import annotations
 
-import sqlalchemy as sql
+from attr import dataclass
 
+from mautrix.client import SyncStore
 from mautrix.types import FilterID, SyncToken, UserID
-from mautrix.util.db import Base
+from mautrix.util.async_db import Database
+
+find_q = "SELECT next_batch, filter_id FROM standalone_next_batch WHERE user_id=$1"
+insert_q = "INSERT INTO standalone_next_batch (user_id, next_batch, filter_id) VALUES ($1, $2, $3)"
+update_nb_q = "UPDATE standalone_next_batch SET next_batch=$1 WHERE user_id=$2"
+update_filter_q = "UPDATE standalone_next_batch SET filter_id=$1 WHERE user_id=$2"
 
 
-class NextBatch(Base):
-    __tablename__ = "standalone_next_batch"
+@dataclass
+class NextBatch(SyncStore):
+    db: Database
+    user_id: UserID
+    next_batch: SyncToken = ""
+    filter_id: FilterID = ""
 
-    user_id: UserID = sql.Column(sql.String(255), primary_key=True)
-    next_batch: SyncToken = sql.Column(sql.String(255))
-    filter_id: FilterID = sql.Column(sql.String(255))
+    async def load(self) -> NextBatch:
+        row = await self.db.fetchrow(find_q, self.user_id)
+        if row is not None:
+            self.next_batch = row["next_batch"]
+            self.filter_id = row["filter_id"]
+        else:
+            await self.db.execute(insert_q, self.user_id, self.next_batch, self.filter_id)
+        return self
 
-    @classmethod
-    def get(cls, user_id: UserID) -> NextBatch | None:
-        return cls._select_one_or_none(cls.c.user_id == user_id)
+    async def put_filter_id(self, filter_id: FilterID) -> None:
+        self.filter_id = filter_id
+        await self.db.execute(update_filter_q, self.filter_id, self.user_id)
+
+    async def put_next_batch(self, next_batch: SyncToken) -> None:
+        self.next_batch = next_batch
+        await self.db.execute(update_nb_q, self.next_batch, self.user_id)
+
+    async def get_next_batch(self) -> SyncToken:
+        return self.next_batch