feat(system): keep the library in ComfyUI's user folder (#318)
* feat(system): resolve every library path in one place utils/library_location.py decides where the library lives: the PROMPT_MANAGER_DATA_DIR folder (validated: absolute, never inside a folder ComfyUI or Prompt Manager serves), an explicit database.default_path as before, a node-folder library that was not moved, else ComfyUI's user/__promptmanager. The database, diagnostics, config.json, logs and the LoRA image cache all ask it, so a new install never writes into the node folder and an existing one keeps every path until the user moves it. "Keep it here" is persisted as library.keep_in_node_folder. * feat(system): move the library as a verified Activity job POST /prompt_manager/system/library/move starts "Move library", an exclusive, non-cancellable job (database/library_move.py). It checks the target volume's free space, copies the database with SQLite's backup API from one read snapshot, verifies the copy (integrity check, every workflow document, the same rows in every table), copies the safety and pre-migration backups and config.json (owner-only, checked by size and SHA-256), then publishes under the old file's write lock and switches the live process (every PromptModel, the review and history observers, the API's plans and health report). The old prompts.db and config.json are renamed *.moved-to-<folder> only after the switch. A failure removes everything it wrote and leaves the old library in use; a marker in the new folder makes a restart after any crash pick the right library. GET /prompt_manager/system/library and the Overview's "location" say where the library lives and why; POST .../keep persists "Keep it here". * feat(system): show where the library lives and offer the move D129. System > Overview's Library health gains a "Library location" row (the database path, why it is there, the "Keep it here" choice and any PROMPT_MANAGER_DATA_DIR problem) with Move library for a node-folder library and Details otherwise. While a node-folder library is not kept there, the workspace shows the concept's .notice above every page: "Keep it here" saves the choice on the server and stops the notice; "Move library" opens the shared review (From, To, Size, what moves, what happens to the old files) and starts the exclusive job in Activity. The notice and the Overview read the location again when the job ends. The Overview comparison maps the inserted row; a new comparison holds the notice to the concept's Sources notice (markup, tokens, alignment, both themes and widths) with a mutation check. * feat(system): close the browser server's log files before its root goes The logs now live in the library's data folder, which for the browser test server is <root>/user/__promptmanager; Windows could not remove the temp root while the server still held prompt_manager.log open. The API test's node folder is spelled in lower case because public paths are case-folded on Windows. * feat(system): show the library location on the legacy admin page A line under the admin page's statistics names the database path and why it is there (ComfyUI's user folder, PROMPT_MANAGER_DATA_DIR, config.json, or the node folder with a pointer to Move library), read once from GET /prompt_manager/system/library. * feat(system): prove a fresh host keeps its library in user/__promptmanager The execution suite starts one pinned CPU ComfyUI with no configuration at all on an OS-chosen port: the library is created in user/__promptmanager (database and logs), nothing but byte code is added to the node folder, the location route reports the default folder, and ComfyUI's /userdata API serves none of it (traversal spellings, a listing, a system-user header). * feat(system): document where the library lives and how it moves README: where the library lives, the upgrade path (Move library / Keep it here), PROMPT_MANAGER_DATA_DIR, explicit paths unchanged, and the way back to an older version. TESTING.md: D129 and the library location contract (resolution order, move steps, crash windows, coded refusals, Windows handling). * wip(system): library move fix round in progress * fix(system): retry a fallback image link that races the library switch The image monitor's fallback and recent-prompt links went through PromptDatabase.link_image_to_prompt, which logs and returns 0 on any error, so survive_switch never saw a "database is locked" or closed connection during Move library and the link was lost in both files. The monitor now links through link_image_to_prompt_or_raise (one transaction per attempt); link_image_to_prompt keeps its log-and-return-0 contract for other callers. A link that returns 0 is no longer logged as "Successfully linked image 0". A switch that fails after repointing the models now closes the old file's detached connections instead of leaving them to GC: switch() adds them to a list the mover owns. The mocked retry test is replaced by a real-lock race asserting the image row lands in the new library only. * fix(system): say which source is the default and where a symlinked library lives Sources & scans: with no source configured, ComfyUI's output folder now reads "ComfyUI/output · watching · default" on its own row instead of a separate note line. GET /settings marks each effective source with in_comfyui (its public path is relative to the ComfyUI base) so only those get the "ComfyUI/" prefix; the scan still sends the server's public path. Library location: a library in the node folder reads through the node folder's unresolved place in custom_nodes, so a node folder symlinked to a checkout outside ComfyUI shows custom_nodes/promptmanager/prompts.db instead of prompts.db (System library row, Move review, admin Library line). Never absolute or "..", else the shared _public_path, which sources and jobs keep using unchanged.
This commit is contained in:
Vito
@@ -196,3 +196,14 @@ class ExecutionHistoryOperations:
|
||||
.fetchone()
|
||||
)
|
||||
return dict(row) if row is not None else None
|
||||
|
||||
|
||||
def rekey_history_restore(old_path: str, new_path: str) -> None:
|
||||
"""Move the observers of ``old_path`` to ``new_path`` (Move library)."""
|
||||
old, new = _database_identity(old_path), _database_identity(new_path)
|
||||
if old == new:
|
||||
return
|
||||
with _restore_lock:
|
||||
observers = _restore_observers.pop(old, None)
|
||||
if observers:
|
||||
_restore_observers.setdefault(new, {}).update(observers)
|
||||
|
||||
@@ -0,0 +1,768 @@
|
||||
"""Move library (#249 S6b): copy the node-folder library, verify it, switch.
|
||||
|
||||
``plan_move`` checks a ``utils.library_location`` location can move and
|
||||
sizes the copy; ``move_library`` does it, in this order:
|
||||
|
||||
1. **Disk check** on the target's volume (the copy, the backups and
|
||||
config.json, plus a margin). The source is never renamed across folders,
|
||||
so another drive is simply a copy.
|
||||
2. **Copy** with SQLite's backup API from a private connection holding one
|
||||
read transaction: the per-table counts, the workflow documents' digest
|
||||
and ``PRAGMA data_version`` describe exactly the copied snapshot. The copy
|
||||
is a temp file in the target folder, made self-contained
|
||||
(``journal_mode=DELETE``), fsynced and owner-only.
|
||||
3. **Verify**: ``verify_database_file`` (integrity check, schema, every
|
||||
workflow document decompressed against its hash, #315), the same row
|
||||
count in every table and the same documents.
|
||||
4. **Copy files**: every pooled safety backup and pre-migration backup,
|
||||
checked by size and SHA-256 after the copy (their times are kept, so the
|
||||
pool prunes in the same order).
|
||||
5. **Publish** under the old database's write lock (``BEGIN IMMEDIATE``) and
|
||||
the config.json lock: if any other connection wrote since the snapshot
|
||||
(``data_version``), steps 2–4 run again (``attempts``, then
|
||||
``library-busy``). Else config.json is copied (owner-only, checked), a
|
||||
``pending`` marker is written, the temp files take their final names
|
||||
(``os.replace`` inside the target folder), the marker turns ``moved``
|
||||
and ``switch`` points the process at the new library while the old file
|
||||
is still locked, so no write can land in it.
|
||||
6. **Retire**: the old file's WAL is checkpointed and the old ``prompts.db``
|
||||
and ``config.json`` are renamed ``*.moved-to-<folder>`` (never deleted;
|
||||
the old backups stay where they are). A file still open elsewhere is left
|
||||
in place and reported (``old-in-use``): the marker already makes a restart
|
||||
use the new library.
|
||||
|
||||
Any failure before the marker says ``moved`` removes everything the move
|
||||
wrote in the target folder and leaves the old library in use, untouched.
|
||||
A crash at any point leaves either the old library in use (the marker is
|
||||
absent or ``pending``; the next move replaces its own leftovers) or the new
|
||||
one (``moved``). ``MoveFailed`` carries a code and a message without paths.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import time
|
||||
import types
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from .models import BUSY_TIMEOUT_SECONDS, PromptModel
|
||||
|
||||
try:
|
||||
from ..utils import library_location as ll
|
||||
from ..utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
from utils import library_location as ll
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
logger = get_logger("prompt_manager.database.library_move")
|
||||
|
||||
TEMP_PREFIX = ".pm-move-"
|
||||
MOVED_INFIX = ".moved-to-"
|
||||
DISK_MARGIN_BYTES = 16 * 2**20
|
||||
DEFAULT_ATTEMPTS = 3
|
||||
_COPY_CHUNK = 4 * 2**20
|
||||
# Backups that live beside the database and move with it.
|
||||
BACKUP_INFIXES = (".backup_", ".pre-workflow-storage_")
|
||||
|
||||
MESSAGES = {
|
||||
"not-movable": (
|
||||
"This library is not in the node folder, so there is nothing to move."
|
||||
),
|
||||
"same-folder": "The new folder is the node folder itself.",
|
||||
"target-invalid": "The new folder cannot hold the library.",
|
||||
"target-has-library": (
|
||||
"The new folder already holds a Prompt Manager library. Nothing was moved."
|
||||
),
|
||||
"target-conflict": (
|
||||
"The new folder already holds files with the same names. Nothing was moved."
|
||||
),
|
||||
"insufficient-space": (
|
||||
"There is not enough free space for the copy. Nothing was moved."
|
||||
),
|
||||
"copy-failed": "The copy could not be written or did not match. Nothing was moved.",
|
||||
"verify-failed": "The copy did not pass verification. Nothing was moved.",
|
||||
"library-busy": (
|
||||
"Prompt Manager kept writing to the library while it was copied. "
|
||||
"Nothing was moved; try again when no workflow is running."
|
||||
),
|
||||
"switch-failed": (
|
||||
"Prompt Manager could not switch to the verified copy, so the copy was "
|
||||
"removed. Nothing was moved."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class MoveFailed(Exception):
|
||||
"""A move that did not happen (or did not finish), with a coded reason."""
|
||||
|
||||
def __init__(self, code: str, message: Optional[str] = None):
|
||||
self.code = code
|
||||
self.message = message or MESSAGES.get(code, "The library was not moved.")
|
||||
super().__init__(self.message)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MovePlan:
|
||||
location: "ll.LibraryLocation"
|
||||
db_path: str
|
||||
source_dir: str
|
||||
target_dir: str
|
||||
config_path: Optional[str]
|
||||
backups: Tuple[str, ...]
|
||||
db_bytes: int
|
||||
files_bytes: int
|
||||
|
||||
@property
|
||||
def required_bytes(self) -> int:
|
||||
return self.db_bytes + self.files_bytes + DISK_MARGIN_BYTES
|
||||
|
||||
@property
|
||||
def target_location(self) -> "ll.LibraryLocation":
|
||||
"""Where the library is once moved (an explicit config.json stays)."""
|
||||
source = self.location.target_source or ll.SOURCE_DEFAULT
|
||||
config = (
|
||||
self.location.config_path
|
||||
if self.location.config_explicit
|
||||
else os.path.join(self.target_dir, ll.CONFIG_NAME)
|
||||
)
|
||||
return ll.LibraryLocation(
|
||||
source,
|
||||
os.path.join(self.target_dir, ll.DATABASE_NAME),
|
||||
config,
|
||||
self.target_dir,
|
||||
self.target_dir,
|
||||
None,
|
||||
source,
|
||||
self.location.config_explicit,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoveResult:
|
||||
location: "ll.LibraryLocation"
|
||||
counts: Dict[str, int]
|
||||
files: List[str]
|
||||
renamed: List[str] = field(default_factory=list)
|
||||
retire_problem: Optional[str] = None
|
||||
|
||||
|
||||
def _size(path: str) -> int:
|
||||
try:
|
||||
return os.path.getsize(path)
|
||||
except OSError:
|
||||
return 0
|
||||
|
||||
|
||||
def _free_bytes(path: str) -> int:
|
||||
return shutil.disk_usage(path).free
|
||||
|
||||
|
||||
def _existing_ancestor(path: str) -> str:
|
||||
current = os.path.abspath(path)
|
||||
while not os.path.isdir(current):
|
||||
parent = os.path.dirname(current)
|
||||
if parent == current:
|
||||
break
|
||||
current = parent
|
||||
return current
|
||||
|
||||
|
||||
def backup_names(db_path: str) -> Tuple[str, ...]:
|
||||
"""The pooled and pre-migration backups beside ``db_path`` (names)."""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
base = os.path.basename(db_path)
|
||||
prefixes = tuple(base + infix for infix in BACKUP_INFIXES)
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return ()
|
||||
return tuple(
|
||||
sorted(
|
||||
name
|
||||
for name in names
|
||||
if name.startswith(prefixes) and os.path.isfile(os.path.join(folder, name))
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def plan_move(location) -> MovePlan:
|
||||
"""Check ``location`` can move to its target; raises MoveFailed."""
|
||||
if not location.movable:
|
||||
raise MoveFailed("not-movable")
|
||||
target = os.path.abspath(location.target_dir)
|
||||
source_dir = os.path.dirname(os.path.abspath(location.db_path))
|
||||
if ll.same_folder(target, source_dir) or os.path.normcase(
|
||||
os.path.realpath(target)
|
||||
) == os.path.normcase(os.path.realpath(source_dir)):
|
||||
raise MoveFailed("same-folder")
|
||||
if os.path.exists(target) and not os.path.isdir(target):
|
||||
raise MoveFailed("target-invalid")
|
||||
backups = backup_names(location.db_path)
|
||||
# config.json moves with a node-folder library; one named by
|
||||
# PROMPT_MANAGER_CONFIG_PATH is the user's and stays where it is.
|
||||
config = None if location.config_explicit else location.config_path
|
||||
config = config if config and os.path.isfile(config) else None
|
||||
if ll.read_marker(target).get("state") != ll.MARKER_PENDING:
|
||||
if os.path.exists(os.path.join(target, ll.DATABASE_NAME)):
|
||||
raise MoveFailed("target-has-library")
|
||||
names = [ll.CONFIG_NAME] if config else []
|
||||
if any(os.path.exists(os.path.join(target, n)) for n in names + list(backups)):
|
||||
raise MoveFailed("target-conflict")
|
||||
db_bytes = _size(location.db_path) + _size(location.db_path + "-wal")
|
||||
files_bytes = sum(_size(os.path.join(source_dir, n)) for n in backups)
|
||||
files_bytes += _size(config) if config else 0
|
||||
return MovePlan(
|
||||
location,
|
||||
os.path.abspath(location.db_path),
|
||||
source_dir,
|
||||
target,
|
||||
config,
|
||||
backups,
|
||||
db_bytes,
|
||||
files_bytes,
|
||||
)
|
||||
|
||||
|
||||
def space_problem(plan: MovePlan) -> Optional[Tuple[int, int]]:
|
||||
"""``(required, free)`` when the target's volume is short, else None."""
|
||||
free = _free_bytes(_existing_ancestor(plan.target_dir))
|
||||
if free < plan.required_bytes:
|
||||
return plan.required_bytes, free
|
||||
return None
|
||||
|
||||
|
||||
# ── copy and verify ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _tables(conn) -> List[str]:
|
||||
return [
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' "
|
||||
"AND name NOT LIKE 'sqlite_%' ORDER BY name"
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _describe(conn) -> Tuple[Dict[str, int], str]:
|
||||
"""Row counts per table and a digest of the workflow documents."""
|
||||
counts = {}
|
||||
for table in _tables(conn):
|
||||
# A name read from sqlite_master, quoted as an identifier.
|
||||
quoted = '"' + table.replace('"', '""') + '"'
|
||||
sql = f"SELECT COUNT(*) FROM {quoted}" # nosec B608
|
||||
counts[table] = conn.execute(sql).fetchone()[0]
|
||||
digest = hashlib.sha256()
|
||||
if "workflow_documents" in counts:
|
||||
for row in conn.execute(
|
||||
"SELECT id, sha256, byte_length FROM workflow_documents ORDER BY id"
|
||||
):
|
||||
digest.update(repr((row[0], bytes(row[1]), row[2])).encode())
|
||||
return counts, digest.hexdigest()
|
||||
|
||||
|
||||
def _fsync_file(path: str) -> None:
|
||||
with open(path, "rb+") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def _fsync_dir(path: str) -> None:
|
||||
if os.name != "posix":
|
||||
return
|
||||
fd = os.open(path, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(fd)
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
|
||||
def _private(path: str) -> None:
|
||||
if os.name == "posix":
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
|
||||
def _copy_snapshot(source: sqlite3.Connection, temp: str):
|
||||
"""Back up ``source``'s current snapshot to ``temp``; return what it holds."""
|
||||
source.execute("BEGIN")
|
||||
try:
|
||||
version = source.execute("PRAGMA data_version").fetchone()[0]
|
||||
counts, documents = _describe(source)
|
||||
destination = sqlite3.connect(temp)
|
||||
try:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
destination.close()
|
||||
finally:
|
||||
source.execute("COMMIT")
|
||||
with contextlib.closing(sqlite3.connect(temp)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
_fsync_file(temp)
|
||||
_private(temp)
|
||||
return version, counts, documents
|
||||
|
||||
|
||||
def _verify_copy(temp: str, counts, documents) -> None:
|
||||
ok, reason = PromptModel.verify_database_file(temp)
|
||||
if not ok:
|
||||
logger.error(f"Library move: the copy did not verify: {reason}")
|
||||
detail = reason if reason.startswith("Workflow document") else None
|
||||
raise MoveFailed(
|
||||
"verify-failed",
|
||||
f"{MESSAGES['verify-failed']} {detail}" if detail else None,
|
||||
)
|
||||
uri = Path(os.path.abspath(temp)).as_uri() + "?mode=ro"
|
||||
try:
|
||||
with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn:
|
||||
copied = _describe(conn)
|
||||
except sqlite3.Error as error:
|
||||
logger.error(f"Library move: the copy could not be read: {error}")
|
||||
raise MoveFailed("verify-failed") from error
|
||||
if copied != (counts, documents):
|
||||
logger.error("Library move: the copy's rows or documents differ")
|
||||
raise MoveFailed("verify-failed")
|
||||
|
||||
|
||||
def _hash_copy(source: str, destination: str) -> Tuple[int, str]:
|
||||
"""Copy ``source`` to ``destination`` (fsynced, owner-only, times kept)."""
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with open(source, "rb") as reader, open(destination, "wb") as writer:
|
||||
while True:
|
||||
chunk = reader.read(_COPY_CHUNK)
|
||||
if not chunk:
|
||||
break
|
||||
digest.update(chunk)
|
||||
size += len(chunk)
|
||||
writer.write(chunk)
|
||||
writer.flush()
|
||||
os.fsync(writer.fileno())
|
||||
stat = os.stat(source)
|
||||
os.utime(destination, ns=(stat.st_atime_ns, stat.st_mtime_ns))
|
||||
_private(destination)
|
||||
return size, digest.hexdigest()
|
||||
|
||||
|
||||
def _file_digest(path: str) -> Tuple[int, str]:
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with open(path, "rb") as reader:
|
||||
while True:
|
||||
chunk = reader.read(_COPY_CHUNK)
|
||||
if not chunk:
|
||||
return size, digest.hexdigest()
|
||||
digest.update(chunk)
|
||||
size += len(chunk)
|
||||
|
||||
|
||||
def _check_copies(copies) -> None:
|
||||
for source, temp, expected in copies:
|
||||
if _file_digest(temp) != expected or _file_digest(source) != expected:
|
||||
logger.error(
|
||||
f"Library move: {os.path.basename(source)} did not copy unchanged"
|
||||
)
|
||||
raise MoveFailed("copy-failed")
|
||||
|
||||
|
||||
def write_marker(
|
||||
directory: str, state: str, source_db: str, source_config: Optional[str] = None
|
||||
) -> str:
|
||||
"""Atomically write the move marker; return its path.
|
||||
|
||||
``from`` and ``config_from`` name the old files, so a restart after a
|
||||
crash between the switch and the rename can finish the rename.
|
||||
"""
|
||||
path = os.path.join(directory, ll.MARKER_NAME)
|
||||
temp = os.path.join(directory, f"{TEMP_PREFIX}{uuid.uuid4().hex}.marker")
|
||||
payload = json.dumps(
|
||||
{
|
||||
"version": 1,
|
||||
"state": state,
|
||||
"from": source_db,
|
||||
"config_from": source_config,
|
||||
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
},
|
||||
indent=2,
|
||||
).encode("utf-8")
|
||||
fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as handle:
|
||||
handle.write(payload)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
_private(temp)
|
||||
os.replace(temp, path)
|
||||
except BaseException:
|
||||
_remove(temp)
|
||||
raise
|
||||
_fsync_dir(directory)
|
||||
return path
|
||||
|
||||
|
||||
def _remove(*paths) -> None:
|
||||
for path in paths:
|
||||
try:
|
||||
if path and os.path.lexists(path):
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
logger.warning(f"Library move: could not remove {os.path.basename(path)}")
|
||||
|
||||
|
||||
def sweep_move_temps(directory: str) -> List[str]:
|
||||
"""Remove the temp files an interrupted move left in ``directory``."""
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(directory)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
if name.startswith(TEMP_PREFIX):
|
||||
_remove(os.path.join(directory, name))
|
||||
removed.append(name)
|
||||
return removed
|
||||
|
||||
|
||||
# ── the move ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _noop(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
|
||||
def switch_process(old_db: str, location, detached: list) -> None:
|
||||
"""Point this process at ``location``: resolver, every model, observers.
|
||||
|
||||
Adds the old file's connections, detached and still open, to
|
||||
``detached`` as soon as they are detached, so the mover closes them
|
||||
(``PromptModel.close_detached``) whether or not the rest succeeds.
|
||||
"""
|
||||
from .execution_history import rekey_history_restore
|
||||
from .review_changes import rekey_review_changes
|
||||
|
||||
ll.switch_to(location)
|
||||
detached.extend(PromptModel.repoint_instances(old_db, location.db_path))
|
||||
rekey_review_changes(old_db, location.db_path)
|
||||
rekey_history_restore(old_db, location.db_path)
|
||||
|
||||
|
||||
def unswitch_process(old_db: str, new_db: str) -> list:
|
||||
"""Undo ``switch_process`` (any part of it); returns detached connections."""
|
||||
from .execution_history import rekey_history_restore
|
||||
from .review_changes import rekey_review_changes
|
||||
|
||||
ll.clear_override()
|
||||
detached = PromptModel.repoint_instances(new_db, old_db)
|
||||
rekey_review_changes(new_db, old_db)
|
||||
rekey_history_restore(new_db, old_db)
|
||||
return detached
|
||||
|
||||
|
||||
class _Move:
|
||||
def __init__(self, plan, *, switch, config_lock, progress, checkpoint):
|
||||
self.plan = plan
|
||||
self.switch = switch
|
||||
self.config_lock = config_lock
|
||||
self.progress = progress
|
||||
self.checkpoint = checkpoint
|
||||
self.created_dir = False
|
||||
self.temps: List[str] = []
|
||||
self.published: List[str] = []
|
||||
self.marker: Optional[str] = None
|
||||
self.moved = False
|
||||
|
||||
def temp(self, suffix: str) -> str:
|
||||
path = os.path.join(
|
||||
self.plan.target_dir, f"{TEMP_PREFIX}{uuid.uuid4().hex}{suffix}"
|
||||
)
|
||||
self.temps.append(path)
|
||||
return path
|
||||
|
||||
def prepare_target(self) -> None:
|
||||
target = self.plan.target_dir
|
||||
self.created_dir = not os.path.isdir(target)
|
||||
ll.ensure_private_dir(target)
|
||||
sweep_move_temps(target)
|
||||
if ll.read_marker(target).get("state") == ll.MARKER_PENDING:
|
||||
# An earlier move stopped before it switched: its files are ours.
|
||||
for name in (ll.DATABASE_NAME, ll.CONFIG_NAME) + self.plan.backups:
|
||||
_remove(os.path.join(target, name))
|
||||
for suffix in ("-wal", "-shm", "-journal"):
|
||||
_remove(os.path.join(target, ll.DATABASE_NAME + suffix))
|
||||
|
||||
def copy_round(self, source):
|
||||
"""Steps 2–4 once: the verified temp copies and the snapshot version."""
|
||||
self.progress("Copying the database")
|
||||
temp_db = self.temp(".db")
|
||||
version, counts, documents = _copy_snapshot(source, temp_db)
|
||||
self.checkpoint("copied")
|
||||
self.progress("Verifying the copy")
|
||||
_verify_copy(temp_db, counts, documents)
|
||||
self.checkpoint("verified")
|
||||
copies = []
|
||||
for name in self.plan.backups:
|
||||
src = os.path.join(self.plan.source_dir, name)
|
||||
temp = self.temp(".bak")
|
||||
try:
|
||||
expected = _hash_copy(src, temp)
|
||||
except OSError as error:
|
||||
logger.error(f"Library move: copying {name} failed: {error}")
|
||||
raise MoveFailed("copy-failed") from error
|
||||
copies.append((src, temp, expected, name))
|
||||
if copies:
|
||||
self.progress(f"Copied {len(copies)} backup file(s)")
|
||||
self.checkpoint("files")
|
||||
_check_copies([(s, t, e) for s, t, e, _ in copies])
|
||||
return version, counts, temp_db, copies
|
||||
|
||||
def discard_round(self, temp_db, copies) -> None:
|
||||
_remove(temp_db, *[temp for _s, temp, _e, _n in copies])
|
||||
|
||||
def publish(self, temp_db, copies) -> List[str]:
|
||||
"""Under the locks: config.json, marker, final names, marker moved."""
|
||||
target = self.plan.target_dir
|
||||
files = []
|
||||
if self.plan.config_path:
|
||||
temp = self.temp(".config")
|
||||
try:
|
||||
expected = _hash_copy(self.plan.config_path, temp)
|
||||
except OSError as error:
|
||||
raise MoveFailed("copy-failed") from error
|
||||
_check_copies([(self.plan.config_path, temp, expected)])
|
||||
copies = copies + [(self.plan.config_path, temp, expected, ll.CONFIG_NAME)]
|
||||
self.marker = write_marker(
|
||||
target, ll.MARKER_PENDING, self.plan.db_path, self.plan.config_path
|
||||
)
|
||||
self.checkpoint("pending")
|
||||
for _src, temp, _expected, name in copies:
|
||||
final = os.path.join(target, name)
|
||||
os.replace(temp, final)
|
||||
self.published.append(final)
|
||||
files.append(name)
|
||||
final_db = os.path.join(target, ll.DATABASE_NAME)
|
||||
os.replace(temp_db, final_db)
|
||||
self.published.append(final_db)
|
||||
files.append(ll.DATABASE_NAME)
|
||||
_fsync_dir(target)
|
||||
self.checkpoint("published")
|
||||
write_marker(target, ll.MARKER_MOVED, self.plan.db_path, self.plan.config_path)
|
||||
self.moved = True
|
||||
self.checkpoint("marked")
|
||||
return files
|
||||
|
||||
def cleanup(self) -> None:
|
||||
if self.moved:
|
||||
return
|
||||
_remove(*self.temps, *self.published, self.marker)
|
||||
if self.created_dir:
|
||||
try:
|
||||
os.rmdir(self.plan.target_dir)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _unique(path: str) -> str:
|
||||
candidate, number = path, 1
|
||||
while os.path.lexists(candidate):
|
||||
candidate = f"{path}_{number}"
|
||||
number += 1
|
||||
return candidate
|
||||
|
||||
|
||||
def retire_old_files(plan: MovePlan) -> Tuple[List[str], Optional[str]]:
|
||||
"""Rename the old database and config.json; ``(renamed, problem code)``.
|
||||
|
||||
Every connection to the old file must be closed. The WAL is checkpointed
|
||||
first; when it cannot be emptied (another process still has the file
|
||||
open) nothing is renamed (``old-in-use``).
|
||||
"""
|
||||
old = plan.db_path
|
||||
folder = os.path.basename(os.path.normpath(plan.target_dir))
|
||||
try:
|
||||
with contextlib.closing(sqlite3.connect(old, timeout=1)) as conn:
|
||||
busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
except sqlite3.Error as error:
|
||||
logger.warning(
|
||||
f"Library move: the old database could not be checkpointed: {error}"
|
||||
)
|
||||
return [], "old-in-use"
|
||||
if busy or _size(old + "-wal") > 0:
|
||||
return [], "old-in-use"
|
||||
renamed = []
|
||||
try:
|
||||
for path in (old, plan.config_path):
|
||||
if not path or not os.path.isfile(path):
|
||||
continue
|
||||
new_name = _unique(path + MOVED_INFIX + folder)
|
||||
os.rename(path, new_name)
|
||||
renamed.append(os.path.basename(new_name))
|
||||
for suffix in ("-wal", "-shm"):
|
||||
if os.path.isfile(old + suffix) and _size(old + suffix) == 0:
|
||||
os.remove(old + suffix)
|
||||
except OSError as error:
|
||||
logger.warning(f"Library move: the old files could not be renamed: {error}")
|
||||
return renamed, "old-locked"
|
||||
return renamed, None
|
||||
|
||||
|
||||
def finish_interrupted_move(location) -> List[str]:
|
||||
"""Rename old files a crash left between the switch and the rename.
|
||||
|
||||
The marker of the folder in use says ``moved`` and names the old files;
|
||||
an old ``prompts.db`` (and config.json) still there under its own name is
|
||||
renamed ``*.moved-to-<folder>`` now, never deleted. Returns the new
|
||||
names (nothing when there is nothing to finish or the file is in use).
|
||||
"""
|
||||
if location.source not in (ll.SOURCE_DEFAULT, ll.SOURCE_ENVIRONMENT):
|
||||
return []
|
||||
marker = ll.read_marker(location.data_dir)
|
||||
old = marker.get("from")
|
||||
if marker.get("state") != ll.MARKER_MOVED or not isinstance(old, str):
|
||||
return []
|
||||
if not os.path.isfile(old) or os.path.basename(old) != ll.DATABASE_NAME:
|
||||
return []
|
||||
if ll.same_folder(os.path.dirname(old), location.data_dir):
|
||||
return []
|
||||
config = marker.get("config_from")
|
||||
if not (
|
||||
isinstance(config, str)
|
||||
and os.path.basename(config) == ll.CONFIG_NAME
|
||||
and os.path.isfile(config)
|
||||
and os.path.normcase(os.path.abspath(config))
|
||||
!= os.path.normcase(os.path.abspath(location.config_path))
|
||||
):
|
||||
config = None
|
||||
pending = types.SimpleNamespace(
|
||||
db_path=old, config_path=config, target_dir=location.data_dir
|
||||
)
|
||||
renamed, problem = retire_old_files(pending)
|
||||
if problem:
|
||||
logger.warning(f"Library move: the old files are still in use ({problem})")
|
||||
elif renamed:
|
||||
logger.info(f"Library move finished: renamed {', '.join(renamed)}")
|
||||
return renamed
|
||||
|
||||
|
||||
def move_library(
|
||||
plan: MovePlan,
|
||||
*,
|
||||
switch: Callable = switch_process,
|
||||
unswitch: Callable = unswitch_process,
|
||||
after_switch: Callable = _noop,
|
||||
config_lock=None,
|
||||
progress: Callable[[str], None] = _noop,
|
||||
checkpoint: Callable[[str], None] = _noop,
|
||||
attempts: int = DEFAULT_ATTEMPTS,
|
||||
) -> MoveResult:
|
||||
"""Move the library as planned (blocking; see the module docstring).
|
||||
|
||||
``switch(old_db, location, detached)`` points the process at the copy
|
||||
and adds the old connections it detached to ``detached``; if it raises,
|
||||
``unswitch(old_db, new_db)`` undoes it while both files are still
|
||||
locked, every detached connection is closed, the copy is removed and
|
||||
the move fails (``switch-failed``). ``after_switch()`` runs
|
||||
once the switch is final; it is best effort (logged, never fails the move).
|
||||
"""
|
||||
short = space_problem(plan)
|
||||
if short is not None:
|
||||
raise MoveFailed("insufficient-space")
|
||||
lock = config_lock if config_lock is not None else contextlib.nullcontext()
|
||||
move = _Move(
|
||||
plan,
|
||||
switch=switch,
|
||||
config_lock=lock,
|
||||
progress=progress,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
move.unswitch = unswitch
|
||||
location = plan.target_location
|
||||
try:
|
||||
move.prepare_target()
|
||||
files, counts = _copy_and_switch(move, location, attempts)
|
||||
except Exception:
|
||||
move.cleanup()
|
||||
raise
|
||||
try:
|
||||
after_switch()
|
||||
except Exception:
|
||||
logger.exception("Library move: a follow-up after the switch failed")
|
||||
from .execution_history import notify_history_restore
|
||||
|
||||
notify_history_restore(location.db_path)
|
||||
progress("Renaming the old files")
|
||||
renamed, problem = retire_old_files(plan)
|
||||
checkpoint("retired")
|
||||
return MoveResult(location, counts, files, renamed, problem)
|
||||
|
||||
|
||||
def _switch_guarded(move: _Move, location) -> list:
|
||||
"""Switch while the copy is locked too; undo it all if any part fails.
|
||||
|
||||
The copy's write lock keeps every repointed thread from writing to it
|
||||
until the switch is whole, so undoing a partial switch loses nothing.
|
||||
Returns the old file's detached connections (still open).
|
||||
"""
|
||||
plan = move.plan
|
||||
guard = sqlite3.connect(
|
||||
location.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None
|
||||
)
|
||||
try:
|
||||
guard.execute("BEGIN IMMEDIATE")
|
||||
detached: list = []
|
||||
try:
|
||||
move.switch(plan.db_path, location, detached)
|
||||
except Exception as error:
|
||||
logger.exception("Library move: switching failed; undoing it")
|
||||
move.moved = False # the copy and its marker are removed
|
||||
try:
|
||||
PromptModel.close_detached(
|
||||
move.unswitch(plan.db_path, location.db_path), logger
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Library move: undoing the switch failed")
|
||||
# The old file's connections a partial switch detached: the
|
||||
# models reopen it on next use, so close them now, not at GC.
|
||||
PromptModel.close_detached(detached, logger)
|
||||
raise MoveFailed("switch-failed") from error
|
||||
move.checkpoint("switched")
|
||||
return detached
|
||||
finally:
|
||||
if guard.in_transaction:
|
||||
guard.execute("ROLLBACK")
|
||||
guard.close()
|
||||
|
||||
|
||||
def _copy_and_switch(move: _Move, location, attempts: int):
|
||||
plan = move.plan
|
||||
for attempt in range(1, attempts + 1):
|
||||
source = sqlite3.connect(
|
||||
plan.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None
|
||||
)
|
||||
try:
|
||||
source.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
|
||||
version, counts, temp_db, copies = move.copy_round(source)
|
||||
source.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
if source.execute("PRAGMA data_version").fetchone()[0] != version:
|
||||
move.progress("The library changed during the copy; copying again")
|
||||
move.discard_round(temp_db, copies)
|
||||
continue
|
||||
move.checkpoint("locked")
|
||||
with move.config_lock:
|
||||
files = move.publish(temp_db, copies)
|
||||
detached = _switch_guarded(move, location)
|
||||
# Still holding the old file's lock: a write that was waiting
|
||||
# on it fails when its connection closes (and runs again on
|
||||
# the new library) instead of landing in the old file.
|
||||
PromptModel.close_detached(detached, logger)
|
||||
finally:
|
||||
if source.in_transaction:
|
||||
source.execute("ROLLBACK")
|
||||
return files, counts
|
||||
finally:
|
||||
source.close()
|
||||
raise MoveFailed("library-busy")
|
||||
+70
-15
@@ -609,19 +609,37 @@ class PromptModel:
|
||||
if conn is not None and self._local.generation == self._generation:
|
||||
return conn
|
||||
|
||||
# check_same_thread=False only so close_all can clean up other
|
||||
# threads' connections; each connection is still used by one thread.
|
||||
conn = sqlite3.connect(
|
||||
self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
||||
)
|
||||
conn.row_factory = sqlite3.Row
|
||||
configure_connection(conn)
|
||||
while True:
|
||||
# The path and generation are read together, and the new
|
||||
# connection is kept only if no Move library repointed the model
|
||||
# meanwhile: otherwise it would write to the old file.
|
||||
with self._conn_lock:
|
||||
path, generation = self.db_path, self._generation
|
||||
# check_same_thread=False only so close_all can clean up other
|
||||
# threads' connections; each connection is still used by one thread.
|
||||
conn = sqlite3.connect(
|
||||
path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
||||
)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
configure_connection(conn)
|
||||
except BaseException:
|
||||
conn.close()
|
||||
raise
|
||||
with self._conn_lock:
|
||||
if generation == self._generation:
|
||||
self._connections.add(conn)
|
||||
break
|
||||
conn.close()
|
||||
self._local.conn = conn
|
||||
self._local.generation = self._generation
|
||||
with self._conn_lock:
|
||||
self._connections.add(conn)
|
||||
self._local.generation = generation
|
||||
return conn
|
||||
|
||||
@property
|
||||
def generation(self) -> int:
|
||||
"""Bumped whenever every connection is closed or the model repointed."""
|
||||
return self._generation
|
||||
|
||||
@classmethod
|
||||
def close_all_instances(cls) -> None:
|
||||
"""Close every connection of every live model in this process."""
|
||||
@@ -630,22 +648,59 @@ class PromptModel:
|
||||
for model in models:
|
||||
model.close_all()
|
||||
|
||||
@classmethod
|
||||
def repoint_instances(cls, old_path: str, new_path: str) -> list:
|
||||
"""Point every model of ``old_path`` at ``new_path`` (Move library).
|
||||
|
||||
Every thread opens the new file on its next ``get_connection``. The
|
||||
old connections are detached, not closed, and returned: the caller
|
||||
closes them once nothing can block on the new file any more. Closing
|
||||
a connection another thread is using waits for its statement; a write
|
||||
waiting on the mover's lock then fails (and, through
|
||||
``PromptDatabase.survive_library_switch``, runs again on the new file).
|
||||
"""
|
||||
old_key = os.path.normcase(os.path.realpath(old_path))
|
||||
with cls._instances_lock:
|
||||
models = list(cls._instances)
|
||||
with cls._init_lock:
|
||||
cls._initialized_paths.add(os.path.normcase(os.path.realpath(new_path)))
|
||||
detached = []
|
||||
for model in models:
|
||||
if model._schema_key() == old_key:
|
||||
detached.extend(model._detach_connections(new_path))
|
||||
return detached
|
||||
|
||||
@staticmethod
|
||||
def close_detached(connections, logger=None) -> None:
|
||||
"""Close connections ``repoint_instances`` detached."""
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
if logger is not None:
|
||||
logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def close_all(self) -> None:
|
||||
"""Close every thread's connection; they reopen lazily on next use.
|
||||
|
||||
Needed before deleting the database file (Windows keeps open files
|
||||
locked) and by tests that ran queries through executor threads.
|
||||
"""
|
||||
self.close_detached(self._detach_connections(), self.logger)
|
||||
|
||||
def _detach_connections(self, new_path: Optional[str] = None) -> list:
|
||||
"""Forget every connection (optionally pointing the model elsewhere).
|
||||
|
||||
Returns them for the caller to close; threads reopen on next use.
|
||||
"""
|
||||
with self._conn_lock:
|
||||
if new_path is not None:
|
||||
self.db_path = new_path
|
||||
connections = list(self._connections)
|
||||
self._connections.clear()
|
||||
self._generation += 1
|
||||
self._local.conn = None
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
return connections
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's connection, if it has one.
|
||||
|
||||
+158
-107
@@ -46,6 +46,7 @@ try:
|
||||
from ..utils.logging_config import get_logger
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
from ..utils.json_strict import dumps_strict
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
@@ -54,6 +55,7 @@ except ImportError:
|
||||
from utils.logging_config import get_logger
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.json_strict import dumps_strict
|
||||
from utils import library_location
|
||||
|
||||
|
||||
# Subquery to fetch tags from junction table, embedded in SELECT statements
|
||||
@@ -99,40 +101,56 @@ def escape_like(term: str) -> str:
|
||||
)
|
||||
|
||||
|
||||
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
||||
"""Resolve the database path from config, falling back to defaults.
|
||||
def _configured_default_path():
|
||||
"""config.json's ``database.default_path``, or None.
|
||||
|
||||
When no explicit path is given, reads PromptManagerConfig.DEFAULT_DB_PATH
|
||||
(set via config.json). Relative paths are resolved against the extension
|
||||
root directory so the database location is predictable regardless of the
|
||||
process working directory (important for Docker deployments).
|
||||
The package-qualified config wins over a standalone ``py.config`` (a
|
||||
custom-node package uses its own configuration).
|
||||
"""
|
||||
try:
|
||||
try:
|
||||
from ..py.config import PromptManagerConfig
|
||||
except ImportError:
|
||||
from py.config import PromptManagerConfig
|
||||
except Exception:
|
||||
return None
|
||||
return getattr(PromptManagerConfig, "DEFAULT_DB_PATH", None)
|
||||
|
||||
|
||||
def resolve_library_location():
|
||||
"""Where the library lives now (``utils.library_location``, #249 S6b).
|
||||
|
||||
The image sources that rule out a ``PROMPT_MANAGER_DATA_DIR`` are read by
|
||||
the resolver itself, exactly as ``config_location`` reads them.
|
||||
"""
|
||||
return library_location.resolve_library(_configured_default_path())
|
||||
|
||||
|
||||
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
||||
"""Resolve the database path through the one library resolver.
|
||||
|
||||
With no explicit path, ``utils.library_location.resolve_library`` decides
|
||||
(``PROMPT_MANAGER_DATA_DIR``, an explicit ``database.default_path``, a
|
||||
node-folder library not moved yet, else ComfyUI's ``user/__promptmanager``).
|
||||
An explicit relative path is resolved against the extension root, as
|
||||
before, so the location does not depend on the working directory.
|
||||
|
||||
Args:
|
||||
db_path: Explicit database path, or None to use config/default.
|
||||
db_path: Explicit database path, or None to use the resolver.
|
||||
|
||||
Returns:
|
||||
Absolute path to the SQLite database file.
|
||||
Absolute path to the SQLite database file; its folder exists (a new
|
||||
data folder is created owner-only).
|
||||
"""
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
if db_path is None:
|
||||
try:
|
||||
try:
|
||||
from ..py.config import PromptManagerConfig
|
||||
except ImportError:
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
db_path = PromptManagerConfig.DEFAULT_DB_PATH
|
||||
except Exception:
|
||||
db_path = "prompts.db"
|
||||
|
||||
if not os.path.isabs(db_path):
|
||||
db_path = resolve_library_location().db_path
|
||||
elif not os.path.isabs(db_path):
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
db_path = os.path.normpath(os.path.join(extension_root, db_path))
|
||||
|
||||
# Ensure parent directory exists for custom paths
|
||||
parent = os.path.dirname(db_path)
|
||||
if parent:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
library_location.ensure_private_dir(parent)
|
||||
|
||||
return db_path
|
||||
|
||||
@@ -365,6 +383,25 @@ class PromptDatabase(
|
||||
self.model = PromptModel(db_path)
|
||||
self.logger.debug("Database operations initialized successfully")
|
||||
|
||||
def survive_library_switch(self, write):
|
||||
"""Run ``write()``; once more if it failed while the library moved.
|
||||
|
||||
A background write (run counting, the image monitor) that was waiting
|
||||
on the old file when Move library switched fails when its connection
|
||||
is closed. When the model's generation changed during the attempt,
|
||||
the write runs again on a fresh connection, which is the new library;
|
||||
any other failure is raised as it was. ``write`` must open its own
|
||||
transaction (every PromptDatabase method does).
|
||||
"""
|
||||
generation = self.model.generation
|
||||
try:
|
||||
return write()
|
||||
except Exception:
|
||||
if self.model.generation == generation:
|
||||
raise
|
||||
self.logger.info("Running a write again on the moved library")
|
||||
return write()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's database connection (see PromptModel.close)."""
|
||||
self.model.close()
|
||||
@@ -1580,103 +1617,117 @@ class PromptDatabase(
|
||||
metadata: Optional metadata about the image
|
||||
|
||||
Returns:
|
||||
int: The ID of the created image record
|
||||
int: The ID of the created image record, or 0 when nothing was
|
||||
linked (invalid or unknown prompt, already linked, or an error,
|
||||
which is logged)
|
||||
"""
|
||||
try:
|
||||
return self.link_image_to_prompt_or_raise(prompt_id, image_path, metadata)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error linking image to prompt {prompt_id}: {e}")
|
||||
return 0
|
||||
|
||||
def link_image_to_prompt_or_raise(
|
||||
self, prompt_id: str, image_path: str, metadata: Optional[Dict[str, Any]] = None
|
||||
) -> int:
|
||||
"""``link_image_to_prompt`` that raises database errors.
|
||||
|
||||
The image monitor links through this so a write that failed because
|
||||
Move library switched reaches ``survive_library_switch`` and runs
|
||||
again on the new library (#249 S6b). Each call is one transaction:
|
||||
a failed attempt rolls back, and a committed one does not raise.
|
||||
Returns 0 for an invalid or unknown prompt or an image already linked.
|
||||
"""
|
||||
# Validate that prompt_id is a valid integer and exists in prompts table
|
||||
try:
|
||||
# Convert prompt_id to integer if it's a string number
|
||||
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
||||
prompt_id_int = int(prompt_id)
|
||||
elif isinstance(prompt_id, int):
|
||||
prompt_id_int = prompt_id
|
||||
# Convert prompt_id to integer if it's a string number
|
||||
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
||||
prompt_id_int = int(prompt_id)
|
||||
elif isinstance(prompt_id, int):
|
||||
prompt_id_int = prompt_id
|
||||
else:
|
||||
# Handle temporary IDs like "temp_123456" - skip linking
|
||||
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
||||
self.logger.debug(
|
||||
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
||||
)
|
||||
return 0
|
||||
else:
|
||||
# Handle temporary IDs like "temp_123456" - skip linking
|
||||
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
||||
self.logger.debug(
|
||||
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
||||
)
|
||||
return 0
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Verify the prompt exists in the database
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
||||
self.logger.warning(
|
||||
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
self.logger.warning(
|
||||
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
||||
"image linking"
|
||||
)
|
||||
return 0
|
||||
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(
|
||||
# Verify the prompt exists in the database
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
self.logger.warning(
|
||||
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
||||
"image linking"
|
||||
)
|
||||
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", {})),
|
||||
},
|
||||
)
|
||||
|
||||
# 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,
|
||||
{
|
||||
"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,
|
||||
(
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
conn.commit()
|
||||
|
||||
if cursor.rowcount == 0:
|
||||
self.logger.debug(
|
||||
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
if cursor.rowcount == 0:
|
||||
self.logger.debug(
|
||||
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error linking image to prompt {prompt_id}: {e}")
|
||||
return 0
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
|
||||
@@ -62,3 +62,14 @@ def notify_review_change(db_path: str) -> None:
|
||||
callback()
|
||||
except Exception:
|
||||
_logger.warning("Needs review change observer failed", exc_info=True)
|
||||
|
||||
|
||||
def rekey_review_changes(old_path: str, new_path: str) -> None:
|
||||
"""Move the observers of ``old_path`` to ``new_path`` (Move library)."""
|
||||
old, new = _identity(old_path), _identity(new_path)
|
||||
if old == new:
|
||||
return
|
||||
with _lock:
|
||||
observers = _observers.pop(old, None)
|
||||
if observers:
|
||||
_observers.setdefault(new, {}).update(observers)
|
||||
|
||||
Reference in New Issue
Block a user