feat(ingestion): store a file's prompt in one transaction, one prompt per file

database/scanned_prompts.py links a file to the prompt of one text, creating
it when new, inside one BEGIN IMMEDIATE transaction: a file linked to
another prompt is a conflict, a failed link leaves no prompt behind and two
writers store one prompt. pair_ingestion gains file_link_rows,
link_conflict and Marks/mark_created; ingest_resolved_pair takes optional
marks for the components it creates. recorded_runs and named_prompt in
run_attribution become public (#247).

database/pair_ingestion.py

@@ -1,7 +1,14 @@
-"""Atomic metadata enrichment for a single known generation context."""
+"""Atomic metadata enrichment for a single known generation context.
+
+Also the helpers both ingestion store paths share (#247): the file's link
+rows and the one-prompt-per-file rule (``file_link_rows``,
+``link_conflict``) and the marks a prompt created by ingestion carries
+(``Marks``, ``mark_created``).
+"""
 
 import os
-from typing import Any, Dict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
 
 from .models import normalize_image_path
 from .operations import _coerce_file_info
@@ -37,16 +44,69 @@ def blank_positive(resolution: Dict[str, Any]) -> bool:
     return not (isinstance(text, str) and text.strip())
 
 
+@dataclass(frozen=True)
+class Marks:
+    """What a prompt ingestion creates carries, so it can be found and removed.
+
+    ``note`` is formatted with ``name``, the file's name.
+    """
+
+    category: str
+    tags: Tuple[str, ...]
+    note: str
+
+
+SCAN_MARKS = Marks("scanned", ("auto-scanned",), "Auto-scanned from {name}")
+RUN_MARKS = Marks("scanned", ("auto-scanned",), "Saved from a workflow run: {name}")
+
+
+def mark_created(conn: Any, db: Any, prompt_id: int, marks: Marks, name: str) -> None:
+    """Apply ``marks`` to a prompt row the caller just created on ``conn``."""
+    conn.execute(
+        "UPDATE prompts SET category = ?, notes = ? WHERE id = ?",
+        (marks.category, marks.note.format(name=name), prompt_id),
+    )
+    db._sync_prompt_tags(conn, prompt_id, list(marks.tags))
+
+
+def file_link_rows(conn: Any, image_path: str) -> List[Any]:
+    """The file's link rows (either spelling), each with its prompt's hash."""
+    return conn.execute(
+        "SELECT gi.*,p.hash AS positive_hash FROM generated_images gi "
+        "JOIN prompts p ON p.id=gi.prompt_id "
+        "WHERE gi.file_path=? OR gi.image_path=?",
+        (normalize_image_path(image_path), image_path),
+    ).fetchall()
+
+
+def link_conflict(rows: List[Any], prompt_hash: str) -> bool:
+    """One prompt per file: more than one row, or a row of another prompt."""
+    return len(rows) > 1 or bool(rows and rows[0]["positive_hash"] != prompt_hash)
+
+
 def ingest_resolved_pair(
-    db: Any, image_path: str, metadata: Dict[str, Any], resolution: Dict[str, Any]
+    db: Any,
+    image_path: str,
+    metadata: Dict[str, Any],
+    resolution: Dict[str, Any],
+    *,
+    marks: Optional[Marks] = None,
 ) -> Dict[str, Any]:
+    """Store a resolved pair and link the file to it in one transaction.
+
+    With ``marks`` (ingestion from files), each pair component this call
+    creates carries them; without (the PromptManager node path) components
+    are created bare, as always.
+    """
     state = resolution.get("state", "unresolved")
     if state == "resolved" and blank_positive(resolution):
         state = "unresolved"
     if state != "resolved":
         return {"status": state, "pairs_added": 0, "links_added": 0}
     with db._pair_transaction() as conn:
-        return _ingest_resolved_pair_in(conn, db, image_path, metadata, resolution)
+        return _ingest_resolved_pair_in(
+            conn, db, image_path, metadata, resolution, marks=marks
+        )
 
 
 def _conflict() -> Dict[str, Any]:
@@ -61,6 +121,7 @@ def _ingest_resolved_pair_in(
     resolution: Dict[str, Any],
     *,
     provenance: Dict[str, Any] | None = None,
+    marks: Optional[Marks] = None,
 ) -> Dict[str, Any]:
     """``ingest_resolved_pair`` inside the caller's open ``_pair_transaction``.
 
@@ -69,7 +130,8 @@ def _ingest_resolved_pair_in(
     negative is a ``conflict`` and nothing is written. The caller owns the
     transaction, so its own later failure rolls the link back too.
     ``provenance`` replaces the link's default resolver provenance (Needs
-    review records the user's selection there).
+    review records the user's selection there). ``marks`` as in
+    ``ingest_resolved_pair``.
     """
     positive, negative = resolution["positive"], resolution["negative"]
     positive_text = positive.get("storage_text", positive["text"])
@@ -82,15 +144,8 @@ def _ingest_resolved_pair_in(
             "negative": negative,
         }
     encoded = dumps_strict(provenance, sort_keys=True, separators=(",", ":"))
-    images = conn.execute(
-        "SELECT gi.*,p.hash AS positive_hash FROM generated_images gi "
-        "JOIN prompts p ON p.id=gi.prompt_id "
-        "WHERE gi.file_path=? OR gi.image_path=?",
-        (normalized, image_path),
-    ).fetchall()
-    if len(images) > 1 or (
-        images and images[0]["positive_hash"] != generate_prompt_hash(positive_text)
-    ):
+    images = file_link_rows(conn, image_path)
+    if link_conflict(images, generate_prompt_hash(positive_text)):
         return _conflict()
     existing = images[0] if images else None
     if existing:
@@ -115,7 +170,12 @@ def _ingest_resolved_pair_in(
             generate_prompt_hash(negative_text) if negative_text.strip() else None,
         ),
     ).fetchone()
