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 = 3


# ── 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)")


async def _migrate_to_v3(db) -> None:
    """Add the moderation queue. Existing rows default to approved, so an
    upgraded deployment behaves exactly as it did before."""
    await db.execute("ALTER TABLE media ADD COLUMN status TEXT NOT NULL DEFAULT 'approved'")
    await db.execute("ALTER TABLE media ADD COLUMN reviewed_by INTEGER")
    await db.execute("ALTER TABLE media ADD COLUMN reviewed_at TIMESTAMP")
    await db.execute("CREATE INDEX IF NOT EXISTS idx_status ON media(status)")


MIGRATIONS = {
    2: _migrate_to_v2,
    3: _migrate_to_v3,
}


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], status: str = "approved"):
    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, status)
            VALUES (?,?,?,?,?,?)
            ON CONFLICT(owner_id, file_unique_id) DO UPDATE SET
                file_id = excluded.file_id,
                media_type = excluded.media_type,
                keywords = excluded.keywords,
                status = excluded.status
            """,
            (owner_id, file_id, file_unique_id, media_type, kw, status),
        )
        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, owner_id, file_id, file_unique_id, media_type, keywords, status"


async def search_media(query: str, owner_id: int | None = None,
                       status: str | None = "approved") -> 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.
    status=None includes items still awaiting moderation."""
    query = query.strip().lower()
    where = []
    params: list = []
    if status is not None:
        where.append("status = ?")
        params.append(status)
    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, status: str | None = "approved") -> list[dict]:
    """Return one page of saved media, newest first."""
    offset = page * page_size
    clause = "WHERE status = ?" if status is not None else ""
    params = ([status] if status is not None else []) + [page_size, offset]
    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, id DESC
            LIMIT ? OFFSET ?
            """,
            params,
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]


async def count_media(status: str | None = "approved") -> int:
    clause = "WHERE status = ?" if status is not None else ""
    params = [status] if status is not None else []
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(f"SELECT COUNT(*) FROM media {clause}", params)
        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


# ── moderation queue ─────────────────────────────────────────────────────────

async def next_pending() -> dict | None:
    """Oldest item awaiting review. The queue keeps no cursor: each call just
    asks for the next one, so a moderator can stop and resume at any time and
    several moderators can work at once."""
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            f"SELECT {COLUMNS} FROM media WHERE status = 'pending' ORDER BY added_at ASC, id ASC LIMIT 1"
        )
        row = await cursor.fetchone()
        return dict(row) if row else None


async def count_pending() -> int:
    return await count_media(status="pending")


async def set_status(media_id: int, status: str, reviewer_id: int) -> bool:
    """Approve or reject a pending item. Returns False if it was not pending
    any more, which is how two moderators reviewing at once are kept apart."""
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute(
            """
            UPDATE media SET status = ?, reviewed_by = ?, reviewed_at = CURRENT_TIMESTAMP
            WHERE id = ? AND status = 'pending'
            """,
            (status, reviewer_id, media_id),
        )
        await db.commit()
        return cur.rowcount > 0