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"])