fix(lora): stage imports before replacing the collection (#221)

The LoRA import deleted the whole lora-manager category first and then saved each record in its own transaction, so a discovery, download or database failure left the previous collection gone or partial. The import now runs in three stages: discover and fingerprint the metadata files, stage every record (CivitAI downloads go into the image cache first, counted), then replace the category in one BEGIN IMMEDIATE transaction (replace_category_prompts). Any failure leaves the previous prompts, tags and image links exactly as they were. Prompts a saved pair uses block the replacement; an empty replacement of a non-empty collection needs a reviewed import. POST /prompt_manager/lora/import/preview reviews an import without writing or requesting anything; POST /lora/import/apply {token} runs it as the Activity job "LoRA reference import" (records admission, cancellable while it downloads, the safety backup policy before the replace, refused with plan-changed when the files changed since the review). The legacy POST /lora/scan runs the same stages with its SSE event shape.
2026-10-25 06:52:34 -07:00
parent 82f7fc8a6f
commit 6a2dae0f78
9 changed files with 1699 additions and 227 deletions
@@ -0,0 +1,170 @@
+"""Atomic replacement of one imported category (#221, #249 System S10).
+
+The LoRA reference import used to delete its whole category and then save
+each record in its own transaction, so a failure part way left the previous
+collection deleted and the new one partial. ``replace_category_prompts``
+does both in one ``BEGIN IMMEDIATE`` transaction: the category's current
+prompts (with their tags and image links) are replaced by the staged
+records, or, on any error, nothing changes.
+
+A staged record is ``{text, hash, tags, images}``. A record whose hash a
+prompt outside the category already has links its images to that prompt
+(the legacy import did the same); a second record with the hash of an
+earlier one links its images to that one. Neither is a new prompt.
+"""
+
+from typing import Any, Dict, List, Optional, Sequence
+
+try:
+    from .maintenance import PlanChanged
+    from .workflow_documents import documents_of_rows, release_documents
+except ImportError:  # pragma: no cover - flat test imports
+    from database.maintenance import PlanChanged
+    from database.workflow_documents import documents_of_rows, release_documents
+
+
+class EmptyReplacement(ValueError):
+    """The replacement has no record and would remove every current import."""
+
+
+# Each record's outcome and the result counter it adds to.
+_COUNTERS = {"imported": "imported", "linked": "linked", "duplicate": "duplicates"}
+
+
+def _placeholders(values: Sequence[Any]) -> str:
+    return ",".join("?" * len(values))
+
+
+class CategoryImportOperations:
+    """Read the state of an imported category and replace it atomically."""
+
+    def category_import_state(self, category: str) -> Dict[str, Any]:
+        """``{ids, paired}``: the category's prompt ids (ascending) and how
+        many of them a saved pair uses (those block a replacement)."""
+        with self.model.get_connection() as conn:
+            ids = [
+                row[0]
+                for row in conn.execute(
+                    "SELECT id FROM prompts WHERE category = ? ORDER BY id",
+                    (category,),
+                ).fetchall()
+            ]
+            paired = 0
+            if ids:
+                marks = _placeholders(ids)
+                paired = conn.execute(
+                    f"SELECT COUNT(*) FROM prompts WHERE id IN ({marks}) AND id IN "
+                    "(SELECT positive_prompt_id FROM prompt_pairs UNION "
+                    "SELECT negative_prompt_id FROM prompt_pairs)",
+                    ids,
+                ).fetchone()[0]
+        return {"ids": ids, "paired": paired}
+
+    def existing_prompt_hashes(
+        self, hashes: Sequence[str], exclude_category: str
+    ) -> set:
+        """The given hashes a prompt outside ``exclude_category`` already has."""
+        found = set()
+        wanted = list(dict.fromkeys(hashes))
+        with self.model.get_connection() as conn:
+            for start in range(0, len(wanted), 500):
+                chunk = wanted[start : start + 500]
+                rows = conn.execute(
+                    f"SELECT hash FROM prompts WHERE hash IN ({_placeholders(chunk)}) "
+                    "AND (category IS NULL OR category != ?)",
+                    [*chunk, exclude_category],
+                ).fetchall()
+                found.update(row[0] for row in rows)
+        return found
+
+    def replace_category_prompts(
+        self,
+        category: str,
+        records: List[Dict[str, Any]],
+        *,
+        expected_ids: Optional[List[int]] = None,
+        allow_empty: bool = False,
+    ) -> Dict[str, Any]:
+        """Replace every prompt of ``category`` with ``records`` in one transaction.
+
+        Raises (and changes nothing): ``PlanChanged`` when ``expected_ids`` is
+        given and the category's prompt ids differ from it;
+        ``PairDependencyError`` when a saved pair uses one of its prompts;
+        ``EmptyReplacement`` when ``records`` is empty while the category has
+        prompts, unless ``allow_empty`` (an empty replacement that was
+        reviewed). Any other error rolls everything back too.
+
+        Returns ``{removed, imported, linked, duplicates, images, rows}``:
+        prompts removed and created, records linked to an existing prompt or
+        to an earlier record with the same text, image links added, and per
+        record ``{name, outcome}`` (``imported``/``linked``/``duplicate``).
+        """
+        with self._pair_transaction() as conn:
+            ids = [
+                row[0]
+                for row in conn.execute(
+                    "SELECT id FROM prompts WHERE category = ? ORDER BY id",
+                    (category,),
+                ).fetchall()
+            ]
+            if expected_ids is not None and ids != list(expected_ids):
+                raise PlanChanged("The imported prompts changed since the review")
+            if not records and ids and not allow_empty:
+                raise EmptyReplacement(
+                    "Nothing to import; the previous imports were kept"
+                )
+            self._ensure_prompts_unpaired(conn, ids)
+            removed = self._delete_category_rows(conn, ids)
+            result = {
+                "removed": removed,
+                "imported": 0,
+                "linked": 0,
+                "duplicates": 0,
+                "images": 0,
+                "rows": [],
+            }
+            created = {}
+            for record in records:
+                self._import_record(conn, category, record, created, result)
+            return result
+
+    def _delete_category_rows(self, conn, ids: List[int]) -> int:
+        """Delete prompts ``ids`` with their image links and tags (no commit)."""
+        if not ids:
+            return 0
+        marks = _placeholders(ids)
+        documents = documents_of_rows(conn, "prompt_id", ids)
+        conn.execute(f"DELETE FROM generated_images WHERE prompt_id IN ({marks})", ids)
+        release_documents(conn, documents)
+        conn.execute(f"DELETE FROM prompt_tags WHERE prompt_id IN ({marks})", ids)
+        return conn.execute(f"DELETE FROM prompts WHERE id IN ({marks})", ids).rowcount
+
+    def _import_record(self, conn, category, record, created, result) -> None:
+        """Save one staged record on ``conn``, or link it to the prompt it repeats."""
+        prompt_hash = record["hash"]
+        name = record.get("name") or record["text"]
+        if prompt_hash in created:
+            prompt_id, outcome = created[prompt_hash], "duplicate"
+        else:
+            row = conn.execute(
+                "SELECT id FROM prompts WHERE hash = ?", (prompt_hash,)
+            ).fetchone()
+            if row is not None:
+                prompt_id, outcome = row[0], "linked"
+            else:
+                prompt_id = self._insert_prompt(
+                    conn,
+                    record["text"],
+                    category,
+                    list(record.get("tags") or []),
+                    None,
+                    None,
+                    prompt_hash,
+                )
+                outcome = "imported"
+            created[prompt_hash] = prompt_id
+        result[_COUNTERS[outcome]] += 1
+        for image in record.get("images") or []:
+            if self._insert_image_link(conn, prompt_id, image):
+                result["images"] += 1
+        result["rows"].append({"name": name, "outcome": outcome})
@@ -26,6 +26,7 @@ from .history_reads import HistoryReadOperations
 from .health import HealthOperations, iter_consistency_issues
 from .duplicate_copies import DuplicateCopyOperations
 from .workflow_storage import WorkflowStorageOperations
+from .category_import import CategoryImportOperations, EmptyReplacement  # noqa: F401
 from .workflow_documents import (
    DocumentReader,
    documents_of_rows,
@@ -366,6 +367,7 @@ class PromptDatabase(
    MaintenanceOperations,
    DuplicateCopyOperations,
    WorkflowStorageOperations,
+    CategoryImportOperations,
 ):
    """Database operations class for managing prompts."""

@@ -455,26 +457,9 @@ class PromptDatabase(

        try:
            with self.model.get_connection() as conn:
-                cursor = conn.execute(
-                    """
-                    INSERT INTO prompts (
-                        text, category, tags, rating, notes, hash, created_at,
-                        updated_at
-                    ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
-                    """,
-                    (
-                        text.strip(),
-                        category,
-                        rating,
-                        notes,
-                        prompt_hash,
-                        datetime.datetime.now(datetime.timezone.utc).isoformat(),
-                        datetime.datetime.now(datetime.timezone.utc).isoformat(),
-                    ),
+                prompt_id = self._insert_prompt(
+                    conn, text, category, tags, rating, notes, prompt_hash
                )
-                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
@@ -490,6 +475,32 @@ class PromptDatabase(
            )
            return existing["id"]

+    def _insert_prompt(
+        self,
+        conn: sqlite3.Connection,
+        text: str,
+        category: Optional[str],
+        tags: Optional[List[str]],
+        rating: Optional[int],
+        notes: Optional[str],
+        prompt_hash: Optional[str],
+    ) -> int:
+        """Insert one prompt row and its tags on ``conn`` (no commit)."""
+        now = datetime.datetime.now(datetime.timezone.utc).isoformat()
+        cursor = conn.execute(
+            """
+            INSERT INTO prompts (
+                text, category, tags, rating, notes, hash, created_at,
+                updated_at
+            ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
+            """,
+            (text.strip(), category, rating, notes, prompt_hash, now, now),
+        )
+        prompt_id = cursor.lastrowid
+        if tags:
+            self._sync_prompt_tags(conn, prompt_id, tags)
+        return prompt_id
+
    def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
        """
        Count runs of a prompt and mark it as most recently used.
@@ -1670,64 +1681,83 @@ class PromptDatabase(
                return 0

            # Proceed with linking
-            if not isinstance(metadata, dict):
-                metadata = {}
-            filename = os.path.basename(image_path)
-            file_path = normalize_image_path(image_path)
-            file_size, width, height, image_format = _coerce_file_info(metadata)
-            # Each workflow and prompt graph is stored once (new form).
-            documents = stored_values(
-                conn,
-                {
-                    "workflow_data": dumps_strict(metadata.get("workflow", {})),
-                    "prompt_metadata": dumps_strict(metadata.get("prompt", {})),
-                },
+            image_id = self._insert_image_link(
+                conn, prompt_id_int, image_path, metadata
            )
+            conn.commit()

-            # INSERT OR IGNORE skips a file already linked to this prompt
-            # (same prompt_id + normalised full path)
-            cursor = conn.execute(
-                """
-                INSERT OR IGNORE INTO generated_images
-                (prompt_id, image_path, filename, file_path, file_size, width,
-                 height, format, workflow_data, prompt_metadata, parameters,
-                 workflow_document_id, prompt_document_id)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
+            if not image_id:
+                self.logger.debug(
+                    f"Image {normalize_image_path(image_path)} already linked to "
+                    f"prompt {prompt_id_int}"
+                )
+                return 0
+
+            return image_id
+
+    def _insert_image_link(
+        self,
+        conn: sqlite3.Connection,
+        prompt_id_int: int,
+        image_path: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> int:
+        """Insert one image link on ``conn`` (no commit).
+
+        Returns the new image row's id, or 0 when the file is already linked
+        to this prompt (``INSERT OR IGNORE``).
+        """
+        if not isinstance(metadata, dict):
+            metadata = {}
+        filename = os.path.basename(image_path)
+        file_path = normalize_image_path(image_path)
+        file_size, width, height, image_format = _coerce_file_info(metadata)
+        # Each workflow and prompt graph is stored once (new form).
+        documents = stored_values(
+            conn,
+            {
+                "workflow_data": dumps_strict(metadata.get("workflow", {})),
+                "prompt_metadata": dumps_strict(metadata.get("prompt", {})),
+            },
+        )
+
+        # INSERT OR IGNORE skips a file already linked to this prompt
+        # (same prompt_id + normalised full path)
+        cursor = conn.execute(
+            """
+            INSERT OR IGNORE INTO generated_images
+            (prompt_id, image_path, filename, file_path, file_size, width,
+             height, format, workflow_data, prompt_metadata, parameters,
+             workflow_document_id, prompt_document_id)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                prompt_id_int,
+                image_path,
+                filename,
+                file_path,
+                file_size,
+                width,
+                height,
+                image_format,
+                documents["workflow_data"],
+                documents["prompt_metadata"],
+                dumps_strict(metadata.get("parameters", {})),
+                documents["workflow_document_id"],
+                documents["prompt_document_id"],
+            ),
+        )
+        if cursor.rowcount == 0:
+            # Already linked: a document made only for this row goes.
+            release_documents(
+                conn,
                (
-                    prompt_id_int,
-                    image_path,
-                    filename,
-                    file_path,
-                    file_size,
-                    width,
-                    height,
-                    image_format,
-                    documents["workflow_data"],
-                    documents["prompt_metadata"],
-                    dumps_strict(metadata.get("parameters", {})),
                    documents["workflow_document_id"],
                    documents["prompt_document_id"],
                ),
            )
-            if cursor.rowcount == 0:
-                # Already linked: a document made only for this row goes.
-                release_documents(
-                    conn,
-                    (
-                        documents["workflow_document_id"],
-                        documents["prompt_document_id"],
-                    ),
-                )
-            conn.commit()
-
-            if cursor.rowcount == 0:
-                self.logger.debug(
-                    f"Image {file_path} already linked to prompt {prompt_id_int}"
-                )
-                return 0
-
-            return cursor.lastrowid
+            return 0
+        return cursor.lastrowid

    def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
        """