ComfyAssets-ComfyUI_PromptManager/database/category_import.py

"""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 .pairs import PairDependencyError
    from .workflow_documents import documents_of_rows, release_documents
except ImportError:  # pragma: no cover - flat test imports
    from database.maintenance import PlanChanged
    from database.pairs import PairDependencyError
    from database.workflow_documents import documents_of_rows, release_documents


class EmptyReplacement(ValueError):
    """The replacement has no record and would remove every current import."""


# The category's prompt ids, as a subquery (one bound parameter).
_IN_CATEGORY = "SELECT id FROM prompts WHERE category = ?"

# 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 = conn.execute(
                "SELECT COUNT(*) FROM prompts WHERE category = ? AND id IN "
                "(SELECT positive_prompt_id FROM prompt_pairs UNION "
                "SELECT negative_prompt_id FROM prompt_pairs)",
                (category,),
            ).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_category_unpaired(conn, category)
            removed = self._delete_category_rows(conn, category, 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, category: str, ids: List[int]) -> int:
        """Delete the category's prompts with their image links and tags.

        No commit. The rows are selected by category (just verified against
        ``ids`` under the write lock), so no statement binds one parameter per
        prompt and SQLite's variable limit never caps a collection.
        """
        if not ids:
            return 0
        documents = documents_of_rows(conn, "prompt_id", ids)  # chunked
        conn.execute(
            f"DELETE FROM generated_images WHERE prompt_id IN ({_IN_CATEGORY})",
            (category,),
        )
        release_documents(conn, documents)
        conn.execute(
            f"DELETE FROM prompt_tags WHERE prompt_id IN ({_IN_CATEGORY})",
            (category,),
        )
        return conn.execute(
            "DELETE FROM prompts WHERE category = ?", (category,)
        ).rowcount

    @staticmethod
    def _ensure_category_unpaired(conn, category: str) -> None:
        """Raise PairDependencyError when a saved pair uses a category prompt."""
        row = conn.execute(
            f"SELECT id FROM prompt_pairs WHERE positive_prompt_id IN ({_IN_CATEGORY}) "
            f"OR negative_prompt_id IN ({_IN_CATEGORY}) LIMIT 1",
            (category, category),
        ).fetchone()
        if row:
            raise PairDependencyError(
                f"Prompt belongs to saved pair {row[0]}; delete the pair first "
                "or save changed text as a new prompt"
            )

    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})