Vito Sansevero
edf7811454
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".
437 lines
15 KiB
Python
437 lines
15 KiB
Python
"""Where the library lives (#249 S6b): the one resolver for every data path.
|
|
|
|
The library is ``prompts.db``, ``config.json`` (it can hold the CivitAI key),
|
|
the safety backups beside the database, the logs and the LoRA image cache.
|
|
Inside the node folder an uninstall, a reinstall or a folder-replacing update
|
|
deletes it, so new installs keep it in ComfyUI's system user folder
|
|
``user/__promptmanager`` (``folder_paths.get_system_user_directory``), which
|
|
ComfyUI's ``/userdata`` HTTP API refuses to serve.
|
|
|
|
Resolution order (``resolve_library``):
|
|
|
|
1. ``PROMPT_MANAGER_DATA_DIR``: an absolute folder outside every folder
|
|
ComfyUI or Prompt Manager serves. It is used when it holds a library or
|
|
when no node-folder library exists; an empty one never replaces a
|
|
node-folder library (that one stays in use, and Move library targets the
|
|
environment folder). An invalid value is ignored and reported.
|
|
2. An explicit ``database.default_path`` in config.json, honoured as before
|
|
(a relative path is relative to the node folder). ``PROMPT_MANAGER_CONFIG_PATH``
|
|
still names config.json itself.
|
|
3. The default data folder once a move finished there (its marker says
|
|
``moved``), even before the old file was renamed.
|
|
4. A node-folder ``prompts.db`` that was not moved: it stays in use until the
|
|
user moves it.
|
|
5. The default data folder (new installs). Outside ComfyUI (no
|
|
``folder_paths``) the node folder, as before.
|
|
|
|
Nothing here moves, copies or creates data: it only reads the file system.
|
|
After a live move the process uses ``switch_to``'s location until restart.
|
|
This module imports only the standard library, so the logger can use it.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import Iterable, List, Mapping, Optional, Tuple
|
|
|
|
DATA_DIR_ENV = "PROMPT_MANAGER_DATA_DIR"
|
|
CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH"
|
|
EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
|
|
SYSTEM_USER_NAME = "promptmanager"
|
|
SYSTEM_USER_PREFIX = "__"
|
|
DATABASE_NAME = "prompts.db"
|
|
CONFIG_NAME = "config.json"
|
|
MARKER_NAME = "library-location.json"
|
|
# config.json has always written its default; that value is not a choice.
|
|
DEFAULT_DB_SETTING = "prompts.db"
|
|
|
|
SOURCE_ENVIRONMENT = "environment"
|
|
SOURCE_CONFIG = "config"
|
|
SOURCE_NODE_FOLDER = "node_folder"
|
|
SOURCE_DEFAULT = "default"
|
|
SOURCES = (SOURCE_ENVIRONMENT, SOURCE_CONFIG, SOURCE_NODE_FOLDER, SOURCE_DEFAULT)
|
|
|
|
MARKER_PENDING = "pending"
|
|
MARKER_MOVED = "moved"
|
|
|
|
EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
_UNSET = object()
|
|
_override_lock = threading.Lock()
|
|
_override: Optional["LibraryLocation"] = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LibraryLocation:
|
|
"""Where the library is and why.
|
|
|
|
``target_dir`` is where Move library would put it (None outside ComfyUI)
|
|
and ``target_source`` why (the environment folder or the default one);
|
|
``problem`` is a code for something the user should know (an invalid or
|
|
empty ``PROMPT_MANAGER_DATA_DIR``, a move that did not finish).
|
|
``config_explicit``: config.json is named by ``PROMPT_MANAGER_CONFIG_PATH``
|
|
(it is the user's and never moves).
|
|
"""
|
|
|
|
source: str
|
|
db_path: str
|
|
config_path: str
|
|
data_dir: str
|
|
target_dir: Optional[str] = None
|
|
problem: Optional[str] = None
|
|
target_source: Optional[str] = None
|
|
config_explicit: bool = False
|
|
|
|
@property
|
|
def movable(self) -> bool:
|
|
return self.source == SOURCE_NODE_FOLDER and self.target_dir is not None
|
|
|
|
|
|
def import_folder_paths():
|
|
"""ComfyUI's ``folder_paths`` module, or None outside ComfyUI."""
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
return None
|
|
return folder_paths
|
|
|
|
|
|
def _folder_paths(value):
|
|
return import_folder_paths() if value is _UNSET else value
|
|
|
|
|
|
def _call(module, name, *args) -> Optional[str]:
|
|
getter = getattr(module, name, None)
|
|
if not callable(getter):
|
|
return None
|
|
try:
|
|
value = getter(*args)
|
|
except Exception: # an older or broken host: treat the folder as unknown
|
|
return None
|
|
return str(value) if value else None
|
|
|
|
|
|
def default_data_dir(folder_paths=_UNSET) -> Optional[str]:
|
|
"""``user/__promptmanager``, or None outside ComfyUI.
|
|
|
|
ComfyUI builds without ``get_system_user_directory`` get the same folder
|
|
joined by hand.
|
|
"""
|
|
module = _folder_paths(folder_paths)
|
|
if module is None:
|
|
return None
|
|
folder = _call(module, "get_system_user_directory", SYSTEM_USER_NAME)
|
|
if folder:
|
|
return os.path.abspath(folder)
|
|
user = _call(module, "get_user_directory")
|
|
if not user:
|
|
return None
|
|
return os.path.abspath(os.path.join(user, SYSTEM_USER_PREFIX + SYSTEM_USER_NAME))
|
|
|
|
|
|
def _identity(path: str) -> Optional[Tuple[int, int]]:
|
|
try:
|
|
stat = os.stat(path)
|
|
except (OSError, ValueError):
|
|
return None
|
|
return (stat.st_dev, stat.st_ino)
|
|
|
|
|
|
def same_folder(first: str, second: str) -> bool:
|
|
"""Whether two spellings name one existing folder (links and case folded)."""
|
|
one, two = _identity(first), _identity(second)
|
|
if one is None or two is None:
|
|
return False
|
|
if one[1] == 0: # no file id on this volume: compare the spelling
|
|
return _spelling(first) == _spelling(second)
|
|
return one == two
|
|
|
|
|
|
def _spelling(path: str) -> str:
|
|
return os.path.normcase(os.path.realpath(path))
|
|
|
|
|
|
def _existing_ancestors(path: str) -> List[str]:
|
|
"""``path`` (resolved) and its parents that exist, nearest first."""
|
|
current = os.path.realpath(path)
|
|
found = []
|
|
while True:
|
|
if os.path.isdir(current):
|
|
found.append(current)
|
|
parent = os.path.dirname(current)
|
|
if parent == current:
|
|
return found
|
|
current = parent
|
|
|
|
|
|
def _inside_any(path: str, folders: Iterable[str]) -> bool:
|
|
"""Whether ``path`` is one of ``folders`` or below one.
|
|
|
|
Compared by spelling (realpath + normcase) and by file identity of every
|
|
existing ancestor, so symlinks and case-insensitive volumes (Windows,
|
|
macOS) cannot hide a served folder.
|
|
"""
|
|
folders = [folder for folder in folders if folder]
|
|
if not folders:
|
|
return False
|
|
spelling = _spelling(path)
|
|
sep = os.path.sep
|
|
for folder in folders:
|
|
root = _spelling(folder)
|
|
if spelling == root or spelling.startswith(root.rstrip(sep) + sep):
|
|
return True
|
|
identities = {_identity(folder) for folder in folders}
|
|
identities.discard(None)
|
|
identities = {item for item in identities if item[1] != 0}
|
|
return any(_identity(a) in identities for a in _existing_ancestors(path))
|
|
|
|
|
|
def _served_folders(module, environ, gallery_roots) -> List[str]:
|
|
"""Folders whose files ComfyUI or Prompt Manager hand out over HTTP."""
|
|
folders = []
|
|
if module is not None:
|
|
for name in (
|
|
"get_output_directory",
|
|
"get_input_directory",
|
|
"get_temp_directory",
|
|
):
|
|
folder = _call(module, name)
|
|
if folder:
|
|
folders.append(folder)
|
|
extra = (environ or {}).get(EXTRA_GALLERY_ROOTS_ENV, "")
|
|
folders.extend(part for part in extra.split(os.pathsep) if part.strip())
|
|
folders.extend(str(root) for root in gallery_roots or () if root)
|
|
return folders
|
|
|
|
|
|
def _in_public_user_folder(path: str, module) -> bool:
|
|
"""Whether ``path`` lies in a ``/userdata``-served user folder (not ``__``)."""
|
|
user = _call(module, "get_user_directory") if module is not None else None
|
|
if not user:
|
|
return False
|
|
if not _inside_any(path, [user]) or same_folder(path, user):
|
|
return False
|
|
real_user = _spelling(user)
|
|
relative = os.path.relpath(_spelling(path), real_user)
|
|
first = relative.split(os.sep, 1)[0]
|
|
if first.startswith(os.pardir):
|
|
# Reached through a link or another spelling: judge by identity.
|
|
for ancestor in _existing_ancestors(path):
|
|
if same_folder(os.path.dirname(ancestor), user):
|
|
first = os.path.basename(ancestor)
|
|
break
|
|
else:
|
|
return False
|
|
return not first.startswith(SYSTEM_USER_PREFIX)
|
|
|
|
|
|
def data_dir_problem(
|
|
path, *, folder_paths=_UNSET, environ=None, gallery_roots=()
|
|
) -> Optional[str]:
|
|
"""Why ``path`` cannot hold the library, as a code, or None.
|
|
|
|
``data-dir-not-absolute``: empty or relative. ``data-dir-served``: inside
|
|
ComfyUI's output, input or temp folder, a public user folder, an extra
|
|
gallery root or a gallery source (their files are served over HTTP).
|
|
"""
|
|
if not isinstance(path, str) or not path.strip() or not os.path.isabs(path):
|
|
return "data-dir-not-absolute"
|
|
module = _folder_paths(folder_paths)
|
|
if _inside_any(path, _served_folders(module, environ, gallery_roots)):
|
|
return "data-dir-served"
|
|
if _in_public_user_folder(path, module):
|
|
return "data-dir-served"
|
|
return None
|
|
|
|
|
|
def read_marker(directory: Optional[str]) -> dict:
|
|
"""The move marker in ``directory``, or ``{}`` when absent or unreadable."""
|
|
if not directory:
|
|
return {}
|
|
try:
|
|
with open(os.path.join(directory, MARKER_NAME), "r", encoding="utf-8") as fh:
|
|
data = json.load(fh)
|
|
except (OSError, ValueError):
|
|
return {}
|
|
if not isinstance(data, dict) or data.get("state") not in (
|
|
MARKER_PENDING,
|
|
MARKER_MOVED,
|
|
):
|
|
return {}
|
|
return data
|
|
|
|
|
|
def _marker_state(directory) -> Optional[str]:
|
|
return read_marker(directory).get("state")
|
|
|
|
|
|
def switch_to(location: LibraryLocation) -> None:
|
|
"""Use ``location`` for the rest of this process (after a live move)."""
|
|
global _override
|
|
with _override_lock:
|
|
_override = location
|
|
|
|
|
|
def clear_override() -> None:
|
|
"""Forget a live move's location (tests)."""
|
|
global _override
|
|
with _override_lock:
|
|
_override = None
|
|
|
|
|
|
def current_override() -> Optional[LibraryLocation]:
|
|
with _override_lock:
|
|
return _override
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Decision:
|
|
data_dir: str
|
|
source: str
|
|
target_dir: Optional[str]
|
|
problem: Optional[str]
|
|
db_dir: Optional[str] = None # when the database is not in data_dir
|
|
target_source: Optional[str] = None
|
|
|
|
|
|
def _decide(environ, extension_root, folder_paths, gallery_roots) -> _Decision:
|
|
"""Which folder holds the library (before an explicit database path)."""
|
|
module = _folder_paths(folder_paths)
|
|
environ = os.environ if environ is None else environ
|
|
legacy_db = os.path.isfile(os.path.join(extension_root, DATABASE_NAME))
|
|
legacy = legacy_db or os.path.isfile(os.path.join(extension_root, CONFIG_NAME))
|
|
default = default_data_dir(module)
|
|
problem = None
|
|
env_dir = None
|
|
raw = environ.get(DATA_DIR_ENV, "")
|
|
if raw.strip():
|
|
problem = data_dir_problem(
|
|
raw.strip(),
|
|
folder_paths=module,
|
|
environ=environ,
|
|
gallery_roots=gallery_roots,
|
|
)
|
|
if problem is None:
|
|
env_dir = os.path.abspath(raw.strip())
|
|
target = env_dir or default
|
|
target_source = SOURCE_ENVIRONMENT if env_dir else SOURCE_DEFAULT
|
|
|
|
def decision(data_dir, source, problem=None, db_dir=None):
|
|
return _Decision(
|
|
data_dir, source, target, problem, db_dir, target_source if target else None
|
|
)
|
|
|
|
if env_dir is not None:
|
|
state = _marker_state(env_dir)
|
|
holds = os.path.isfile(os.path.join(env_dir, DATABASE_NAME))
|
|
if state != MARKER_PENDING and (holds or not legacy_db):
|
|
return decision(env_dir, SOURCE_ENVIRONMENT)
|
|
problem = "move-unfinished" if state == MARKER_PENDING else "data-dir-empty"
|
|
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
|
|
if default is None:
|
|
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
|
|
if _marker_state(default) == MARKER_MOVED and os.path.isfile(
|
|
os.path.join(default, DATABASE_NAME)
|
|
):
|
|
return decision(default, SOURCE_DEFAULT, problem)
|
|
if legacy:
|
|
if _marker_state(default) == MARKER_PENDING:
|
|
problem = problem or "move-unfinished"
|
|
if legacy_db:
|
|
return decision(extension_root, SOURCE_NODE_FOLDER, problem)
|
|
# Only a node-folder config.json: keep reading it, but a new database
|
|
# never goes into the node folder.
|
|
return decision(extension_root, SOURCE_DEFAULT, problem, default)
|
|
return decision(default, SOURCE_DEFAULT, problem)
|
|
|
|
|
|
def config_location(
|
|
environ=None, extension_root=None, folder_paths=_UNSET, gallery_roots=()
|
|
) -> Tuple[str, str]:
|
|
"""``(config.json path, data folder)``; needs no config value (no cycle)."""
|
|
override = current_override()
|
|
if override is not None:
|
|
return override.config_path, override.data_dir
|
|
environ = os.environ if environ is None else environ
|
|
root = extension_root or EXTENSION_ROOT
|
|
decision = _decide(environ, root, folder_paths, gallery_roots)
|
|
explicit = environ.get(CONFIG_PATH_ENV, "").strip()
|
|
config_path = explicit or os.path.join(decision.data_dir, CONFIG_NAME)
|
|
return config_path, decision.data_dir
|
|
|
|
|
|
def is_explicit_db_setting(value) -> bool:
|
|
"""Whether config.json's ``database.default_path`` names a chosen file."""
|
|
return isinstance(value, str) and value.strip() not in ("", DEFAULT_DB_SETTING)
|
|
|
|
|
|
def resolve_library(
|
|
default_path=None,
|
|
*,
|
|
environ: Optional[Mapping[str, str]] = None,
|
|
extension_root: Optional[str] = None,
|
|
folder_paths=_UNSET,
|
|
gallery_roots=(),
|
|
) -> LibraryLocation:
|
|
"""The library's location now (the module docstring has the order)."""
|
|
override = current_override()
|
|
if override is not None:
|
|
return override
|
|
environ = os.environ if environ is None else environ
|
|
root = extension_root or EXTENSION_ROOT
|
|
decision = _decide(environ, root, folder_paths, gallery_roots)
|
|
explicit_config = environ.get(CONFIG_PATH_ENV, "").strip()
|
|
config_path = explicit_config or os.path.join(decision.data_dir, CONFIG_NAME)
|
|
common = {
|
|
"target_source": decision.target_source,
|
|
"config_explicit": bool(explicit_config),
|
|
}
|
|
if decision.source != SOURCE_ENVIRONMENT and is_explicit_db_setting(default_path):
|
|
db_path = default_path.strip()
|
|
if not os.path.isabs(db_path):
|
|
db_path = os.path.normpath(os.path.join(root, db_path))
|
|
return LibraryLocation(
|
|
SOURCE_CONFIG,
|
|
db_path,
|
|
config_path,
|
|
decision.data_dir,
|
|
decision.target_dir,
|
|
decision.problem,
|
|
**common,
|
|
)
|
|
return LibraryLocation(
|
|
decision.source,
|
|
os.path.join(decision.db_dir or decision.data_dir, DATABASE_NAME),
|
|
config_path,
|
|
decision.data_dir,
|
|
decision.target_dir,
|
|
decision.problem,
|
|
**common,
|
|
)
|
|
|
|
|
|
def ensure_private_dir(path: str) -> str:
|
|
"""Create ``path`` (and parents) if missing; a new folder is owner-only.
|
|
|
|
An existing folder is left as it is (its permissions are the user's).
|
|
"""
|
|
if not os.path.isdir(path):
|
|
os.makedirs(path, mode=0o700, exist_ok=True)
|
|
if os.name == "posix":
|
|
try:
|
|
os.chmod(path, 0o700)
|
|
except OSError:
|
|
pass
|
|
return path
|
|
|
|
|
|
def logs_dir(**kwargs) -> str:
|
|
"""The folder for Prompt Manager's log files."""
|
|
return os.path.join(config_location(**kwargs)[1], "logs")
|
|
|
|
|
|
def cache_dir(name: str, **kwargs) -> str:
|
|
"""A re-creatable cache folder, ``<data folder>/data/<name>``."""
|
|
return os.path.join(config_location(**kwargs)[1], "data", name)
|