diff --git a/database/operations.py b/database/operations.py index d83a2b2..918e2a9 100644 --- a/database/operations.py +++ b/database/operations.py @@ -141,31 +141,43 @@ class PromptDatabase: f"Saving prompt: text_length={len(text)}, category={category}, tags={tags}, rating={rating}" ) - with self.model.get_connection() as conn: - cursor = conn.execute( - """ - INSERT INTO prompts ( - text, category, tags, rating, notes, hash, created_at, updated_at, - last_used_at - ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?) - """, - ( - text.strip(), - category, - rating, - notes, - prompt_hash, - datetime.datetime.now(datetime.timezone.utc).isoformat(), - datetime.datetime.now(datetime.timezone.utc).isoformat(), - utc_now_iso(), - ), + try: + with self.model.get_connection() as conn: + cursor = conn.execute( + """ + INSERT INTO prompts ( + text, category, tags, rating, notes, hash, created_at, + updated_at, last_used_at + ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?) + """, + ( + text.strip(), + category, + rating, + notes, + prompt_hash, + datetime.datetime.now(datetime.timezone.utc).isoformat(), + datetime.datetime.now(datetime.timezone.utc).isoformat(), + utc_now_iso(), + ), + ) + prompt_id = cursor.lastrowid + if tags: + self._sync_prompt_tags(conn, prompt_id, tags) + conn.commit() + self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}") + return prompt_id + except sqlite3.IntegrityError: + # Another caller inserted the same hash between the caller's + # duplicate check and this INSERT (the with-block rolled back). + existing = self.get_prompt_by_hash(prompt_hash) if prompt_hash else None + if existing is None: + raise + self.logger.debug( + f"Prompt with hash {prompt_hash} was saved concurrently; " + f"returning existing ID {existing['id']}" ) - prompt_id = cursor.lastrowid - if tags: - self._sync_prompt_tags(conn, prompt_id, tags) - conn.commit() - self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}") - return prompt_id + return existing["id"] def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool: """ diff --git a/tests/test_database.py b/tests/test_database.py index b26fdaa..b03daff 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -10,6 +10,7 @@ import os 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__)))) @@ -49,6 +50,31 @@ class DatabaseTestCase(unittest.TestCase): class TestPromptCRUD(DatabaseTestCase): """Test basic create, read, update, delete operations.""" + def test_save_existing_prompt_returns_its_id_when_check_raced(self): + # Two callers check the hash, both see nothing, both insert: the second + # insert hits the UNIQUE(hash) constraint and must yield the first id. + text = "raced prompt" + prompt_hash = generate_prompt_hash(text) + original = self.db.get_prompt_by_hash + calls = [] + + def racing_lookup(h): + calls.append(h) + return None if len(calls) == 1 else original(h) + + first_id = self._save(text) + with patch.object(self.db, "get_prompt_by_hash", racing_lookup): + self.assertIsNone(self.db.get_prompt_by_hash(prompt_hash)) + second_id = self.db.save_prompt(text=text, prompt_hash=prompt_hash) + + self.assertEqual(second_id, first_id) + self.assertEqual(len(self.db.search_prompts(text=text)), 1) + + def test_save_without_hash_never_collides(self): + first = self.db.save_prompt(text="no hash a") + second = self.db.save_prompt(text="no hash b") + self.assertNotEqual(first, second) + def test_save_and_retrieve(self): pid = self._save( "A beautiful sunset", category="nature", tags=["sunset", "sky"], rating=5 diff --git a/tests/test_database_concurrency.py b/tests/test_database_concurrency.py index c13781a..0f39c6d 100644 --- a/tests/test_database_concurrency.py +++ b/tests/test_database_concurrency.py @@ -105,6 +105,26 @@ class TestConcurrentWrites(ConcurrencyTestCase): self.assertEqual(len(prompts), THREADS * CALLS_PER_THREAD) self.assertTrue(all(p["run_count"] == 1 for p in prompts)) + def test_saving_the_same_prompt_from_many_threads_keeps_one_row(self): + text = "the same prompt from every thread" + prompt_hash = generate_prompt_hash(text) + ids = [] + + def work(_index): + for _ in range(CALLS_PER_THREAD): + pid = self.db.save_prompt(text=text, prompt_hash=prompt_hash) + ids.append(pid) + self.assertTrue(self.db.record_prompt_use(pid)) + + self.assertEqual(self._run_threads(work), []) + + self.assertEqual(len(ids), THREADS * CALLS_PER_THREAD) + self.assertEqual(len(set(ids)), 1) + self.assertEqual(self.db.model.get_database_info()["total_prompts"], 1) + prompt = self.db.get_prompt_by_hash(prompt_hash) + self.assertEqual(prompt["id"], ids[0]) + self.assertEqual(prompt["run_count"], THREADS * CALLS_PER_THREAD) + if __name__ == "__main__": unittest.main()