database.py

from pathlib import Path

import aiosqlite


BASE_DIR = Path(__file__).resolve().parent
DB_PATH = BASE_DIR / "data" / "annotator.db"


async def _ensure_column(
    db: aiosqlite.Connection,
    table: str,
    column: str,
    definition: str,
) -> None:
    cursor = await db.execute(f"PRAGMA table_info({table})")
    columns = {row[1] for row in await cursor.fetchall()}
    if column not in columns:
        await db.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")


async def init_db() -> None:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)

    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("PRAGMA foreign_keys = ON")
        await db.executescript(
            """
            CREATE TABLE IF NOT EXISTS sources (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT,
                file_path TEXT,
                created_at DATETIME,
                source_type TEXT DEFAULT 'md',
                last_viewed_block_index INTEGER DEFAULT 0,
                last_viewed_at TEXT,
                completed_at TEXT
            );

            CREATE TABLE IF NOT EXISTS blocks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                source_id INTEGER NOT NULL,
                block_index INTEGER,
                content TEXT,
                speaker TEXT,
                message_timestamp TEXT,
                message_id TEXT,
                source TEXT,
                conversation_id TEXT,
                conversation_name TEXT,
                FOREIGN KEY (source_id) REFERENCES sources(id)
                    ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS annotations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                source_id INTEGER NOT NULL,
                block_id INTEGER NOT NULL,
                parent_id INTEGER,
                start_offset INTEGER,
                end_offset INTEGER,
                quote TEXT,
                content TEXT,
                created_at DATETIME,
                archived_at TEXT,
                FOREIGN KEY (source_id) REFERENCES sources(id)
                    ON DELETE CASCADE,
                FOREIGN KEY (block_id) REFERENCES blocks(id)
                    ON DELETE CASCADE,
                FOREIGN KEY (parent_id) REFERENCES annotations(id)
                    ON DELETE CASCADE
            );
            """
        )
        await _ensure_column(db, "sources", "source_type", "TEXT DEFAULT 'md'")
        await _ensure_column(
            db,
            "sources",
            "last_viewed_block_index",
            "INTEGER DEFAULT 0",
        )
        await _ensure_column(db, "sources", "last_viewed_at", "TEXT")
        await _ensure_column(db, "sources", "completed_at", "TEXT")
        await _ensure_column(db, "blocks", "speaker", "TEXT")
        await _ensure_column(db, "blocks", "message_timestamp", "TEXT")
        await _ensure_column(db, "blocks", "message_id", "TEXT")
        await _ensure_column(db, "blocks", "source", "TEXT")
        await _ensure_column(db, "blocks", "conversation_id", "TEXT")
        await _ensure_column(db, "blocks", "conversation_name", "TEXT")
        await _ensure_column(db, "annotations", "archived_at", "TEXT")
        await db.commit()