memes_bot/db.py

import aiosqlite
import os

DB_PATH = os.getenv("DB_PATH", "/data/media.db")

# Bump this and add a matching _migrate_to_vN() when the schema changes.
SCHEMA_VERSION = 2


# ── schema / migrations ──────────────────────────────────────────────────────

async def _table_exists(db, name: str) -> bool:
    cursor = await db.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (name,)
    )
    return await cursor.fetchone() is not None


async def _create_v1(db) -> None:
    """The original schema. Fresh installs start here and migrate forward, so
    the migration path is exercised on every new deployment, not just old ones."""
    await db.execute("""
        CREATE TABLE media (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            owner_id INTEGER NOT NULL,
            file_id TEXT NOT NULL,
            file_unique_id TEXT NOT NULL UNIQUE,
            media_type TEXT NOT NULL,
            keywords TEXT NOT NULL,
            added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    await db.execute("CREATE INDEX IF NOT EXISTS idx_owner ON media(owner_id)")


async def _migrate_to_v2(db) -> None:
    """Scope uniqueness to the owner so two users can each save the same file.

    SQLite cannot drop a column constraint in place, so the table is rebuilt.
    Row ids are carried over, which keeps any /delete <id> the user noted valid.
    """
    await db.execute("""
        CREATE TABLE media_v2 (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            owner_id INTEGER NOT NULL,
            file_id TEXT NOT NULL,
            file_unique_id TEXT NOT NULL,
            media_type TEXT NOT NULL,
            keywords TEXT NOT NULL,
            added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE(owner_id, file_unique_id)
        )
    """)
    await db.execute("""
        INSERT INTO media_v2 (id, owner_id, file_id, file_unique_id, media_type, keywords, added_at)
        SELECT id, owner_id, file_id, file_unique_id, media_type, keywords, added_at FROM media
    """)
    await db.execute("DROP TABLE media")
    await db.execute("ALTER TABLE media_v2 RENAME TO media")
    await db.execute("CREATE INDEX IF NOT EXISTS idx_owner ON media(owner_id)")


MIGRATIONS = {
    2: _migrate_to_v2,
}


async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute("PRAGMA user_version")
        version = (await cursor.fetchone())[0]

        # A pre-versioning deployment also reports 0, so tell the two apart by
        # whether the table is already there.
        if version == 0:
            if not await _table_exists(db, "media"):
                await _create_v1(db)
            version = 1
            await db.execute(f"PRAGMA user_version = {version}")

        for target in range(version + 1, SCHEMA_VERSION + 1):
            migrate = MIGRATIONS.get(target)
            if migrate is None:
                continue
            await migrate(db)
            await db.execute(f"PRAGMA user_version = {target}")

        await db.commit()


# ── writes ───────────────────────────────────────────────────────────────────

async def add_media(owner_id: int, file_id: str, file_unique_id: str, media_type: str, keywords: list[str]):
    kw = ",".join(k.strip().lower() for k in keywords if k.strip())
    async with aiosqlite.connect(DB_PATH) as db:
        # Upsert rather than INSERT OR REPLACE: replacing would delete and
        # reinsert the row, handing it a new id and breaking /delete <id>.
        await db.execute(
            """
            INSERT INTO media (owner_id, file_id, file_unique_id, media_type, keywords)
            VALUES (?,?,?,?,?)
            ON CONFLICT(owner_id, file_unique_id) DO UPDATE SET
                file_id = excluded.file_id,
                media_type = excluded.media_type,
                keywords = excluded.keywords
            """,
            (owner_id, file_id, file_unique_id, media_type, kw),
        )
        await db.commit()


async def update_media_keywords(media_id: int, keywords: list[str]) -> bool:
    kw = ",".join(k.strip().lower() for k in keywords if k.strip())
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute(
            "UPDATE media SET keywords = ? WHERE id = ?",
            (kw, media_id),
        )
        await db.commit()
        return cur.rowcount > 0


async def delete_media(media_id: int) -> bool:
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute("DELETE FROM media WHERE id = ?", (media_id,))
        await db.commit()
        return cur.rowcount > 0


# ── reads ────────────────────────────────────────────────────────────────────

COLUMNS = "id, file_id, file_unique_id, media_type, keywords"


async def search_media(query: str, owner_id: int | None = None) -> list[dict]:
    """Search by keyword prefix. owner_id=None searches every owner, which is
    the current shared-library behaviour; public mode will pass a scope."""
    query = query.strip().lower()
    where = []
    params: list = []
    if query:
        where.append("(',' || keywords || ',') LIKE ?")
        params.append(f"%,{query}%")
    if owner_id is not None:
        where.append("owner_id = ?")
        params.append(owner_id)
    clause = f"WHERE {' AND '.join(where)}" if where else ""

    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            f"SELECT {COLUMNS} FROM media {clause} ORDER BY added_at DESC LIMIT 50",
            params,
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]


async def list_media(page: int, page_size: int) -> list[dict]:
    """Return one page of saved media, newest first."""
    offset = page * page_size
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            f"""
            SELECT {COLUMNS}
            FROM media
            ORDER BY added_at DESC, id DESC
            LIMIT ? OFFSET ?
            """,
            (page_size, offset),
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]


async def count_media() -> int:
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute("SELECT COUNT(*) FROM media")
        row = await cursor.fetchone()
        return row[0]


async def get_media_by_id(media_id: int) -> dict | None:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            f"SELECT {COLUMNS} FROM media WHERE id = ?",
            (media_id,),
        )
        row = await cursor.fetchone()
        return dict(row) if row else None