diff --git a/db.py b/db.py index 4d0b297..863180b 100644 --- a/db.py +++ b/db.py @@ -3,58 +3,155 @@ 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 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: - await db.execute(""" - CREATE TABLE IF NOT EXISTS 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)") + 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 . await db.execute( - "INSERT OR REPLACE INTO media (owner_id, file_id, file_unique_id, media_type, keywords) VALUES (?,?,?,?,?)", + """ + 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 search_media(query: str) -> list[dict]: +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 - if query: - cursor = await db.execute( - """ - SELECT id, file_id, file_unique_id, media_type, keywords - FROM media - WHERE (',' || keywords || ',') LIKE ? - ORDER BY added_at DESC - LIMIT 50 - """, - (f"%,{query}%",), - ) - else: - cursor = await db.execute( - """ - SELECT id, file_id, file_unique_id, media_type, keywords - FROM media - ORDER BY added_at DESC - LIMIT 50 - """, - ) + 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] @@ -65,8 +162,8 @@ async def list_media(page: int, page_size: int) -> list[dict]: async with aiosqlite.connect(DB_PATH) as db: db.row_factory = aiosqlite.Row cursor = await db.execute( - """ - SELECT id, file_id, file_unique_id, media_type, keywords + f""" + SELECT {COLUMNS} FROM media ORDER BY added_at DESC, id DESC LIMIT ? OFFSET ? @@ -88,26 +185,8 @@ 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( - "SELECT id, file_id, file_unique_id, media_type, keywords FROM media WHERE id = ?", + f"SELECT {COLUMNS} FROM media WHERE id = ?", (media_id,), ) row = await cursor.fetchone() return dict(row) if row else None - - -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 diff --git a/test_db.py b/test_db.py new file mode 100644 index 0000000..fc1c41d --- /dev/null +++ b/test_db.py @@ -0,0 +1,108 @@ +"""Migration tests. Run with: python test_db.py + +The v1->v2 migration rebuilds the media table, so it drops the original. +These cases cover the paths that rebuild touches before it ever runs against +a real bot_db/media.db. +""" +import asyncio +import importlib +import os +import sqlite3 +import sys +import tempfile + +V1_SCHEMA = """ + 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 + ); +""" + + +def user_version(path: str) -> int: + with sqlite3.connect(path) as c: + return c.execute("PRAGMA user_version").fetchone()[0] + + +def load_db(path: str): + os.environ["DB_PATH"] = path + import db + return importlib.reload(db) + + +async def test_legacy_upgrade(tmp: str): + path = os.path.join(tmp, "legacy.db") + with sqlite3.connect(path) as c: + c.executescript(V1_SCHEMA) + c.execute("INSERT INTO media (id,owner_id,file_id,file_unique_id,media_type,keywords)" + " VALUES (7,111,'FID_A','UNI_A','photo','cat,funny')") + c.execute("INSERT INTO media (id,owner_id,file_id,file_unique_id,media_type,keywords)" + " VALUES (9,111,'FID_B','UNI_B','gif','dog')") + + db = load_db(path) + await db.init_db() + + assert user_version(path) == 2 + rows = sorted((r["id"], r["keywords"]) for r in await db.search_media("")) + assert rows == [(7, "cat,funny"), (9, "dog")], f"rows not preserved: {rows}" + print(" legacy v1 database upgrades to v2 with ids and keywords intact") + + +async def test_per_owner_uniqueness(tmp: str): + path = os.path.join(tmp, "owners.db") + db = load_db(path) + await db.init_db() + + await db.add_media(111, "FID_A", "UNI_A", "photo", ["cat"]) + await db.add_media(222, "FID_A2", "UNI_A", "photo", ["cat"]) + assert len(await db.search_media("cat")) == 2, "two owners must each keep their own row" + assert len(await db.search_media("", owner_id=222)) == 1, "owner scoping failed" + print(" two owners can save the same file; owner scoping filters correctly") + + +async def test_readd_keeps_id(tmp: str): + path = os.path.join(tmp, "readd.db") + db = load_db(path) + await db.init_db() + + await db.add_media(111, "FID_A", "UNI_A", "photo", ["cat"]) + original = (await db.search_media("cat"))[0]["id"] + await db.add_media(111, "FID_NEW", "UNI_A", "photo", ["kitten"]) + row = await db.get_media_by_id(original) + + assert row is not None, "re-adding the same file orphaned its id" + assert row["keywords"] == "kitten" and row["file_id"] == "FID_NEW" + print(" re-adding a file updates in place and keeps its id") + + +async def test_fresh_and_idempotent(tmp: str): + path = os.path.join(tmp, "fresh.db") + db = load_db(path) + await db.init_db() + + with sqlite3.connect(path) as c: + sql = c.execute("SELECT sql FROM sqlite_master WHERE name='media'").fetchone()[0] + assert "UNIQUE(owner_id, file_unique_id)" in sql + assert user_version(path) == 2 + + await db.init_db() + await db.init_db() + print(" fresh install lands on v2; init_db is idempotent") + + +async def main(): + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + with tempfile.TemporaryDirectory() as tmp: + for test in (test_legacy_upgrade, test_per_owner_uniqueness, + test_readd_keeps_id, test_fresh_and_idempotent): + await test(tmp) + print("\nAll database tests passed.") + + +if __name__ == "__main__": + asyncio.run(main())