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


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


async def _migrate_to_v4(db) -> None:
    """The vault: posts kept whole, as pointers to the messages the bot already
    holds in its private chat with the owner. Nothing here is ever shared, so it
    lives in its own table rather than as a flag on media."""
    await db.execute("""
        CREATE TABLE IF NOT EXISTS vault_items (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            owner_id INTEGER NOT NULL,
            chat_id INTEGER NOT NULL,
            message_ids TEXT NOT NULL,
            caption TEXT NOT NULL DEFAULT '',
            keywords TEXT NOT NULL DEFAULT '',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    await db.execute("CREATE INDEX IF NOT EXISTS idx_vault_owner ON vault_items(owner_id)")


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


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


# ── vault ────────────────────────────────────────────────────────────────────
#
# Every query here takes an owner_id and filters on it. A vault item belongs to
# one person and is never searchable, never moderated and never federated; the
# separate table is what makes that structural rather than a flag to remember.

VAULT_COLUMNS = "id, owner_id, chat_id, message_ids, caption, keywords, created_at"


def pack_message_ids(message_ids: list[int]) -> str:
    return ",".join(str(m) for m in message_ids)


def unpack_message_ids(packed: str) -> list[int]:
    return [int(m) for m in packed.split(",") if m.strip()]


async def add_vault_item(owner_id: int, chat_id: int, message_ids: list[int],
                         caption: str, keywords: list[str]) -> int:
    kw = ",".join(k.strip().lower() for k in keywords if k.strip())
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute(
            """
            INSERT INTO vault_items (owner_id, chat_id, message_ids, caption, keywords)
            VALUES (?,?,?,?,?)
            """,
            (owner_id, chat_id, pack_message_ids(message_ids), caption or "", kw),
        )
        await db.commit()
        return cur.lastrowid


async def search_vault(owner_id: int, query: str, limit: int = 20) -> list[dict]:
    """Match keywords by prefix or the original caption anywhere in its text —
    the point of a vault is finding a post you never got round to tagging."""
    query = query.strip().lower()
    params: list = [owner_id]
    clause = ""
    if query:
        clause = "AND ((',' || keywords || ',') LIKE ? OR LOWER(caption) LIKE ?)"
        params += [f"%,{query}%", f"%{query}%"]
    params.append(limit)

    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            f"""
            SELECT {VAULT_COLUMNS} FROM vault_items
            WHERE owner_id = ? {clause}
            ORDER BY created_at DESC, id DESC
            LIMIT ?
            """,
            params,
        )
        return [dict(r) for r in await cursor.fetchall()]


async def get_vault_item(owner_id: int, item_id: int) -> dict | None:
    """Scoped by owner so one user cannot fetch another user's saved post."""
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            f"SELECT {VAULT_COLUMNS} FROM vault_items WHERE id = ? AND owner_id = ?",
            (item_id, owner_id),
        )
        row = await cursor.fetchone()
        return dict(row) if row else None


async def delete_vault_item(owner_id: int, item_id: int) -> bool:
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute(
            "DELETE FROM vault_items WHERE id = ? AND owner_id = ?", (item_id, owner_id)
        )
        await db.commit()
        return cur.rowcount > 0


async def count_vault(owner_id: int) -> int:
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            "SELECT COUNT(*) FROM vault_items WHERE owner_id = ?", (owner_id,)
        )
        return (await cursor.fetchone())[0]