114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
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")
|