refactor: isolate database migrations and model defaults

2026-07-27 17:51:16 +03:00
parent ab029a4c67
commit c20d893255
7 changed files with 238 additions and 311 deletions
--- a/assistant_bot/migrations.py
+++ b/assistant_bot/migrations.py
@@ -0,0 +1,113 @@
+import sqlite3
+
+
+LATEST_SCHEMA_VERSION = 2
+
+INITIAL_SCHEMA = """
+    CREATE TABLE IF NOT EXISTS user_settings (
+        user_id INTEGER PRIMARY KEY,
+        ollama_model TEXT,
+        updated_at TEXT NOT NULL
+    );
+
+    CREATE TABLE IF NOT EXISTS authorized_users (
+        user_id INTEGER PRIMARY KEY,
+        authorized_at TEXT NOT NULL
+    );
+
+    CREATE TABLE IF NOT EXISTS memories (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        user_id INTEGER NOT NULL,
+        text TEXT NOT NULL,
+        created_at TEXT NOT NULL
+    );
+
+    CREATE TABLE IF NOT EXISTS notes (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        user_id INTEGER NOT NULL,
+        text TEXT NOT NULL,
+        created_at TEXT NOT NULL
+    );
+
+    CREATE TABLE IF NOT EXISTS reminders (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        user_id INTEGER NOT NULL,
+        chat_id INTEGER NOT NULL,
+        text TEXT NOT NULL,
+        remind_at TEXT NOT NULL,
+        status TEXT NOT NULL DEFAULT 'pending',
+        created_at TEXT NOT NULL,
+        sent_at TEXT
+    );
+
+    CREATE TABLE IF NOT EXISTS tracked_items (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        user_id INTEGER NOT NULL,
+        title TEXT NOT NULL,
+        status TEXT NOT NULL,
+        created_at TEXT NOT NULL,
+        updated_at TEXT NOT NULL
+    );
+
+    CREATE TABLE IF NOT EXISTS conversation_messages (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        user_id INTEGER NOT NULL,
+        chat_id INTEGER NOT NULL,
+        role TEXT NOT NULL CHECK(role IN ('user', 'assistant')),
+        content TEXT NOT NULL,
+        created_at TEXT NOT NULL
+    );
+
+    CREATE TABLE IF NOT EXISTS conversation_contexts (
+        user_id INTEGER NOT NULL,
+        chat_id INTEGER NOT NULL,
+        started_after_id INTEGER NOT NULL DEFAULT 0,
+        updated_at TEXT NOT NULL,
+        PRIMARY KEY (user_id, chat_id)
+    );
+
+    CREATE INDEX IF NOT EXISTS idx_memories_user_created
+        ON memories(user_id, created_at DESC);
+    CREATE INDEX IF NOT EXISTS idx_notes_user_created
+        ON notes(user_id, created_at DESC);
+    CREATE INDEX IF NOT EXISTS idx_reminders_due
+        ON reminders(status, remind_at);
+    CREATE INDEX IF NOT EXISTS idx_tracked_user_updated
+        ON tracked_items(user_id, updated_at DESC);
+    CREATE INDEX IF NOT EXISTS idx_conversation_chat_id
+        ON conversation_messages(user_id, chat_id, id DESC);
+"""
+
+
+def _schema_version(connection: sqlite3.Connection) -> int:
+    row = connection.execute("PRAGMA user_version").fetchone()
+    return int(row[0])
+
+
+def _has_column(
+    connection: sqlite3.Connection,
+    table: str,
+    column: str,
+) -> bool:
+    return any(
+        str(row["name"]) == column
+        for row in connection.execute(f"PRAGMA table_info({table})")
+    )
+
+
+def migrate_database(connection: sqlite3.Connection) -> None:
+    """Bring a new or existing database to the latest known schema."""
+    connection.execute("PRAGMA journal_mode=WAL")
+    version = _schema_version(connection)
+
+    if version < 1:
+        connection.executescript(INITIAL_SCHEMA)
+        connection.execute("PRAGMA user_version = 1")
+        version = 1
+
+    if version < 2:
+        if not _has_column(connection, "user_settings", "yandex_model"):
+            connection.execute(
+                "ALTER TABLE user_settings ADD COLUMN yandex_model TEXT"
+            )
+        connection.execute("PRAGMA user_version = 2")