From e9e325a8717b0be8515c41892afd8cf941fc41b1 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 27 Mar 2022 20:27:54 +0300
Subject: [PATCH] Create standalone_next_batch table automatically

---
 maubot/standalone/__main__.py |  3 ++-
 maubot/standalone/database.py | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/maubot/standalone/__main__.py b/maubot/standalone/__main__.py
index 3e9c3aa..c799fcf 100644
--- a/maubot/standalone/__main__.py
+++ b/maubot/standalone/__main__.py
@@ -49,7 +49,7 @@ from ..plugin_base import Plugin
 from ..plugin_server import PluginWebApp, PrefixResource
 from ..server import AccessLogger
 from .config import Config
-from .database import NextBatch
+from .database import NextBatch, upgrade_table
 from .loader import FileSystemLoader
 
 crypto_import_error = None
@@ -124,6 +124,7 @@ db = Database.create(
     db_args=config.get("database_opts", {}),
     ignore_foreign_tables=True,
     log=logging.getLogger("maubot.db"),
+    upgrade_table=upgrade_table,
 )
 
 user_id = config["user.credentials.id"]
diff --git a/maubot/standalone/database.py b/maubot/standalone/database.py
index e69a220..117634b 100644
--- a/maubot/standalone/database.py
+++ b/maubot/standalone/database.py
@@ -15,11 +15,30 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from __future__ import annotations
 
+import logging
+
 from attr import dataclass
 
 from mautrix.client import SyncStore
 from mautrix.types import FilterID, SyncToken, UserID
-from mautrix.util.async_db import Database
+from mautrix.util.async_db import Connection, Database, UpgradeTable
+
+upgrade_table = UpgradeTable(
+    version_table_name="standalone_version", log=logging.getLogger("maubot.db")
+)
+
+
+@upgrade_table.register(description="Initial revision")
+async def upgrade_v1(conn: Connection) -> None:
+    await conn.execute(
+        """CREATE TABLE IF NOT EXISTS standalone_next_batch (
+            user_id    TEXT NOT NULL,
+            next_batch TEXT,
+            filter_id  TEXT,
+            PRIMARY KEY (user_id)
+        )"""
+    )
+
 
 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)"