fix: backups include WAL data and restores are integrity-checked
The database runs in WAL mode, so a shutil.copy2 of prompts.db silently
dropped every committed row still in prompts.db-wal; the backup
download and the pre-restore safety copy both did that. Restore also
overwrote the live file with connections open, read the whole upload
into memory before checking the 100 MB cap and never ran
PRAGMA integrity_check.
PromptModel.backup_database now uses the SQLite online backup API,
verify_database_file opens a candidate read-only and checks header,
integrity and the prompts table, and restore_from_file backs up the
live database, closes every thread's connection, copies the file,
removes stale -wal/-shm and re-runs migrations. The admin handlers
stream the upload to a temp file with an early size cutoff, reject
anything verify_database_file refuses with a 400 and keep the
{success: ...} JSON shape. The schema-init connection is now closed
instead of leaking.
This commit is contained in:
Vito Sansevero
+143
-9
@@ -3,10 +3,12 @@ Database schema and models for KikoTextEncode prompt storage.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import shutil
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
from typing import Set
|
||||
from pathlib import Path
|
||||
from typing import Set, Tuple
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
@@ -31,6 +33,10 @@ def utc_now_iso() -> str:
|
||||
)
|
||||
|
||||
|
||||
SQLITE_HEADER = b"SQLite format 3\x00"
|
||||
REQUIRED_PROMPT_COLUMNS = ("id", "text", "created_at")
|
||||
|
||||
|
||||
class PromptModel:
|
||||
"""Database model for prompt storage and schema management."""
|
||||
|
||||
@@ -49,6 +55,9 @@ class PromptModel:
|
||||
self._local = threading.local()
|
||||
self._connections: Set[sqlite3.Connection] = set()
|
||||
self._conn_lock = threading.Lock()
|
||||
# Bumped by _close_all_connections so threads holding a closed
|
||||
# connection reopen on their next get_connection().
|
||||
self._generation = 0
|
||||
self._ensure_database_exists()
|
||||
|
||||
def _ensure_database_exists(self) -> None:
|
||||
@@ -62,13 +71,16 @@ class PromptModel:
|
||||
Exception: If database creation fails
|
||||
"""
|
||||
try:
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute("PRAGMA journal_mode = WAL")
|
||||
conn.execute("PRAGMA busy_timeout = 5000")
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
self._create_tables(conn)
|
||||
self._create_indexes(conn)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error creating database: {e}")
|
||||
raise
|
||||
@@ -200,7 +212,7 @@ class PromptModel:
|
||||
sqlite3.Connection: The current thread's configured connection
|
||||
"""
|
||||
conn = getattr(self._local, "conn", None)
|
||||
if conn is not None:
|
||||
if conn is not None and self._local.generation == self._generation:
|
||||
return conn
|
||||
|
||||
# check_same_thread=False only so restore_from_file can close other
|
||||
@@ -211,10 +223,24 @@ class PromptModel:
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
conn.execute("PRAGMA busy_timeout = 5000")
|
||||
self._local.conn = conn
|
||||
self._local.generation = self._generation
|
||||
with self._conn_lock:
|
||||
self._connections.add(conn)
|
||||
return conn
|
||||
|
||||
def _close_all_connections(self) -> None:
|
||||
"""Close every thread's connection; they reopen lazily on next use."""
|
||||
with self._conn_lock:
|
||||
connections = list(self._connections)
|
||||
self._connections.clear()
|
||||
self._generation += 1
|
||||
self._local.conn = None
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's connection, if it has one.
|
||||
|
||||
@@ -599,19 +625,127 @@ class PromptModel:
|
||||
|
||||
def backup_database(self, backup_path: str) -> bool:
|
||||
"""
|
||||
Create a backup of the database.
|
||||
Write a consistent copy of the database to backup_path.
|
||||
|
||||
Uses the SQLite online backup API through this thread's connection,
|
||||
so rows that are committed but still sit in the -wal file are
|
||||
included and other threads may keep reading and writing meanwhile.
|
||||
|
||||
Args:
|
||||
backup_path: Path where the backup should be saved
|
||||
backup_path: Path where the backup should be saved (overwritten)
|
||||
|
||||
Returns:
|
||||
bool: True if backup was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
import shutil
|
||||
|
||||
shutil.copy2(self.db_path, backup_path)
|
||||
if os.path.exists(backup_path) and os.path.samefile(
|
||||
backup_path, self.db_path
|
||||
):
|
||||
self.logger.error("Refusing to back up the database onto itself")
|
||||
return False
|
||||
if os.path.exists(backup_path):
|
||||
os.remove(backup_path)
|
||||
source = self.get_connection()
|
||||
destination = sqlite3.connect(backup_path)
|
||||
try:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
destination.close()
|
||||
return True
|
||||
except Exception as e:
|
||||
except (OSError, sqlite3.Error) as e:
|
||||
self.logger.error(f"Error creating database backup: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def verify_database_file(path: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
Check that a file is a healthy prompts database before restoring it.
|
||||
|
||||
Opens the file read-only, runs PRAGMA integrity_check and confirms the
|
||||
prompts table has the columns every version of this extension needs.
|
||||
|
||||
Args:
|
||||
path: File to inspect
|
||||
|
||||
Returns:
|
||||
(ok, reason): reason is "ok" when the file is acceptable, otherwise
|
||||
a human-readable explanation. Never raises for bad input.
|
||||
"""
|
||||
if not os.path.isfile(path):
|
||||
return False, "File not found"
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
header = fh.read(len(SQLITE_HEADER))
|
||||
except OSError as e:
|
||||
return False, f"Cannot read file: {e}"
|
||||
if header != SQLITE_HEADER:
|
||||
return False, "Not a SQLite database file"
|
||||
|
||||
uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro"
|
||||
try:
|
||||
conn = sqlite3.connect(uri, uri=True)
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
try:
|
||||
row = conn.execute("PRAGMA integrity_check").fetchone()
|
||||
verdict = row[0] if row else "no result"
|
||||
if verdict != "ok":
|
||||
return False, f"Integrity check failed: {verdict}"
|
||||
table = conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'"
|
||||
).fetchone()
|
||||
if not table:
|
||||
return False, "Database does not contain a 'prompts' table"
|
||||
columns = {r[1] for r in conn.execute("PRAGMA table_info(prompts)")}
|
||||
missing = [c for c in REQUIRED_PROMPT_COLUMNS if c not in columns]
|
||||
if missing:
|
||||
return (
|
||||
False,
|
||||
"prompts table is missing required columns: " + ", ".join(missing),
|
||||
)
|
||||
return True, "ok"
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def restore_from_file(self, src_path: str) -> str:
|
||||
"""
|
||||
Replace the live database with the file at src_path.
|
||||
|
||||
The current database is first backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), every
|
||||
thread's connection is closed, the file is copied over, stale -wal and
|
||||
-shm files are removed and the schema migrations run on the restored
|
||||
file. Connections reopen lazily on the next get_connection().
|
||||
|
||||
Callers should run verify_database_file(src_path) first.
|
||||
|
||||
Args:
|
||||
src_path: SQLite file to restore
|
||||
|
||||
Returns:
|
||||
Path of the backup taken of the previous database, or "" when
|
||||
there was no database to back up.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the previous database could not be backed up
|
||||
OSError: If the restored file could not be copied into place
|
||||
"""
|
||||
backup_path = ""
|
||||
if os.path.exists(self.db_path):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = f"{self.db_path}.backup_{timestamp}"
|
||||
if not self.backup_database(backup_path):
|
||||
raise RuntimeError("Could not back up the current database")
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
|
||||
self._close_all_connections()
|
||||
shutil.copy2(src_path, self.db_path)
|
||||
for suffix in ("-wal", "-shm"):
|
||||
stale = self.db_path + suffix
|
||||
if os.path.exists(stale):
|
||||
os.remove(stale)
|
||||
self._ensure_database_exists()
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
return backup_path
|
||||
|
||||
+75
-92
@@ -18,6 +18,11 @@ try:
|
||||
except ImportError:
|
||||
from utils.validators import validate_result_timeout
|
||||
|
||||
try:
|
||||
from ...database.operations import PromptDatabase
|
||||
except ImportError:
|
||||
from database.operations import PromptDatabase
|
||||
|
||||
|
||||
class AdminRoutesMixin:
|
||||
"""Mixin providing admin, diagnostics, and maintenance API endpoints."""
|
||||
@@ -883,24 +888,33 @@ class AdminRoutesMixin:
|
||||
)
|
||||
|
||||
async def backup_database(self, request):
|
||||
"""Backup the entire prompts.db database file."""
|
||||
"""Download a consistent copy of the prompts database (WAL included)."""
|
||||
try:
|
||||
db_path = self.db.model.db_path
|
||||
model = self.db.model
|
||||
db_path = model.db_path
|
||||
|
||||
if not os.path.exists(db_path):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Database file not found"}, status=404
|
||||
)
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as temp_file:
|
||||
temp_path = temp_file.name
|
||||
|
||||
shutil.copy2(db_path, temp_path)
|
||||
|
||||
with open(temp_path, "rb") as f:
|
||||
file_data = f.read()
|
||||
|
||||
os.unlink(temp_path)
|
||||
fd, temp_path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
ok = await loop.run_in_executor(None, model.backup_database, temp_path)
|
||||
if not ok:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Failed to create database backup"},
|
||||
status=500,
|
||||
)
|
||||
# Read once (the file is deleted right after) rather than
|
||||
# streaming so the temp file never outlives the request.
|
||||
with open(temp_path, "rb") as fh:
|
||||
file_data = fh.read()
|
||||
finally:
|
||||
if os.path.exists(temp_path):
|
||||
os.unlink(temp_path)
|
||||
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"prompts_backup_{timestamp}.db"
|
||||
@@ -922,7 +936,14 @@ class AdminRoutesMixin:
|
||||
)
|
||||
|
||||
async def restore_database(self, request):
|
||||
"""Restore the prompts.db database from uploaded file."""
|
||||
"""Restore the prompts database from an uploaded SQLite file.
|
||||
|
||||
The upload is streamed to a temp file (capped at restore_max_bytes),
|
||||
verified with PromptModel.verify_database_file and only then copied
|
||||
over the live database, which is backed up first.
|
||||
"""
|
||||
max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024)
|
||||
temp_path = None
|
||||
try:
|
||||
reader = await request.multipart()
|
||||
field = await reader.next()
|
||||
@@ -936,95 +957,52 @@ class AdminRoutesMixin:
|
||||
status=400,
|
||||
)
|
||||
|
||||
MAX_RESTORE_SIZE = 100 * 1024 * 1024 # 100MB
|
||||
file_data = await field.read()
|
||||
fd, temp_path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
size = 0
|
||||
with open(temp_path, "wb") as out:
|
||||
while True:
|
||||
chunk = await field.read_chunk()
|
||||
if not chunk:
|
||||
break
|
||||
size += len(chunk)
|
||||
if size > max_bytes:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"File too large. Maximum size is {max_bytes // (1024 * 1024)}MB",
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
out.write(chunk)
|
||||
|
||||
if not file_data:
|
||||
if size == 0:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Uploaded file is empty"}, status=400
|
||||
)
|
||||
|
||||
if len(file_data) > MAX_RESTORE_SIZE:
|
||||
model = self.db.model
|
||||
ok, reason = model.verify_database_file(temp_path)
|
||||
if not ok:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"File too large. Maximum size is {MAX_RESTORE_SIZE // (1024*1024)}MB",
|
||||
},
|
||||
status=400,
|
||||
{"success": False, "error": reason}, status=400
|
||||
)
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as temp_file:
|
||||
temp_path = temp_file.name
|
||||
temp_file.write(file_data)
|
||||
loop = asyncio.get_running_loop()
|
||||
backup_path = await loop.run_in_executor(
|
||||
None, model.restore_from_file, temp_path
|
||||
)
|
||||
self.db = PromptDatabase(model.db_path)
|
||||
prompt_count = self.db.model.get_database_info().get("total_prompts", 0)
|
||||
|
||||
try:
|
||||
with sqlite3.connect(temp_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'"
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
raise ValueError("Database does not contain a 'prompts' table")
|
||||
|
||||
cursor = conn.execute("PRAGMA table_info(prompts)")
|
||||
columns = [row["name"] for row in cursor.fetchall()]
|
||||
required_columns = ["id", "text", "created_at"]
|
||||
|
||||
for col in required_columns:
|
||||
if col not in columns:
|
||||
raise ValueError(f"Database missing required column: {col}")
|
||||
|
||||
cursor = conn.execute("SELECT COUNT(*) as count FROM prompts")
|
||||
prompt_count = cursor.fetchone()["count"]
|
||||
|
||||
db_path = self.db.model.db_path
|
||||
backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||||
|
||||
if os.path.exists(db_path):
|
||||
shutil.copy2(db_path, backup_path)
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
|
||||
shutil.copy2(temp_path, db_path)
|
||||
|
||||
# Reinitialize the database connection
|
||||
try:
|
||||
from ...database.operations import PromptDatabase
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0,
|
||||
os.path.dirname(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
),
|
||||
)
|
||||
from database.operations import PromptDatabase
|
||||
self.db = PromptDatabase()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": f"Database restored successfully. Found {prompt_count} prompts.",
|
||||
"prompt_count": prompt_count,
|
||||
"backup_created": (
|
||||
backup_path if os.path.exists(db_path) else None
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
return web.json_response(
|
||||
{"success": False, "error": f"Invalid SQLite database: {str(e)}"},
|
||||
status=400,
|
||||
)
|
||||
except ValueError as e:
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(e)}, status=400
|
||||
)
|
||||
finally:
|
||||
if os.path.exists(temp_path):
|
||||
os.unlink(temp_path)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": f"Database restored successfully. Found {prompt_count} prompts.",
|
||||
"prompt_count": prompt_count,
|
||||
"backup_created": backup_path or None,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Restore error: {e}", exc_info=True)
|
||||
@@ -1032,6 +1010,11 @@ class AdminRoutesMixin:
|
||||
{"success": False, "error": f"Failed to restore database: {str(e)}"},
|
||||
status=500,
|
||||
)
|
||||
finally:
|
||||
if temp_path:
|
||||
for suffix in ("", "-wal", "-shm", "-journal"):
|
||||
if os.path.exists(temp_path + suffix):
|
||||
os.unlink(temp_path + suffix)
|
||||
|
||||
async def scan_images(self, request):
|
||||
"""Scan ComfyUI output images for prompt metadata and add them to the database."""
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
"""Route tests for GET /prompt_manager/backup and POST /prompt_manager/restore.
|
||||
|
||||
The backup must include rows still sitting in the WAL, the restore must
|
||||
stream the upload to disk, reject anything that is not a healthy prompts
|
||||
database and leave the live database untouched when it does.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import AioHTTPTestCase
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from py.api import PromptManagerAPI
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
def _texts(path):
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class BackupRestoreTestCase(AioHTTPTestCase):
|
||||
async def get_application(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self.db_path)
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
async def tearDownAsync(self):
|
||||
self.api.db.close()
|
||||
for name in os.listdir(self.tmpdir):
|
||||
os.unlink(os.path.join(self.tmpdir, name))
|
||||
os.rmdir(self.tmpdir)
|
||||
|
||||
def _save(self, text):
|
||||
return self.api.db.save_prompt(
|
||||
text=text, prompt_hash=generate_prompt_hash(text)
|
||||
)
|
||||
|
||||
def _write(self, name, data):
|
||||
path = os.path.join(self.tmpdir, name)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(data)
|
||||
return path
|
||||
|
||||
def _make_db_file(self, texts, name="upload.db"):
|
||||
path = os.path.join(self.tmpdir, name)
|
||||
db = PromptDatabase(path)
|
||||
for text in texts:
|
||||
db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
db.close()
|
||||
return path
|
||||
|
||||
async def _restore(self, path, field="database_file"):
|
||||
with open(path, "rb") as fh:
|
||||
form = aiohttp.FormData()
|
||||
form.add_field(
|
||||
field,
|
||||
fh.read(),
|
||||
filename=os.path.basename(path),
|
||||
content_type="application/octet-stream",
|
||||
)
|
||||
return await self.client.post("/prompt_manager/restore", data=form)
|
||||
|
||||
|
||||
class TestBackupRoute(BackupRestoreTestCase):
|
||||
async def test_backup_download_contains_rows_still_in_the_wal(self):
|
||||
for n in range(3):
|
||||
self._save(f"backup me {n}")
|
||||
self.assertGreater(os.path.getsize(self.db_path + "-wal"), 0)
|
||||
created = []
|
||||
real_mkstemp = tempfile.mkstemp
|
||||
|
||||
def spy_mkstemp(*args, **kwargs):
|
||||
result = real_mkstemp(*args, **kwargs)
|
||||
created.append(result[1])
|
||||
return result
|
||||
|
||||
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
|
||||
resp = await self.client.get("/prompt_manager/backup")
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(resp.content_type, "application/octet-stream")
|
||||
disposition = resp.headers["Content-Disposition"]
|
||||
self.assertIn("attachment", disposition)
|
||||
self.assertIn("prompts_backup_", disposition)
|
||||
self.assertTrue(disposition.rstrip('"').endswith(".db"))
|
||||
|
||||
body = await resp.read()
|
||||
downloaded = self._write("downloaded.db", body)
|
||||
self.assertEqual(
|
||||
_texts(downloaded), ["backup me 0", "backup me 1", "backup me 2"]
|
||||
)
|
||||
# The server-side temp file is gone once the response is complete
|
||||
self.assertEqual(len(created), 1)
|
||||
self.assertFalse(os.path.exists(created[0]))
|
||||
|
||||
async def test_backup_of_missing_database_is_404(self):
|
||||
self.api.db.model.db_path = os.path.join(self.tmpdir, "nope.db")
|
||||
resp = await self.client.get("/prompt_manager/backup")
|
||||
self.assertEqual(resp.status, 404)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
async def test_backup_failure_is_500(self):
|
||||
self._save("x")
|
||||
with patch.object(self.api.db.model, "backup_database", return_value=False):
|
||||
resp = await self.client.get("/prompt_manager/backup")
|
||||
self.assertEqual(resp.status, 500)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("backup", data["error"].lower())
|
||||
|
||||
|
||||
class TestRestoreRoute(BackupRestoreTestCase):
|
||||
async def test_restore_replaces_database_and_reports_backup(self):
|
||||
self._save("old prompt")
|
||||
old_db = self.api.db
|
||||
upload = self._make_db_file(["restored a", "restored b"])
|
||||
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["prompt_count"], 2)
|
||||
self.assertTrue(os.path.exists(data["backup_created"]))
|
||||
self.assertEqual(_texts(data["backup_created"]), ["old prompt"])
|
||||
self.assertIsNot(self.api.db, old_db)
|
||||
self.assertEqual(self.api.db.model.db_path, self.db_path)
|
||||
|
||||
texts = sorted(p["text"] for p in self.api.db.search_prompts())
|
||||
self.assertEqual(texts, ["restored a", "restored b"])
|
||||
|
||||
async def test_restore_leaves_no_temp_files_behind(self):
|
||||
upload = self._make_db_file(["restored"])
|
||||
created = []
|
||||
real_mkstemp = tempfile.mkstemp
|
||||
|
||||
def spy_mkstemp(*args, **kwargs):
|
||||
result = real_mkstemp(*args, **kwargs)
|
||||
created.append(result[1])
|
||||
return result
|
||||
|
||||
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(len(created), 1)
|
||||
for suffix in ("", "-wal", "-shm", "-journal"):
|
||||
self.assertFalse(os.path.exists(created[0] + suffix))
|
||||
|
||||
async def test_text_file_is_rejected_and_database_untouched(self):
|
||||
self._save("keep me")
|
||||
upload = self._write("notes.txt", b"this is not a database")
|
||||
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("SQLite", data["error"])
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
self.assertFalse(
|
||||
any(n.startswith("prompts.db.backup_") for n in os.listdir(self.tmpdir))
|
||||
)
|
||||
|
||||
async def test_database_without_prompts_table_is_rejected(self):
|
||||
self._save("keep me")
|
||||
path = os.path.join(self.tmpdir, "other.db")
|
||||
conn = sqlite3.connect(path)
|
||||
conn.execute("CREATE TABLE things (id INTEGER)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
resp = await self._restore(path)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("prompts", data["error"])
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
|
||||
async def test_wrong_field_name_is_rejected(self):
|
||||
upload = self._make_db_file(["restored"])
|
||||
resp = await self._restore(upload, field="file")
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("database_file", data["error"])
|
||||
|
||||
async def test_empty_upload_is_rejected(self):
|
||||
upload = self._write("empty.db", b"")
|
||||
resp = await self._restore(upload)
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("empty", data["error"].lower())
|
||||
|
||||
async def test_upload_over_the_size_limit_is_rejected_while_streaming(self):
|
||||
self._save("keep me")
|
||||
self.api.restore_max_bytes = 1024
|
||||
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
|
||||
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("too large", data["error"].lower())
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
|
||||
async def test_non_multipart_request_is_rejected(self):
|
||||
resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
|
||||
self.assertIn(resp.status, (400, 500))
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,243 @@
|
||||
"""Backup, verification and restore of the prompts database.
|
||||
|
||||
The database runs in WAL mode, so committed rows can sit in prompts.db-wal
|
||||
for a long time. A backup must include them, a restore must reject files
|
||||
that are not healthy SQLite databases, and the live database must stay
|
||||
usable afterwards.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database.models import PromptModel
|
||||
from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
def _count_prompts(path):
|
||||
"""Row count read through a brand-new connection (closed before return)."""
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
return conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _texts(path):
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class BackupTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.path = os.path.join(self.tmpdir, "prompts.db")
|
||||
self.db = PromptDatabase(self.path)
|
||||
self.addCleanup(self._cleanup)
|
||||
|
||||
def _cleanup(self):
|
||||
self.db.close()
|
||||
for name in os.listdir(self.tmpdir):
|
||||
os.unlink(os.path.join(self.tmpdir, name))
|
||||
os.rmdir(self.tmpdir)
|
||||
|
||||
def _save(self, db, text):
|
||||
return db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def _make_source_db(self, texts):
|
||||
"""A second database file holding the given prompts, fully closed."""
|
||||
src_path = os.path.join(self.tmpdir, "source.db")
|
||||
src = PromptDatabase(src_path)
|
||||
for text in texts:
|
||||
self._save(src, text)
|
||||
src.close()
|
||||
return src_path
|
||||
|
||||
|
||||
class TestBackupDatabase(BackupTestCase):
|
||||
def test_backup_contains_rows_that_only_exist_in_the_wal(self):
|
||||
for n in range(5):
|
||||
self._save(self.db, f"wal prompt {n}")
|
||||
wal = self.path + "-wal"
|
||||
self.assertTrue(os.path.exists(wal) and os.path.getsize(wal) > 0)
|
||||
|
||||
dst = os.path.join(self.tmpdir, "backup.db")
|
||||
self.assertTrue(self.db.model.backup_database(dst))
|
||||
|
||||
self.assertEqual(_count_prompts(dst), 5)
|
||||
self.assertFalse(os.path.exists(dst + "-wal"))
|
||||
# The live database is untouched and still usable
|
||||
self.assertEqual(len(self.db.search_prompts(text="wal prompt")), 5)
|
||||
|
||||
def test_backup_overwrites_a_stale_destination(self):
|
||||
self._save(self.db, "one")
|
||||
dst = os.path.join(self.tmpdir, "backup.db")
|
||||
with open(dst, "wb") as fh:
|
||||
fh.write(b"not a database at all")
|
||||
self.assertTrue(self.db.model.backup_database(dst))
|
||||
self.assertEqual(_count_prompts(dst), 1)
|
||||
|
||||
def test_backup_onto_the_live_file_is_refused(self):
|
||||
self._save(self.db, "one")
|
||||
self.assertFalse(self.db.model.backup_database(self.path))
|
||||
self.assertEqual(len(self.db.search_prompts(text="one")), 1)
|
||||
|
||||
def test_backup_to_an_unwritable_location_returns_false(self):
|
||||
dst = os.path.join(self.tmpdir, "missing", "dir", "backup.db")
|
||||
self.assertFalse(self.db.model.backup_database(dst))
|
||||
|
||||
|
||||
class TestVerifyDatabaseFile(BackupTestCase):
|
||||
def test_text_file_is_rejected_without_raising(self):
|
||||
path = os.path.join(self.tmpdir, "notes.txt")
|
||||
with open(path, "w", encoding="utf-8") as fh:
|
||||
fh.write("just some text, definitely not sqlite")
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("SQLite", reason)
|
||||
|
||||
def test_missing_file_is_rejected(self):
|
||||
ok, reason = PromptModel.verify_database_file(
|
||||
os.path.join(self.tmpdir, "nope.db")
|
||||
)
|
||||
self.assertFalse(ok)
|
||||
self.assertTrue(reason)
|
||||
|
||||
def test_database_without_prompts_table_is_rejected(self):
|
||||
path = os.path.join(self.tmpdir, "other.db")
|
||||
conn = sqlite3.connect(path)
|
||||
conn.execute("CREATE TABLE something (id INTEGER)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("prompts", reason)
|
||||
|
||||
def test_prompts_table_missing_required_columns_is_rejected(self):
|
||||
path = os.path.join(self.tmpdir, "partial.db")
|
||||
conn = sqlite3.connect(path)
|
||||
conn.execute("CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("created_at", reason)
|
||||
|
||||
def test_corrupted_database_fails_integrity_check(self):
|
||||
path = self._make_source_db([f"prompt {n}" for n in range(40)])
|
||||
with open(path, "r+b") as fh:
|
||||
page_size = int.from_bytes(fh.read(18)[16:18], "big") or 4096
|
||||
fh.seek(page_size * 2) # scribble over pages 3 and 4 entirely
|
||||
fh.write(b"\xff" * (page_size * 2))
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertFalse(ok)
|
||||
self.assertTrue(reason)
|
||||
|
||||
def test_valid_database_passes(self):
|
||||
path = self._make_source_db(["a", "b"])
|
||||
ok, reason = PromptModel.verify_database_file(path)
|
||||
self.assertTrue(ok, reason)
|
||||
|
||||
def test_verification_does_not_lock_or_modify_the_file(self):
|
||||
path = self._make_source_db(["a"])
|
||||
before = os.path.getsize(path)
|
||||
PromptModel.verify_database_file(path)
|
||||
os.unlink(path) # would fail on Windows if a connection were left open
|
||||
self.assertEqual(before, before)
|
||||
|
||||
|
||||
class TestRestoreFromFile(BackupTestCase):
|
||||
def test_restore_replaces_content_and_backs_up_the_old_database(self):
|
||||
self._save(self.db, "old prompt")
|
||||
src = self._make_source_db(["restored one", "restored two"])
|
||||
|
||||
backup_path = self.db.model.restore_from_file(src)
|
||||
|
||||
self.assertTrue(backup_path.startswith(self.path + ".backup_"))
|
||||
self.assertTrue(os.path.exists(backup_path))
|
||||
self.assertEqual(_texts(backup_path), ["old prompt"])
|
||||
self.assertFalse(os.path.exists(self.path + "-wal"))
|
||||
self.assertFalse(os.path.exists(self.path + "-shm"))
|
||||
self.assertEqual(_texts(self.path), ["restored one", "restored two"])
|
||||
|
||||
def test_database_is_usable_from_every_thread_after_restore(self):
|
||||
self._save(self.db, "old prompt")
|
||||
src = self._make_source_db(["restored"])
|
||||
|
||||
ready = threading.Event()
|
||||
proceed = threading.Event()
|
||||
results = {}
|
||||
|
||||
def worker():
|
||||
self.db.search_prompts(text="old") # opens this thread's connection
|
||||
ready.set()
|
||||
proceed.wait(5)
|
||||
try:
|
||||
results["texts"] = [p["text"] for p in self.db.search_prompts()]
|
||||
except Exception as exc: # noqa: BLE001 - reported through results
|
||||
results["error"] = exc
|
||||
finally:
|
||||
self.db.close()
|
||||
|
||||
thread = threading.Thread(target=worker)
|
||||
thread.start()
|
||||
ready.wait(5)
|
||||
|
||||
self.db.model.restore_from_file(src)
|
||||
proceed.set()
|
||||
thread.join(5)
|
||||
|
||||
self.assertNotIn("error", results)
|
||||
self.assertEqual(results["texts"], ["restored"])
|
||||
self.assertEqual([p["text"] for p in self.db.search_prompts()], ["restored"])
|
||||
self.assertEqual(self._save(self.db, "after restore"), 2)
|
||||
|
||||
def test_restore_applies_migrations_to_an_old_schema(self):
|
||||
src = os.path.join(self.tmpdir, "legacy.db")
|
||||
conn = sqlite3.connect(src)
|
||||
# The 3.0-era schema: workflow_name still present, no usage columns
|
||||
conn.execute(
|
||||
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
||||
" text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP,"
|
||||
" category TEXT, tags TEXT, rating INTEGER, notes TEXT,"
|
||||
" hash TEXT UNIQUE, workflow_name TEXT)"
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO prompts (text, created_at, tags, workflow_name)"
|
||||
" VALUES ('legacy', '2024-01-01T00:00:00', '[\"old\"]', 'wf')"
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
self.db.model.restore_from_file(src)
|
||||
|
||||
prompt = self.db.get_prompt_by_id(1)
|
||||
self.assertEqual(prompt["text"], "legacy")
|
||||
self.assertEqual(prompt["tags"], ["old"])
|
||||
self.assertEqual(prompt["run_count"], 1)
|
||||
self.assertNotIn("workflow_name", prompt)
|
||||
self.assertEqual(self.db.get_prompt_images(1), [])
|
||||
|
||||
def test_restore_without_a_live_database_returns_empty_backup_path(self):
|
||||
src = self._make_source_db(["fresh"])
|
||||
self.db.close()
|
||||
os.unlink(self.path)
|
||||
for suffix in ("-wal", "-shm"):
|
||||
if os.path.exists(self.path + suffix):
|
||||
os.unlink(self.path + suffix)
|
||||
|
||||
self.assertEqual(self.db.model.restore_from_file(src), "")
|
||||
self.assertEqual(_texts(self.path), ["fresh"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user