Vito Sansevero
d247e4c2e7
Review fixes for the staged LoRA import (#221): - The metadata files are listed and fingerprinted again right before the replace transaction, so a file changed, removed or added while the downloads ran aborts the import with nothing deleted (System job: plan-changed, Review again; legacy scan: an error frame in its own words). The tests now change real files behind a held download. - A legacy scan whose page closes stops staging before its next LoRA and keeps its records hold until staging has stopped; CivitAI images are written to a temp file beside the cache entry and moved in place, so a half-written image is never seen as cached. - Start refuses with plan-changed when the integration was turned off or moved to another install since the review. - A failed job keeps the skipped count and the names of the LoRAs that could not be prepared. - The replace binds the category instead of one parameter per prompt (SQLite's variable limit), /lora/status says whether a detected install can be used, and the legacy scan logs expected refusals as warnings.
199 lines
8.2 KiB
Python
199 lines
8.2 KiB
Python
"""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})
|