-    pair_id = db._save_pair(conn, positive_text, negative_text)
+    if marks is None:
+        pair_id = db._save_pair(conn, positive_text, negative_text)
+    else:
+        pair_id = _save_marked_pair(
+            conn, db, positive_text, negative_text, marks, os.path.basename(image_path)
+        )
     pair = conn.execute("SELECT * FROM prompt_pairs WHERE id=?", (pair_id,)).fetchone()
     pairs_added = int(prior_pair is None)
     file_size, width, height, image_format = _coerce_file_info(metadata)
@@ -197,6 +257,26 @@ def _ingest_resolved_pair_in(
     }
 
 
+def _save_marked_pair(
+    conn: Any,
+    db: Any,
+    positive_text: str,
+    negative_text: str,
+    marks: Marks,
+    name: str,
+) -> int:
+    """``_save_pair`` that marks each component it creates."""
+    texts = [positive_text] + ([negative_text] if negative_text.strip() else [])
+    ids = []
+    for text in texts:
+        prompt_id, created = db._pair_component_outcome(conn, text)
+        if created:
+            mark_created(conn, db, prompt_id, marks, name)
+        ids.append(prompt_id)
+    negative_id = ids[1] if len(ids) > 1 else None
+    return db._save_pair_components(conn, ids[0], negative_id)[0]
+
+
 def _document_form(conn: Any, values: Dict[str, Any]) -> Dict[str, Any]:
     """``values`` with the workflow and prompt graph texts stored once.
 

database/run_attribution.py

@@ -164,7 +164,7 @@ def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]:
     return found
 
 
-def _recorded_runs(
+def recorded_runs(
     db: Any,
     image_paths: Iterable[str],
     output_root: Optional[str],
@@ -201,7 +201,8 @@ def _recorded_runs(
     return attributed
 
 
-def _named_prompt(db: Any, run_id: str) -> Optional[int]:
+def named_prompt(db: Any, run_id: str) -> Optional[int]:
+    """The one stored prompt every positive context of ``run_id`` names, else None."""
     contexts = db.get_execution_contexts(run_id).get("contexts") or []
     named = {
         context.get("positive_prompt_id")
@@ -228,11 +229,11 @@ def stored_run_prompts(
     paths = list(dict.fromkeys(path for path in image_paths if path))
     result: Dict[str, Optional[int]] = dict.fromkeys(paths)
     prompts: Dict[str, Optional[int]] = {}
-    for path, (run_id, _nodes) in _recorded_runs(
+    for path, (run_id, _nodes) in recorded_runs(
         db, paths, output_root, path_style
     ).items():
         if run_id not in prompts:
-            prompts[run_id] = _named_prompt(db, run_id)
+            prompts[run_id] = named_prompt(db, run_id)
         result[path] = prompts[run_id]
     return result
 
@@ -275,7 +276,7 @@ def recorded_output(
 ) -> Optional[RecordedOutput]:
     """The recorded run facts for ``image_path``, else None (#386).
 
-    The file must be attributed to one run (``_recorded_runs``: one run
+    The file must be attributed to one run (``recorded_runs``: one run
     reported its location and it was written during that run). Read-only; it
     never links anything.
     """
@@ -286,7 +287,7 @@ def recorded_output(
     if own_transaction:  # outputs and texts from one state of history
         conn.execute("BEGIN")
     try:
-        found = _recorded_runs(db, [image_path], output_root, path_style)
+        found = recorded_runs(db, [image_path], output_root, path_style)
         if image_path not in found:
             return None
         run_id, nodes = found[image_path]

database/scanned_prompts.py

@@ -0,0 +1,57 @@
+"""Link an output file to the prompt of one text, creating it when new (#247).
+
+The store step of ingestion from file metadata when the decision is a
+positive text (``utils.file_prompt``). The prompt insert and the link run in
+one ``BEGIN IMMEDIATE`` transaction: two writers never store one text twice,
+a failed link leaves no new prompt behind, and a file stays linked to one
+prompt (a file already linked to another prompt is a ``conflict``).
+"""
+
+import logging
+import os
+from typing import Any, Dict, Optional, Tuple
+
+from .pair_ingestion import Marks, file_link_rows, link_conflict, mark_created
+
+try:
+    from ..utils.hashing import generate_prompt_hash
+except ImportError:
+    from utils.hashing import generate_prompt_hash
+
+_LOGGER = logging.getLogger("prompt_manager.database.scanned_prompts")
+
+
+def store_scanned_prompt_or_raise(
+    db: Any, path: str, canonical: Dict[str, Any], text: str, marks: Marks
+) -> Tuple[str, Optional[int]]:
+    """``(status, prompt id)`` after linking ``path`` to ``text``'s prompt.
+
+    ``added`` (prompt created, carrying ``marks``), ``linked`` (stored text
+    reused), ``unchanged`` (already linked to it) or ``conflict`` (linked to
+    another prompt; nothing written, id None). Database errors are raised
+    and roll the whole store back, so a caller can run it again.
+    """
+    with db._pair_transaction() as conn:
+        rows = file_link_rows(conn, path)
+        if link_conflict(rows, generate_prompt_hash(text)):
+            return "conflict", None
+        if rows:
+            return "unchanged", rows[0]["prompt_id"]
+        prompt_id, created = db._pair_component_outcome(conn, text)
+        if created:
+            mark_created(conn, db, prompt_id, marks, os.path.basename(path))
+        db._insert_image_link(conn, prompt_id, path, canonical)
+        return ("added" if created else "linked"), prompt_id
+
+
+def store_scanned_prompt(
+    db: Any, path: str, canonical: Dict[str, Any], text: str, marks: Marks
+) -> Tuple[str, Optional[int]]:
+    """The scan's form: an error is logged and reported as ``("found", None)``."""
+    try:
+        return store_scanned_prompt_or_raise(db, path, canonical, text, marks)
+    except Exception as error:
+        _LOGGER.error(
+            f"Failed to store the prompt of {os.path.basename(path)}: {error}"
+        )
+        return "found", None

tests/test_run_attribution.py

@@ -326,7 +326,7 @@ class TestRecordedOutputNode(RecordedRuns, unittest.TestCase):
             self.assertIsNone(recorded_output_lookup(broken)(self.image))
             with (
                 mock.patch(
-                    "database.run_attribution._recorded_runs",
+                    "database.run_attribution.recorded_runs",
                     side_effect=sqlite3.OperationalError("locked"),
                 ),
                 self.assertLogs("prompt_manager.run_attribution", "WARNING"),

tests/test_scanned_prompts.py

@@ -0,0 +1,218 @@
+"""Storing a file's prompt from its text: one transaction, one prompt each (#247)."""
+
+import os
+import sqlite3
+import tempfile
+import threading
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from unittest import mock
+
+try:
+    from tests.open_handles import assert_closed
+    from tests import test_execution_history_database as history_tests
+except ImportError:  # discovered with tests/ as the top-level directory
+    from open_handles import assert_closed
+    import test_execution_history_database as history_tests
+
+from database.operations import PromptDatabase
+from database.pair_ingestion import (
+    RUN_MARKS,
+    SCAN_MARKS,
+    Marks,
+    ingest_resolved_pair,
+)
+from database.scanned_prompts import store_scanned_prompt, store_scanned_prompt_or_raise
+from utils.hashing import generate_prompt_hash
+
+META = {"prompt": {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "x"}}}}
+
+
+def resolved(positive="cat", negative="blur"):
+    return {
+        "state": "resolved",
+        "positive": {"text": positive, "provenance": {"source": "api"}},
+        "negative": {"text": negative, "provenance": {"source": "api"}},
+    }
+
+
+class StoreCase(unittest.TestCase):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.db = PromptDatabase(str(Path(self.temp.name) / "test.db"))
+        self.path = str(Path(self.temp.name) / "gen.png")
+
+    def tearDown(self):
+        self.db.close_all()
+        assert_closed(self, self.temp.name)
+        self.temp.cleanup()
+
+    def prompts(self):
+        rows = self.db.model.get_connection().execute(
+            "SELECT id, text FROM prompts ORDER BY id"
+        )
+        return [tuple(row) for row in rows]
+
+    def links(self):
+        rows = self.db.model.get_connection().execute(
+            "SELECT prompt_id, image_path FROM generated_images ORDER BY id"
+        )
+        return [tuple(row) for row in rows]
+
+    def save(self, text, **fields):
+        return self.db.save_prompt(
+            text, prompt_hash=generate_prompt_hash(text), **fields
+        )
+
+
+class TestStoreScannedPrompt(StoreCase):
+    def test_a_new_text_is_created_with_the_scan_marks_and_linked(self):
+        status, prompt_id = store_scanned_prompt_or_raise(
+            self.db, self.path, META, " a cat ", SCAN_MARKS
+        )
+        self.assertEqual(status, "added")
+        prompt = self.db.get_prompt_by_id(prompt_id)
+        self.assertEqual(prompt["text"], "a cat")
+        self.assertEqual(prompt["category"], "scanned")
+        self.assertEqual(prompt["tags"], ["auto-scanned"])
+        self.assertEqual(prompt["notes"], "Auto-scanned from gen.png")
+        self.assertEqual(prompt["run_count"], 0)
+        self.assertEqual(self.links(), [(prompt_id, self.path)])
+
+    def test_a_run_created_prompt_names_the_workflow_run(self):
+        _status, prompt_id = store_scanned_prompt_or_raise(
+            self.db, self.path, META, "a cat", RUN_MARKS
+        )
+        prompt = self.db.get_prompt_by_id(prompt_id)
+        self.assertEqual(prompt["notes"], "Saved from a workflow run: gen.png")
+        self.assertEqual(prompt["tags"], ["auto-scanned"])
+
+    def test_a_stored_text_is_reused_by_hash_and_keeps_its_fields(self):
+        existing = self.save("A Cat", category="mine", notes="kept")
+        status, prompt_id = store_scanned_prompt_or_raise(
+            self.db, self.path, META, "a cat ", SCAN_MARKS
+        )
+        self.assertEqual((status, prompt_id), ("linked", existing))
+        prompt = self.db.get_prompt_by_id(existing)
+        self.assertEqual((prompt["category"], prompt["notes"]), ("mine", "kept"))
+        self.assertEqual(prompt["tags"], [])
+        self.assertEqual(len(self.prompts()), 1)
+
+    def test_a_repeat_is_unchanged(self):
+        first = store_scanned_prompt_or_raise(
+            self.db, self.path, META, "a cat", SCAN_MARKS
+        )
+        again = store_scanned_prompt_or_raise(
+            self.db, self.path, META, "a cat", SCAN_MARKS
+        )
+        self.assertEqual(again, ("unchanged", first[1]))
+        self.assertEqual(len(self.links()), 1)
+
+    def test_a_file_linked_to_another_prompt_is_a_conflict(self):
+        other = self.save("a dog")
+        # Linked under another spelling of the same file.
+        spelled = os.path.join(self.temp.name, ".", "gen.png")
+        self.db.link_image_to_prompt(other, spelled)
+        result = store_scanned_prompt_or_raise(
+            self.db, self.path, META, "a cat", SCAN_MARKS
+        )
+        self.assertEqual(result, ("conflict", None))
+        self.assertEqual(self.prompts(), [(other, "a dog")])
+        self.assertEqual(self.links(), [(other, spelled)])
+
+    def test_a_failed_link_raises_and_leaves_no_prompt_behind(self):
+        with mock.patch.object(
+            self.db, "_insert_image_link", side_effect=sqlite3.OperationalError("x")
+        ):
+            with self.assertRaises(sqlite3.OperationalError):
+                store_scanned_prompt_or_raise(
+                    self.db, self.path, META, "a cat", SCAN_MARKS
+                )
+            with self.assertLogs("prompt_manager", "ERROR"):
+                result = store_scanned_prompt(
+                    self.db, self.path, META, "a cat", SCAN_MARKS
+                )
+        self.assertEqual(result, ("found", None))
+        self.assertEqual(self.prompts(), [])
+        self.assertEqual(self.links(), [])
+
+    def test_two_concurrent_writers_store_one_prompt_and_one_link(self):
+        hold = history_tests.FirstWriterHold()
+        outcome = self.db._pair_component_outcome
+        started = threading.Barrier(2)
+
+        def held_open(*args, **kwargs):
+            hold()
+            return outcome(*args, **kwargs)
+
+        def write(_index):
+            started.wait(timeout=30)
+            try:
+                return store_scanned_prompt_or_raise(
+                    self.db, self.path, META, "a cat", SCAN_MARKS
+                )[0]
+            finally:
+                self.db.model.close()
+
+        with (
+            mock.patch.object(
+                self.db, "_pair_component_outcome", side_effect=held_open
+            ),
+            ThreadPoolExecutor(max_workers=2) as pool,
+        ):
+            statuses = sorted(pool.map(write, range(2)))
+        self.assertEqual(statuses, ["added", "unchanged"])
+        self.assertEqual(len(self.prompts()), 1)
+        self.assertEqual(len(self.links()), 1)
+
+
+class TestPairMarks(StoreCase):
+    def test_marks_reach_only_created_components(self):
+        kept = self.save("cat", category="mine")
+        result = ingest_resolved_pair(
+            self.db, self.path, META, resolved(), marks=RUN_MARKS
+        )
+        self.assertEqual(result["status"], "added")
+        positive = self.db.get_prompt_by_id(kept)
+        self.assertEqual((positive["category"], positive["tags"]), ("mine", []))
+        negative_id = next(pid for pid, text in self.prompts() if text == "blur")
+        negative = self.db.get_prompt_by_id(negative_id)
+        self.assertEqual(negative["category"], "scanned")
+        self.assertEqual(negative["tags"], ["auto-scanned"])
+        self.assertEqual(negative["notes"], "Saved from a workflow run: gen.png")
+
+    def test_without_marks_created_components_stay_bare(self):
+        """Regression guard (passes on the base): the node path is unchanged."""
+        ingest_resolved_pair(self.db, self.path, META, resolved())
+        for prompt_id, _text in self.prompts():
+            prompt = self.db.get_prompt_by_id(prompt_id)
+            self.assertIsNone(prompt["category"])
+            self.assertEqual(prompt["tags"], [])
+
+    def test_a_blank_negative_creates_only_the_positive(self):
+        result = ingest_resolved_pair(
+            self.db, self.path, META, resolved(negative=""), marks=SCAN_MARKS
+        )
+        self.assertEqual(result["status"], "added")
+        self.assertEqual([text for _id, text in self.prompts()], ["cat"])
+
+    def test_marks_are_a_frozen_value(self):
+        self.assertEqual(
+            SCAN_MARKS, Marks("scanned", ("auto-scanned",), SCAN_MARKS.note)
+        )
+        with self.assertRaises(AttributeError):
+            SCAN_MARKS.category = "x"
+
+
+class TestPublicRunAttribution(unittest.TestCase):
+    def test_recorded_runs_and_named_prompt_are_public(self):
+        from database import run_attribution
+
+        self.assertTrue(callable(run_attribution.recorded_runs))
+        self.assertTrue(callable(run_attribution.named_prompt))
+        self.assertFalse(hasattr(run_attribution, "_recorded_runs"))
+
+
+if __name__ == "__main__":
+    unittest.main()