diff --git a/database/models.py b/database/models.py index 5dc2cfe..14a491b 100644 --- a/database/models.py +++ b/database/models.py @@ -6,7 +6,7 @@ import datetime import sqlite3 import os import threading -from typing import Optional +from typing import Set # Import logging system try: @@ -44,7 +44,10 @@ class PromptModel: self.logger = get_logger("prompt_manager.database.models") self.logger.debug(f"Initializing database model with path: {db_path}") self.db_path = db_path - self._conn: Optional[sqlite3.Connection] = None + # One connection per thread (see get_connection); every open connection + # is also tracked here so restore_from_file can close them all. + self._local = threading.local() + self._connections: Set[sqlite3.Connection] = set() self._conn_lock = threading.Lock() self._ensure_database_exists() @@ -182,24 +185,52 @@ class PromptModel: def get_connection(self) -> sqlite3.Connection: """ - Get the persistent database connection. + Get the calling thread's database connection. - Returns a thread-safe, reusable connection protected by a lock. - The connection is created on first call and reused thereafter. - Callers use it as a context manager; the lock is held for the - duration of the ``with`` block. + Each thread gets its own ``sqlite3.Connection`` (kept in + ``threading.local``), created on first use and reused by that thread + thereafter. Every connection is configured with WAL journaling, + ``row_factory = sqlite3.Row``, ``PRAGMA foreign_keys = ON`` and + ``PRAGMA busy_timeout = 5000`` so concurrent writers wait instead of + failing. Callers use it as a context manager (``with conn:``), which + commits or rolls back that thread's transaction and leaves the + connection open. No lock is held while a connection is in use. Returns: - sqlite3.Connection: Configured database connection + sqlite3.Connection: The current thread's configured connection """ + conn = getattr(self._local, "conn", None) + if conn is not None: + return conn + + # check_same_thread=False only so restore_from_file can close other + # threads' connections; each connection is still used by one thread. + conn = sqlite3.connect(self.db_path, check_same_thread=False) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode = WAL") + conn.execute("PRAGMA foreign_keys = ON") + conn.execute("PRAGMA busy_timeout = 5000") + self._local.conn = conn with self._conn_lock: - if self._conn is None: - self._conn = sqlite3.connect(self.db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - self._conn.execute("PRAGMA journal_mode = WAL") - self._conn.execute("PRAGMA foreign_keys = ON") - self._conn.execute("PRAGMA busy_timeout = 5000") - return self._conn + self._connections.add(conn) + return conn + + def close(self) -> None: + """Close the calling thread's connection, if it has one. + + Other threads keep their connections. Call this before deleting or + replacing the database file: Windows refuses to delete an open file. + """ + conn = getattr(self._local, "conn", None) + if conn is None: + return + self._local.conn = None + with self._conn_lock: + self._connections.discard(conn) + try: + conn.close() + except sqlite3.Error as e: + self.logger.warning(f"Error closing database connection: {e}") def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None: """ diff --git a/database/operations.py b/database/operations.py index f3c0765..d83a2b2 100644 --- a/database/operations.py +++ b/database/operations.py @@ -100,6 +100,10 @@ class PromptDatabase: self.model = PromptModel(db_path) self.logger.debug("Database operations initialized successfully") + def close(self) -> None: + """Close the calling thread's database connection (see PromptModel.close).""" + self.model.close() + def save_prompt( self, text: str, diff --git a/tests/test_database.py b/tests/test_database.py index 80087cd..b26fdaa 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -26,13 +26,11 @@ class DatabaseTestCase(unittest.TestCase): self.db = PromptDatabase(self.temp_db.name) def tearDown(self): - if os.path.exists(self.temp_db.name): - os.unlink(self.temp_db.name) - wal = self.temp_db.name + "-wal" - shm = self.temp_db.name + "-shm" - for f in (wal, shm): - if os.path.exists(f): - os.unlink(f) + self.db.close() + for suffix in ("", "-wal", "-shm"): + path = self.temp_db.name + suffix + if os.path.exists(path): + os.unlink(path) def _save( self, text="Test prompt", category=None, tags=None, rating=None, notes=None diff --git a/tests/test_database_concurrency.py b/tests/test_database_concurrency.py new file mode 100644 index 0000000..c13781a --- /dev/null +++ b/tests/test_database_concurrency.py @@ -0,0 +1,110 @@ +"""Concurrent access to PromptDatabase from several threads. + +ComfyUI touches the database from the server loop, executor threads, the +watchdog image monitor and the queue hook at the same time. Every thread must +get its own connection so one thread's commit or rollback never lands on +another thread's half-done write. +""" + +import os +import sys +import tempfile +import threading +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash + +THREADS = 8 +CALLS_PER_THREAD = 50 + + +class ConcurrencyTestCase(unittest.TestCase): + def setUp(self): + fd, self.path = tempfile.mkstemp(suffix=".db") + os.close(fd) + self.db = PromptDatabase(self.path) + self.addCleanup(self._cleanup) + + def _cleanup(self): + self.db.close() + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.path + suffix): + os.unlink(self.path + suffix) + + def _run_threads(self, target): + errors = [] + + def wrapped(index): + try: + target(index) + except Exception as exc: # noqa: BLE001 - collected for the assertion + errors.append(exc) + finally: + self.db.close() + + threads = [threading.Thread(target=wrapped, args=(i,)) for i in range(THREADS)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + return errors + + +class TestPerThreadConnections(ConcurrencyTestCase): + def test_each_thread_gets_its_own_connection(self): + main_conn = self.db.model.get_connection() + self.assertIs(self.db.model.get_connection(), main_conn) + + seen = [] + + def grab(_index): + first = self.db.model.get_connection() + seen.append((first, self.db.model.get_connection())) + + self.assertEqual(self._run_threads(grab), []) + self.assertEqual(len(seen), THREADS) + for first, second in seen: + self.assertIs(first, second) + self.assertIsNot(first, main_conn) + self.assertEqual(len({id(first) for first, _ in seen}), THREADS) + + def test_close_releases_only_the_calling_threads_connection(self): + main_conn = self.db.model.get_connection() + other = [] + + def grab(_index): + other.append(self.db.model.get_connection()) + self.db.close() + + self.assertEqual(self._run_threads(grab), []) + # The main thread's connection survives closes made by other threads + self.assertIs(self.db.model.get_connection(), main_conn) + main_conn.execute("SELECT 1").fetchone() + self.db.close() + self.assertIsNot(self.db.model.get_connection(), main_conn) + + +class TestConcurrentWrites(ConcurrencyTestCase): + def test_distinct_prompts_saved_and_counted_from_many_threads(self): + def work(index): + for n in range(CALLS_PER_THREAD): + text = f"thread {index} prompt {n}" + pid = self.db.save_prompt( + text=text, prompt_hash=generate_prompt_hash(text) + ) + self.assertTrue(self.db.record_prompt_use(pid)) + + self.assertEqual(self._run_threads(work), []) + + info = self.db.model.get_database_info() + self.assertEqual(info["total_prompts"], THREADS * CALLS_PER_THREAD) + prompts = self.db.search_prompts(limit=THREADS * CALLS_PER_THREAD + 1) + self.assertEqual(len(prompts), THREADS * CALLS_PER_THREAD) + self.assertTrue(all(p["run_count"] == 1 for p in prompts)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lora_database.py b/tests/test_lora_database.py index 2ede390..5be8c9a 100644 --- a/tests/test_lora_database.py +++ b/tests/test_lora_database.py @@ -26,6 +26,7 @@ class LoraDBTestCase(unittest.TestCase): self.db = PromptDatabase(self.temp_db.name) def tearDown(self): + self.db.close() for suffix in ("", "-wal", "-shm"): path = self.temp_db.name + suffix if os.path.exists(path): diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py index 5a0ce6c..d11a74b 100644 --- a/tests/test_prompt_usage.py +++ b/tests/test_prompt_usage.py @@ -20,6 +20,8 @@ class UsageTestCase(unittest.TestCase): self.db = PromptDatabase(self.path) def tearDown(self): + if getattr(self, "db", None) is not None: + self.db.close() for suffix in ("", "-wal", "-shm"): if os.path.exists(self.path + suffix): os.unlink(self.path + suffix) @@ -160,6 +162,8 @@ class TestUsageMigration(unittest.TestCase): conn.execute("ALTER TABLE prompts DROP COLUMN run_count") def tearDown(self): + if getattr(self, "db", None) is not None: + self.db.close() for suffix in ("", "-wal", "-shm"): if os.path.exists(self.path + suffix): os.unlink(self.path + suffix)