From 1f422d2c43d82bf60d62c1b007e71b2c17a5524c Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 19:34:29 -0700 Subject: [PATCH] fix: restore through SQLite's backup API, close_all for shared files, safer output scans restore_from_file copied the uploaded file over the live database and deleted its WAL while other PromptModel instances (nodes, image monitor, queue hook) still held connections, which corrupts on POSIX and fails on Windows. It now copies pages through the online backup API into the live connection so every open handle sees the new content. PromptDatabase.close_all() closes the connections of every thread, which tests need before deleting a database. Output scans skip symlinked files, filesystem entries get a stable digest id instead of a per-process hash(), and thumbnail streaming stops when the client disconnects. --- database/models.py | 55 ++++++++++++++++++++++++------- database/operations.py | 4 +++ py/api/images.py | 20 ++++++++--- tests/test_api_images.py | 49 +++++++++++++++++++++++++++ tests/test_database_backup.py | 62 +++++++++++++++++++++++++++++++++-- 5 files changed, 171 insertions(+), 19 deletions(-) diff --git a/database/models.py b/database/models.py index 7a2c362..3c1d4c7 100644 --- a/database/models.py +++ b/database/models.py @@ -292,8 +292,12 @@ class PromptModel: self._connections.add(conn) return conn - def _close_all_connections(self) -> None: - """Close every thread's connection; they reopen lazily on next use.""" + def close_all(self) -> None: + """Close every thread's connection; they reopen lazily on next use. + + Needed before deleting the database file (Windows keeps open files + locked) and by tests that ran queries through executor threads. + """ with self._conn_lock: connections = list(self._connections) self._connections.clear() @@ -755,15 +759,47 @@ class PromptModel: finally: conn.close() + def _replace_contents_from(self, src_path: str) -> None: + """Copy every page of *src_path* into the live database. + + Goes through SQLite's online backup API into this thread's normal + connection, so SQLite takes the locks itself and the other + PromptModel instances that share the file (nodes, image monitor, + queue hook, API) simply see the new pages on their next statement. + Nothing is swapped underneath open handles and no -wal/-shm file has + to be deleted, which also keeps Windows happy. When SQLite refuses + (page sizes differ while the target is in WAL mode) fall back to a + file copy with every connection of this instance closed first. + """ + source = sqlite3.connect( + Path(src_path).resolve().as_uri() + "?mode=ro", uri=True + ) + try: + try: + source.backup(self.get_connection()) + return + except sqlite3.OperationalError as e: + self.logger.warning( + f"Online restore failed ({e}); falling back to a file copy" + ) + finally: + source.close() + self.close_all() + 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) + 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 - ``.backup_`` (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(). + ``.backup_`` (WAL content included), the + pages of ``src_path`` are copied into the live database through + SQLite's backup API (visible to every open connection) and the schema + migrations run on the restored content. Callers should run verify_database_file(src_path) first. @@ -786,12 +822,7 @@ class PromptModel: 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._replace_contents_from(src_path) with PromptModel._init_lock: PromptModel._initialized_paths.discard(self._schema_key()) self._ensure_database_exists() diff --git a/database/operations.py b/database/operations.py index 6809303..927b1f7 100644 --- a/database/operations.py +++ b/database/operations.py @@ -120,6 +120,10 @@ class PromptDatabase: """Close the calling thread's database connection (see PromptModel.close).""" self.model.close() + def close_all(self) -> None: + """Close every thread's connection (see PromptModel.close_all).""" + self.model.close_all() + def save_prompt( self, text: str, diff --git a/py/api/images.py b/py/api/images.py index 2002711..5076540 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -1,6 +1,7 @@ """Image and gallery API routes for PromptManager.""" import asyncio +import hashlib import json import os import re @@ -216,6 +217,8 @@ def _iter_media_files(output_path, max_files): for name in sorted(files): if Path(name).suffix.lower() not in MEDIA_EXTENSIONS: continue + if (root_path / name).is_symlink(): + continue # never read through a link out of the output tree if produced >= max_files: return produced += 1 @@ -240,7 +243,7 @@ def _output_image_entry(media_path, output_path, root_index): ) return { - "id": str(hash(str(media_path))), + "id": hashlib.sha1(str(media_path).encode("utf-8")).hexdigest()[:16], "filename": media_path.name, "path": str(media_path), "relative_path": str(rel_path), @@ -819,7 +822,7 @@ class ImageRoutesMixin: """Process *targets* one per executor job, emitting SSE after each.""" total = len(targets) video_count = sum(1 for _, _, is_video in targets if is_video) - await self._emit_sse( + connected = await self._emit_sse( response, "start", { @@ -833,6 +836,8 @@ class ImageRoutesMixin: ), }, ) + if not connected: + return stats = _new_thumbnail_stats() start = _time.monotonic() for index, (src, dst, is_video) in enumerate(targets, start=1): @@ -846,23 +851,28 @@ class ImageRoutesMixin: "file_error", {"file": result["file"], "error": result["error"]}, ) - await self._emit_sse( + connected = await self._emit_sse( response, "progress", _thumbnail_progress_payload(index, total, stats, result, start), ) + if not connected: + self.logger.info("Thumbnail client disconnected; stopping early") + return await asyncio.sleep(0) # let other requests run between files payload = _thumbnail_complete_payload(total, stats, start) await self._emit_sse(response, "complete", payload) self.logger.info(payload["message"]) - async def _emit_sse(self, response, event, data): - """Write one SSE frame; a closed client is logged, not raised.""" + async def _emit_sse(self, response, event, data) -> bool: + """Write one SSE frame. Returns False once the client has gone away.""" try: frame = f"event: {event}\ndata: {json.dumps(data)}\n\n" await response.write(frame.encode("utf-8")) + return True except Exception as e: self.logger.warning(f"Failed to send SSE message: {e}") + return False def _generate_video_thumbnail(self, video_path, thumbnail_path, thumbnail_size): """Generate thumbnail from video file. Returns True if successful.""" diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 9454559..835d4ba 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -1291,3 +1291,52 @@ class TestDeleteAndLinkErrors(ImageRouteCoverageCase): if __name__ == "__main__": unittest.main() + + +class TestOutputScanHelpers(ImageAPITestCase): + """Stable ids for filesystem entries, no symlink reads, stop on disconnect.""" + + def test_output_entry_id_is_a_stable_digest_of_the_path(self): + import hashlib + + output = Path(self.output_dir) + media = output / "stable.png" + make_png(media) + entry = images_module._output_image_entry(media, output, 0) + expected = hashlib.sha1(str(media).encode("utf-8")).hexdigest()[:16] + self.assertEqual(entry["id"], expected) + + def test_iter_media_files_skips_symlinked_files(self): + output = Path(self.output_dir) + make_png(output / "real.png") + with tempfile.TemporaryDirectory() as other: + outside = Path(other) / "outside.png" + make_png(outside) + try: + os.symlink(outside, output / "link.png") + except (OSError, NotImplementedError): + self.skipTest("symlinks not available") + names = sorted(p.name for p in images_module._iter_media_files(output, 100)) + self.assertEqual(names, ["real.png"]) + + async def test_thumbnail_stream_stops_when_the_client_disconnects(self): + generated = [] + + def fake_generate(src, dst, size, is_video): + generated.append(src) + return { + "action": "generated", + "file": src.name, + "dir": src.parent.name, + "type": "image", + } + + self.api._generate_one = fake_generate + + class GoneResponse: + async def write(self, data): + raise ConnectionResetError("client went away") + + targets = [(Path(f"/x/{i}.png"), Path(f"/x/t{i}.jpg"), False) for i in range(3)] + await self.api._stream_thumbnails(GoneResponse(), targets, (256, 256)) + self.assertEqual(generated, []) diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 8326b3b..c2bff76 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -165,8 +165,9 @@ class TestRestoreFromFile(BackupTestCase): 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")) + # The restore goes through SQLite's backup API into the live file, so + # the WAL stays valid (and shared with other instances) rather than + # being deleted under them; the content is what matters. self.assertEqual(_texts(self.path), ["restored one", "restored two"]) def test_database_is_usable_from_every_thread_after_restore(self): @@ -241,3 +242,60 @@ class TestRestoreFromFile(BackupTestCase): if __name__ == "__main__": unittest.main() + + +class TestConnectionsAcrossInstances(BackupTestCase): + """Several PromptModel instances share one file (nodes, API, monitor).""" + + def test_close_all_closes_connections_opened_on_other_threads(self): + opened = [] + + def open_one(): + opened.append(self.db.model.get_connection()) + + threads = [threading.Thread(target=open_one) for _ in range(3)] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEqual(len(opened), 3) + + self.db.close_all() + + for conn in opened: + with self.assertRaises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + + def test_restore_is_seen_by_another_instance_holding_an_open_connection(self): + self.db.save_prompt(text="before", prompt_hash=generate_prompt_hash("before")) + snapshot = os.path.join(self.tmpdir, "snapshot.db") + other_db = PromptDatabase(snapshot) + for i in range(3): + other_db.save_prompt( + text=f"after {i}", prompt_hash=generate_prompt_hash(f"after {i}") + ) + other_db.close_all() + + # A second instance on the live file, with a connection kept open on + # a worker thread for the whole restore (like the image monitor). + bystander = PromptDatabase(self.path) + ready, go, seen = threading.Event(), threading.Event(), [] + + def hold_and_read(): + bystander.model.get_connection().execute("SELECT COUNT(*) FROM prompts") + ready.set() + go.wait(10) + seen.append(len(bystander.search_prompts(limit=50))) + + worker = threading.Thread(target=hold_and_read) + worker.start() + self.assertTrue(ready.wait(10)) + try: + self.db.model.restore_from_file(snapshot) + finally: + go.set() + worker.join(10) + bystander.close_all() + + self.assertEqual(seen, [3]) + self.assertEqual(sorted(_texts(self.path)), ["after 0", "after 1", "after 2"])