348 lines
13 KiB
Python
348 lines
13 KiB
Python
"""Library location routes (#249 S6b): where the library lives, Move, Keep.
|
|
|
|
- ``GET /prompt_manager/system/library`` → ``{success, location}``: the
|
|
database's path (relative to ComfyUI, as the statistics show it), why it is
|
|
there (``source``: ``default``, ``environment``, ``config`` or
|
|
``node_folder``), whether it can move and where to (``target``), the
|
|
"Keep it here" choice (``keep``), whether to offer the move (``offer``:
|
|
movable and not kept), a ``problem`` code and the database's size. The
|
|
System Overview carries the same object as ``location``.
|
|
- ``POST /prompt_manager/system/library/move`` → 202 ``{job}``: the
|
|
``library_move`` job, ``exclusive`` (it runs alone and every other job,
|
|
legacy-page operations included, is refused while it runs), not
|
|
cancellable. It copies, verifies and switches the live library
|
|
(``database/library_move.py``) and ends with one truthful line; a failure
|
|
leaves the library in the node folder, unchanged. Refused before any job
|
|
starts: 400 ``not-movable``/``same-folder``/``target-invalid``, 409
|
|
``target-has-library``/``target-conflict``/``busy``, 507
|
|
``insufficient-space``.
|
|
- ``POST /prompt_manager/system/library/keep {keep: bool}`` persists
|
|
``library.keep_in_node_folder`` through the config store and answers the
|
|
new location.
|
|
|
|
Errors are ``{success:false, code, error}`` and never name a path.
|
|
"""
|
|
|
|
import asyncio
|
|
import functools
|
|
import os
|
|
|
|
from aiohttp import web
|
|
|
|
from .job_routes import busy_response
|
|
from .jobs import JobBusy
|
|
|
|
try:
|
|
from ...database import library_move as lm
|
|
from ...database.operations import resolve_library_location
|
|
from ...utils import library_location as ll
|
|
except ImportError:
|
|
from database import library_move as lm
|
|
from database.operations import resolve_library_location
|
|
from utils import library_location as ll
|
|
|
|
LIBRARY_MOVE_JOB_KIND = "library_move"
|
|
LIBRARY_MOVE_TITLE = "Move library"
|
|
LIBRARY_MOVE_GROUPS = ("files", "errors")
|
|
KEEP_BODY_MAX = 1024
|
|
|
|
_REFUSAL_STATUS = {
|
|
"not-movable": 400,
|
|
"same-folder": 400,
|
|
"target-invalid": 400,
|
|
"target-has-library": 409,
|
|
"target-conflict": 409,
|
|
"insufficient-space": 507,
|
|
}
|
|
STILL_THERE = "The library is still in the node folder, unchanged."
|
|
NOT_MOVED = "The library was not moved."
|
|
OLD_LEFT = {
|
|
"old-in-use": (
|
|
"The old database is still open elsewhere, so it was left in the node "
|
|
"folder; it is no longer used."
|
|
),
|
|
"old-locked": (
|
|
"The old files could not be renamed, so they were left in the node "
|
|
"folder; they are no longer used."
|
|
),
|
|
}
|
|
|
|
|
|
def _failure(message, status, code):
|
|
return web.json_response(
|
|
{"success": False, "code": code, "error": message}, status=status
|
|
)
|
|
|
|
|
|
def _plural(count, one, many=None):
|
|
return f"{count:,} {one if count == 1 else (many or one + 's')}"
|
|
|
|
|
|
def _same_file(first, second):
|
|
try:
|
|
return os.path.samefile(first, second)
|
|
except OSError:
|
|
return os.path.normcase(os.path.realpath(first)) == os.path.normcase(
|
|
os.path.realpath(second)
|
|
)
|
|
|
|
|
|
class LibraryLocationRoutesMixin:
|
|
"""Mixin for the /prompt_manager/system/library endpoints."""
|
|
|
|
# Replaced in tests with a resolver over a temp tree.
|
|
_library_resolver = staticmethod(resolve_library_location)
|
|
|
|
def _register_library_location_routes(self, routes):
|
|
@routes.get("/prompt_manager/system/library")
|
|
async def library_location_route(request):
|
|
return await self.get_library_location(request)
|
|
|
|
@routes.post("/prompt_manager/system/library/move")
|
|
async def library_move_route(request):
|
|
return await self.start_library_move(request)
|
|
|
|
@routes.post("/prompt_manager/system/library/keep")
|
|
async def library_keep_route(request):
|
|
return await self.keep_library(request)
|
|
|
|
# ── location ──────────────────────────────────────────────────────
|
|
|
|
def _live_location(self):
|
|
"""The resolved location, unless the live database is another file.
|
|
|
|
A database the API was handed (tests, the browser fixture server) or
|
|
one the resolver no longer names is reported as configured and is
|
|
never moved.
|
|
"""
|
|
location = self._library_resolver()
|
|
live = self.db.model.db_path
|
|
if _same_file(location.db_path, live):
|
|
return location
|
|
return ll.LibraryLocation(
|
|
ll.SOURCE_CONFIG,
|
|
live,
|
|
location.config_path,
|
|
os.path.dirname(live),
|
|
config_explicit=location.config_explicit,
|
|
)
|
|
|
|
def _location_payload(self, location):
|
|
"""The public view of ``location`` (blocking: stats the database)."""
|
|
from ..config import PromptManagerConfig
|
|
|
|
keep = bool(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
|
|
size = 0
|
|
try:
|
|
size = os.path.getsize(location.db_path)
|
|
except OSError:
|
|
pass
|
|
return {
|
|
"path": self._public_path(location.db_path),
|
|
# Where Move library would put it, and whether config.json goes too.
|
|
"target_source": location.target_source if location.movable else None,
|
|
"config_moves": location.movable and not location.config_explicit,
|
|
"source": location.source,
|
|
"movable": location.movable,
|
|
"keep": keep,
|
|
"offer": location.movable and not keep,
|
|
"target": (
|
|
self._public_path(location.target_dir) if location.movable else None
|
|
),
|
|
"problem": location.problem or self._node_folder_risk(location),
|
|
"size_bytes": size,
|
|
}
|
|
|
|
# The node folder (tests point it at a temp tree).
|
|
_extension_root = ll.EXTENSION_ROOT
|
|
|
|
def _node_folder_risk(self, location):
|
|
"""``inside-node-folder`` for a library that stays in the node folder.
|
|
|
|
A configured database path, or a config.json the resolver still reads
|
|
from the node folder, is deleted with the node by a reinstall, and
|
|
it is never offered a move.
|
|
"""
|
|
if location.source == ll.SOURCE_NODE_FOLDER:
|
|
return None
|
|
paths = [location.db_path] + (
|
|
[] if location.config_explicit else [location.config_path]
|
|
)
|
|
root = self._extension_root
|
|
for path in paths:
|
|
parent = os.path.dirname(os.path.abspath(path))
|
|
if ll.same_folder(parent, root) or os.path.normcase(
|
|
os.path.realpath(parent)
|
|
).startswith(os.path.normcase(os.path.realpath(root)) + os.sep):
|
|
return "inside-node-folder"
|
|
return None
|
|
|
|
def _library_payload(self):
|
|
return self._location_payload(self._live_location())
|
|
|
|
async def get_library_location(self, request):
|
|
try:
|
|
payload = await self._run_in_executor(self._library_payload)
|
|
except Exception:
|
|
self.logger.exception("Library location could not be read")
|
|
return _failure(
|
|
"The library location could not be read.", 500, "location-failed"
|
|
)
|
|
return web.json_response(
|
|
{"success": True, "location": payload},
|
|
headers={"Cache-Control": "no-store"},
|
|
)
|
|
|
|
# ── move ──────────────────────────────────────────────────────────
|
|
|
|
def _plan_library_move(self):
|
|
"""The checked plan (blocking); raises MoveFailed."""
|
|
plan = lm.plan_move(self._live_location())
|
|
if lm.space_problem(plan) is not None:
|
|
raise lm.MoveFailed("insufficient-space")
|
|
return plan
|
|
|
|
async def start_library_move(self, request):
|
|
blocker = self._jobs.blocker(LIBRARY_MOVE_JOB_KIND, admission="exclusive")
|
|
if blocker is not None:
|
|
return busy_response(JobBusy(blocker))
|
|
try:
|
|
await self._run_in_executor(self._plan_library_move)
|
|
except lm.MoveFailed as refused:
|
|
status = _REFUSAL_STATUS.get(refused.code, 400)
|
|
return _failure(refused.message, status, refused.code)
|
|
except Exception:
|
|
self.logger.exception("Move library could not be planned")
|
|
return _failure(NOT_MOVED, 500, "move-failed")
|
|
try:
|
|
job = self._jobs.start(
|
|
LIBRARY_MOVE_JOB_KIND,
|
|
LIBRARY_MOVE_TITLE,
|
|
admission="exclusive",
|
|
cancellable=False,
|
|
scope={},
|
|
groups=LIBRARY_MOVE_GROUPS,
|
|
runner=self._run_library_move,
|
|
)
|
|
except JobBusy as busy:
|
|
return busy_response(busy)
|
|
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
|
|
|
def _after_library_switch(self):
|
|
"""Forget what described the old file (best effort, after the switch).
|
|
|
|
Plans, the health report and the history read fence named the old
|
|
file; the fence is closed before it is dropped.
|
|
"""
|
|
self._database_replaced()
|
|
fence, self._history_read_fence = self._history_read_fence, None
|
|
if fence is not None:
|
|
fence.close()
|
|
|
|
def finish_interrupted_library_move(self):
|
|
"""Startup: rename old files a crash left after a move (blocking)."""
|
|
try:
|
|
location = self._library_resolver()
|
|
if _same_file(location.db_path, self.db.model.db_path):
|
|
return lm.finish_interrupted_move(location)
|
|
except Exception:
|
|
self.logger.exception("Could not finish an interrupted library move")
|
|
return []
|
|
|
|
async def _run_library_move(self, job):
|
|
from ..config_store import _COMMIT_LOCK
|
|
|
|
loop = asyncio.get_running_loop()
|
|
|
|
def progress(line):
|
|
def show():
|
|
job.log_line(line)
|
|
job.update(status_line=line)
|
|
|
|
loop.call_soon_threadsafe(show)
|
|
|
|
job.update(status_line="Checking the new folder")
|
|
try:
|
|
plan = await self._run_in_executor(self._plan_library_move)
|
|
target = self._public_path(plan.target_dir)
|
|
job.log_line(f"Moving the library to {target}")
|
|
result = await self._run_in_executor(
|
|
functools.partial(
|
|
lm.move_library,
|
|
plan,
|
|
after_switch=self._after_library_switch,
|
|
config_lock=_COMMIT_LOCK,
|
|
progress=progress,
|
|
)
|
|
)
|
|
except lm.MoveFailed as failed:
|
|
self.logger.error(f"Move library failed: {failed.code}")
|
|
return self._fail_library_move(
|
|
job, failed.code, f"{failed.message} {STILL_THERE}"
|
|
)
|
|
except Exception:
|
|
self.logger.exception("Move library failed")
|
|
return self._fail_library_move(job, None, f"{NOT_MOVED} {STILL_THERE}")
|
|
await asyncio.sleep(0) # the progress lines come first
|
|
job.record_counted("files", len(result.files))
|
|
line = (
|
|
f"Moved the library to {target}: "
|
|
f"{_plural(result.counts.get('prompts', 0), 'prompt')} and "
|
|
f"{_plural(result.counts.get('generated_images', 0), 'image')} verified"
|
|
)
|
|
if result.retire_problem:
|
|
line += ". " + OLD_LEFT.get(result.retire_problem, OLD_LEFT["old-locked"])
|
|
elif result.renamed:
|
|
line += (
|
|
f". The old files were renamed {', '.join(result.renamed)}; "
|
|
"nothing was deleted"
|
|
)
|
|
line += ". Logs move at the next restart."
|
|
job.log_line(line)
|
|
job.update(progress=100, status_line=line)
|
|
|
|
@staticmethod
|
|
def _fail_library_move(job, code, message):
|
|
from .maintenance import _fail_job
|
|
|
|
_fail_job(job, code, message)
|
|
|
|
# ── keep it here ─────────────────────────────────────────────────
|
|
|
|
async def keep_library(self, request):
|
|
from ..config import PromptManagerConfig
|
|
from ..config_store import get_config_store
|
|
|
|
try:
|
|
if (request.content_length or 0) > KEEP_BODY_MAX:
|
|
raise ValueError("too large")
|
|
data = await request.json()
|
|
except Exception:
|
|
data = None
|
|
if not isinstance(data, dict) or not isinstance(data.get("keep"), bool):
|
|
return _failure("keep must be true or false.", 400, "invalid-request")
|
|
keep = data["keep"]
|
|
|
|
def mutate(config):
|
|
library = config.get("library")
|
|
library = dict(library) if isinstance(library, dict) else {}
|
|
library["keep_in_node_folder"] = keep
|
|
config["library"] = library
|
|
return config
|
|
|
|
def publish(_config):
|
|
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep
|
|
|
|
try:
|
|
await self._run_in_executor(
|
|
functools.partial(
|
|
get_config_store().commit,
|
|
mutate,
|
|
expected_revision=None,
|
|
publish=publish,
|
|
)
|
|
)
|
|
except Exception as exc:
|
|
return self._config_commit_failure(exc)
|
|
return await self.get_library_location(request)
|