ComfyAssets-ComfyUI_PromptManager/database/pair_ingestion.py

"""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 dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

from .models import normalize_image_path
from .operations import _coerce_file_info
from .workflow_documents import (
    TEXT_COLUMNS,
    DocumentReader,
    release_documents,
    resolve_image_row,
    stored_values,
)

try:
    from ..utils.hashing import generate_prompt_hash
    from ..utils.json_strict import dumps_strict
    from ..utils.prompt_pair_resolver import RESOLVER_VERSION
except ImportError:
    from utils.hashing import generate_prompt_hash
    from utils.json_strict import dumps_strict
    from utils.prompt_pair_resolver import RESOLVER_VERSION


def blank_positive(resolution: Dict[str, Any]) -> bool:
    """A resolved pair whose stored positive text is empty (#342).

    A PromptManager node whose editor is empty encodes only its prepend and
    append text, which is never stored as a prompt: such a result names no
    prompt to link.
    """
    if resolution.get("state") != "resolved":
        return False
    positive = resolution.get("positive") or {}
    text = positive.get("storage_text", positive.get("text"))
    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, each with its prompt's hash.

    Matched under the path and its resolved spelling (an output folder
    reached through a symlink or junction), each normalized and raw;
    deduplicated by row id.
    """
    rows: Dict[Any, Any] = {}
    for spelling in dict.fromkeys((image_path, _resolved(image_path))):
        for row in 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(spelling), spelling),
        ).fetchall():
            rows.setdefault(row["id"], row)
    return list(rows.values())


def _resolved(path: str) -> str:
    try:
        return os.path.realpath(path)
    except (OSError, ValueError):
        return path


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],
    *,
    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, marks=marks
        )


def _conflict() -> Dict[str, Any]:
    return {"status": "conflict", "pairs_added": 0, "links_added": 0}


def _ingest_resolved_pair_in(
    conn: Any,
    db: Any,
    image_path: str,
    metadata: Dict[str, Any],
    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``.

    The compatible-only rules are the public ingester's: an image with more
    than one exact row, another positive, or a prior pair with another
    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). ``marks`` as in
    ``ingest_resolved_pair``.
    """
    positive, negative = resolution["positive"], resolution["negative"]
    positive_text = positive.get("storage_text", positive["text"])
    negative_text = negative.get("storage_text", negative["text"])
    normalized = normalize_image_path(image_path)
    if provenance is None:
        provenance = {
            "resolver_version": RESOLVER_VERSION,
            "positive": positive,
            "negative": negative,
        }
    encoded = dumps_strict(provenance, sort_keys=True, separators=(",", ":"))
    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:
        prior = conn.execute(
            "SELECT n.hash AS negative_hash FROM image_prompt_pairs ip "
            "JOIN prompt_pairs pp ON pp.id=ip.pair_id "
            "LEFT JOIN prompts n ON n.id=pp.negative_prompt_id WHERE ip.image_id=?",
            (existing["id"],),
        ).fetchall()
        expected = (
            generate_prompt_hash(negative_text) if negative_text.strip() else None
        )
        if any(row["negative_hash"] != expected for row in prior):
            return _conflict()
    prior_pair = conn.execute(
        "SELECT pp.id FROM prompt_pairs pp "
        "JOIN prompts p ON p.id=pp.positive_prompt_id "
        "LEFT JOIN prompts n ON n.id=pp.negative_prompt_id "
        "WHERE p.hash=? AND n.hash IS ?",
        (
            generate_prompt_hash(positive_text),
            generate_prompt_hash(negative_text) if negative_text.strip() else None,
        ),
    ).fetchone()
    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)
    # ComfyUI's NaN/Infinity are stored as null: the snapshot is strict JSON.
    snapshots = {
        column: dumps_strict(metadata[key], sort_keys=True)
        for column, key in (
            ("workflow_data", "workflow"),
            ("prompt_metadata", "prompt"),
            ("parameters", "parameters"),
        )
        if key in metadata
    }
    info = {
        "file_size": file_size,
        "width": width,
        "height": height,
        "format": image_format,
    }
    snapshots.update({key: value for key, value in info.items() if value is not None})
    changed = False
    if existing:
        image_id = existing["id"]
        # Compared as stored text, whichever form the row is in.
        current = resolve_image_row(DocumentReader(conn), existing)
        updates = {
            key: value for key, value in snapshots.items() if current[key] != value
        }
        if updates:
            replaced = [
                existing[reference]
                for column, reference in TEXT_COLUMNS
                if column in updates
            ]
            updates = _document_form(conn, updates)
            conn.execute(
                "UPDATE generated_images SET "
                + ",".join(f"{key}=?" for key in updates)
                + " WHERE id=?",
                [*updates.values(), image_id],
            )
            release_documents(conn, replaced)
            changed = True
    else:
        stored = _document_form(conn, snapshots)
        columns = ["prompt_id", "image_path", "filename", "file_path", *stored]
        values = [
            pair["positive_prompt_id"],
            image_path,
            os.path.basename(image_path),
            normalized,
            *stored.values(),
        ]
        image_id = conn.execute(
            "INSERT INTO generated_images("
            + ",".join(columns)
            + ") VALUES("
            + ",".join("?" for _ in columns)
            + ")",
            values,
        ).lastrowid
    linked = conn.execute(
        "SELECT 1 FROM image_prompt_pairs WHERE pair_id=? "
        "AND image_id=? AND provenance=?",
        (pair_id, image_id, encoded),
    ).fetchone()
    db._link_image_to_pair(conn, pair_id, image_id, provenance)
    status = (
        "added"
        if existing is None
        else "enriched" if changed or not linked else "unchanged"
    )
    return {
        "status": status,
        "pairs_added": pairs_added,
        "links_added": int(not linked),
        "pair_id": pair_id,
        "image_id": image_id,
    }


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.

    Each of those columns becomes NULL beside its document id (new form);
    every other column is unchanged.
    """
    texts = {
        column: values[column]
        for column, _reference in TEXT_COLUMNS
        if column in values
    }
    rest = {key: value for key, value in values.items() if key not in texts}
    return {**rest, **stored_values(conn, texts)}