Add schema migrations and scope uniqueness to the owner
Step 0 of the roadmap: everything from the moderation queue onward adds columns, and init_db could only ever run CREATE TABLE IF NOT EXISTS, so it could not alter a live deployment. - Track schema state in PRAGMA user_version and run ordered migrations. Fresh installs create the v1 table and migrate forward, so the upgrade path runs on every new deployment rather than only on old ones. - Rebuild media with UNIQUE(owner_id, file_unique_id). The old global unique constraint stopped two users from each saving the same file, which blocks public mode. - Replace INSERT OR REPLACE with an upsert. Replacing deleted and reinserted the row, giving it a new id and silently invalidating any id the user had noted from /list. - Give search_media an optional owner_id scope, defaulting to the current shared-library behaviour. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
goyban
@@ -3,11 +3,24 @@ 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
|
||||
|
||||
async def init_db():
|
||||
async with aiosqlite.connect(DB_PATH) as db:
|
||||
|
||||
# ── 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 IF NOT EXISTS media (
|
||||
CREATE TABLE media (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
owner_id INTEGER NOT NULL,
|
||||
file_id TEXT NOT NULL,
|
||||
@@ -18,42 +31,126 @@ async def init_db():
|
||||
)
|
||||
""")
|
||||
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 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
|
||||
""",
|
||||
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
|
||||
|
||||
+108
@@ -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())
|
||||
Reference in New Issue
Block a user