diff --git a/database/operations.py b/database/operations.py index aaaf2bd3..65dc2552 100644 --- a/database/operations.py +++ b/database/operations.py @@ -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,67 @@ 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_library(): + """``(database.default_path, gallery roots)`` from config, 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, [] + roots = [] + try: + try: + from ..py.config import GalleryConfig + except ImportError: + from py.config import GalleryConfig + directories = list(GalleryConfig.MONITORING_DIRECTORIES or []) + except Exception: + directories = [] + for directory in directories: + try: + roots.append(GalleryConfig.resolve_gallery_root(directory)) + except Exception: + continue + return getattr(PromptManagerConfig, "DEFAULT_DB_PATH", None), roots + + +def resolve_library_location(): + """Where the library lives now (``utils.library_location``, #249 S6b).""" + default_path, roots = _configured_library() + return library_location.resolve_library(default_path, gallery_roots=roots) + + +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 diff --git a/py/config.py b/py/config.py index 18f94f10..11bcdf28 100644 --- a/py/config.py +++ b/py/config.py @@ -36,12 +36,14 @@ from typing import Any, Dict, List, Optional, Tuple # noqa: E402 # Import logging system try: from ..utils.logging_config import get_logger + from ..utils import library_location except ImportError: import sys current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, current_dir) from utils.logging_config import get_logger + from utils import library_location # Initialize logger for config operations config_logger = get_logger("prompt_manager.config") @@ -663,6 +665,10 @@ class PromptManagerConfig: # database first (System › Backup & export › Recovery policy, #249 S6). SAFETY_BACKUP = True + # "Keep it here" (#249 S6b): the user chose to keep a node-folder library + # where it is, so the workspace stops offering to move it. + KEEP_LIBRARY_IN_NODE_FOLDER = False + @classmethod def max_worker_threads(cls) -> int: """Largest accepted worker-thread count: the detected core count.""" @@ -718,6 +724,7 @@ class PromptManagerConfig: "worker_threads": cls.WORKER_THREADS, }, "maintenance": {"safety_backup": cls.SAFETY_BACKUP}, + "library": {"keep_in_node_folder": cls.KEEP_LIBRARY_IN_NODE_FOLDER}, "gallery": GalleryConfig.get_config(), "integrations": IntegrationConfig.get_config(), "autotag": AutotagConfig.get_config(), @@ -727,14 +734,12 @@ class PromptManagerConfig: def get_config_path(cls) -> str: """Path of the persisted config.json. - Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file - lives at the repository root next to ``pyproject.toml``. + Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise it lives in + the library's data folder (``utils.library_location``): the node + folder for a library that was not moved, else ComfyUI's + ``user/__promptmanager``. """ - override = os.environ.get(CONFIG_PATH_ENV, "").strip() - if override: - return override - repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - return os.path.join(repo_root, "config.json") + return library_location.config_location()[0] @classmethod def load_from_file(cls, config_path: Optional[str] = None): @@ -859,6 +864,12 @@ class PromptManagerConfig: ): cls.SAFETY_BACKUP = maintenance["safety_backup"] + library = new_config.get("library", {}) + if isinstance(library, dict) and isinstance( + library.get("keep_in_node_folder"), bool + ): + cls.KEEP_LIBRARY_IN_NODE_FOLDER = library["keep_in_node_folder"] + # Update gallery config if "gallery" in new_config: GalleryConfig.update_config(new_config["gallery"]) diff --git a/py/lora_utils.py b/py/lora_utils.py index 454c182a..83091dbe 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -567,10 +567,21 @@ def download_civitai_images( def get_lora_image_cache_dir() -> Path: - """Get the directory used to cache downloaded LoRA example images.""" - # Store in the extension's own directory + """Get the directory used to cache downloaded LoRA example images. + + ``/data/lora_images`` (``utils.library_location``): the node + folder for a library that was not moved, else ComfyUI's + ``user/__promptmanager``. + """ + try: + from ..utils import library_location + except ImportError: + from utils import library_location + ext_root = Path(__file__).resolve().parent.parent - cache = ext_root / "data" / "lora_images" + cache = Path( + library_location.cache_dir("lora_images", extension_root=str(ext_root)) + ) cache.mkdir(parents=True, exist_ok=True) return cache diff --git a/tests/test_library_location.py b/tests/test_library_location.py new file mode 100644 index 00000000..44202bbd --- /dev/null +++ b/tests/test_library_location.py @@ -0,0 +1,326 @@ +"""Where the library lives (#249 S6b): the one resolver, ``utils/library_location``. + +Every case builds its own extension root (the node folder) and ComfyUI user +folder in a temp directory and hands the resolver a fake ``folder_paths``, so +no test reads the real environment, the repository's own files or ComfyUI. +""" + +import json +import os +import sys +import tempfile +import types +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils import library_location as ll # noqa: E402 + + +def _touch(path, data=b"x"): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as handle: + handle.write(data) + + +def _folder_paths(base, *, helper=True): + """A ComfyUI ``folder_paths`` stand-in rooted at ``base``.""" + module = types.SimpleNamespace(base_path=base) + for name in ("output", "input", "temp", "user"): + directory = os.path.join(base, name) + os.makedirs(directory, exist_ok=True) + setattr(module, f"get_{name}_directory", lambda d=directory: d) + if helper: + module.get_system_user_directory = lambda name="system": os.path.join( + base, "user", "__" + name + ) + return module + + +class _Case(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.base = os.path.realpath(self._tmp.name) + self.comfy = os.path.join(self.base, "ComfyUI") + self.ext = os.path.join(self.comfy, "custom_nodes", "ComfyUI_PromptManager") + os.makedirs(self.ext) + self.fp = _folder_paths(self.comfy) + self.default_dir = os.path.join(self.comfy, "user", "__promptmanager") + ll.clear_override() + self.addCleanup(ll.clear_override) + + def resolve(self, default_path=None, environ=None, **kwargs): + kwargs.setdefault("folder_paths", self.fp) + return ll.resolve_library( + default_path, environ=environ or {}, extension_root=self.ext, **kwargs + ) + + def config_location(self, environ=None, **kwargs): + kwargs.setdefault("folder_paths", self.fp) + return ll.config_location( + environ=environ or {}, extension_root=self.ext, **kwargs + ) + + +class TestDefaultDataFolder(_Case): + def test_uses_comfyui_system_user_helper(self): + self.assertEqual(ll.default_data_dir(self.fp), self.default_dir) + + def test_falls_back_to_user_directory_without_helper(self): + fp = _folder_paths(self.comfy, helper=False) + self.assertEqual(ll.default_data_dir(fp), self.default_dir) + + def test_none_outside_comfyui(self): + self.assertIsNone(ll.default_data_dir(None)) + self.assertIsNone(ll.default_data_dir(types.SimpleNamespace())) + + +class TestResolutionOrder(_Case): + def test_new_install_uses_default_folder_and_creates_nothing_in_node_folder(self): + location = self.resolve() + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db")) + self.assertEqual( + location.config_path, os.path.join(self.default_dir, "config.json") + ) + self.assertFalse(location.movable) + self.assertEqual(os.listdir(self.ext), []) + + def test_legacy_node_folder_library_stays_in_use(self): + _touch(os.path.join(self.ext, "prompts.db")) + location = self.resolve() + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db")) + self.assertEqual(location.config_path, os.path.join(self.ext, "config.json")) + self.assertEqual(location.data_dir, self.ext) + self.assertTrue(location.movable) + self.assertEqual(location.target_dir, self.default_dir) + + def test_legacy_wins_over_an_unmarked_default_folder_library(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + + def test_explicit_default_path_is_honoured_as_today(self): + _touch(os.path.join(self.ext, "prompts.db")) + elsewhere = os.path.join(self.base, "elsewhere", "pm.db") + location = self.resolve(elsewhere) + self.assertEqual(location.source, ll.SOURCE_CONFIG) + self.assertEqual(location.db_path, elsewhere) + self.assertFalse(location.movable) + relative = self.resolve("data/my.db") + self.assertEqual(relative.db_path, os.path.join(self.ext, "data", "my.db")) + self.assertEqual(relative.source, ll.SOURCE_CONFIG) + + def test_default_setting_value_is_not_explicit(self): + for value in (None, "", "prompts.db"): + self.assertEqual(self.resolve(value).source, ll.SOURCE_DEFAULT, value) + + def test_environment_folder_wins(self): + data = os.path.join(self.base, "pm-data") + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_ENVIRONMENT) + self.assertEqual(location.db_path, os.path.join(data, "prompts.db")) + self.assertEqual(location.config_path, os.path.join(data, "config.json")) + + def test_environment_folder_with_a_library_wins_over_legacy(self): + data = os.path.join(self.base, "pm-data") + _touch(os.path.join(data, "prompts.db")) + _touch(os.path.join(self.ext, "prompts.db")) + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_ENVIRONMENT) + + def test_empty_environment_folder_never_replaces_a_legacy_library(self): + data = os.path.join(self.base, "pm-data") + _touch(os.path.join(self.ext, "prompts.db")) + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db")) + self.assertEqual(location.problem, "data-dir-empty") + self.assertEqual(location.target_dir, data) + self.assertTrue(location.movable) + + def test_config_path_environment_is_honoured(self): + custom = os.path.join(self.base, "cfg", "custom.json") + location = self.resolve(environ={ll.CONFIG_PATH_ENV: custom}) + self.assertEqual(location.config_path, custom) + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + + def test_without_comfyui_the_node_folder_is_used(self): + location = self.resolve(folder_paths=None) + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db")) + self.assertFalse(location.movable) + + def test_override_after_a_live_move_wins(self): + _touch(os.path.join(self.ext, "prompts.db")) + moved = ll.LibraryLocation( + source=ll.SOURCE_DEFAULT, + db_path=os.path.join(self.default_dir, "prompts.db"), + config_path=os.path.join(self.default_dir, "config.json"), + data_dir=self.default_dir, + target_dir=self.default_dir, + ) + ll.switch_to(moved) + self.assertEqual(self.resolve(), moved) + self.assertEqual(self.config_location(), (moved.config_path, moved.data_dir)) + ll.clear_override() + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + + +class TestMarker(_Case): + def _marker(self, directory, state): + os.makedirs(directory, exist_ok=True) + with open(os.path.join(directory, ll.MARKER_NAME), "w") as handle: + json.dump({"version": 1, "state": state}, handle) + + def test_moved_marker_wins_while_the_old_file_was_not_renamed_yet(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self._marker(self.default_dir, "moved") + location = self.resolve() + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db")) + self.assertEqual( + location.config_path, os.path.join(self.default_dir, "config.json") + ) + + def test_pending_marker_keeps_the_old_library(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self._marker(self.default_dir, "pending") + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + data = os.path.join(self.base, "pm-data") + _touch(os.path.join(data, "prompts.db")) + self._marker(data, "pending") + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.problem, "move-unfinished") + + def test_unreadable_marker_is_not_trusted(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + _touch(os.path.join(self.default_dir, ll.MARKER_NAME), b"{not json") + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(ll.read_marker(self.default_dir), {}) + + def test_renamed_old_file_is_never_picked(self): + _touch(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self.assertEqual(self.resolve().source, ll.SOURCE_DEFAULT) + + +class TestConfigLocation(_Case): + def test_legacy_config_without_a_database_keeps_being_read(self): + _touch(os.path.join(self.ext, "config.json"), b"{}") + path, data_dir = self.config_location() + self.assertEqual(path, os.path.join(self.ext, "config.json")) + self.assertEqual(data_dir, self.ext) + + def test_legacy_config_alone_never_puts_a_new_database_in_the_node_folder(self): + _touch(os.path.join(self.ext, "config.json"), b"{}") + location = self.resolve() + self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db")) + self.assertEqual(location.config_path, os.path.join(self.ext, "config.json")) + self.assertFalse(location.movable) + + def test_new_install_reads_config_from_the_data_folder(self): + path, data_dir = self.config_location() + self.assertEqual(path, os.path.join(self.default_dir, "config.json")) + self.assertEqual(data_dir, self.default_dir) + + def test_logs_and_cache_follow_the_data_folder(self): + self.assertEqual( + ll.logs_dir(environ={}, extension_root=self.ext, folder_paths=self.fp), + os.path.join(self.default_dir, "logs"), + ) + _touch(os.path.join(self.ext, "prompts.db")) + self.assertEqual( + ll.logs_dir(environ={}, extension_root=self.ext, folder_paths=self.fp), + os.path.join(self.ext, "logs"), + ) + self.assertEqual( + ll.cache_dir( + "lora_images", environ={}, extension_root=self.ext, folder_paths=self.fp + ), + os.path.join(self.ext, "data", "lora_images"), + ) + + +class TestDataDirValidation(_Case): + def problem(self, path, **kwargs): + kwargs.setdefault("folder_paths", self.fp) + return ll.data_dir_problem(path, **kwargs) + + def test_valid_folder(self): + self.assertIsNone(self.problem(os.path.join(self.base, "pm-data"))) + self.assertIsNone(self.problem(self.default_dir)) + + def test_relative_and_empty_paths_are_refused(self): + self.assertEqual(self.problem("pm-data"), "data-dir-not-absolute") + self.assertEqual(self.problem(""), "data-dir-not-absolute") + + def test_served_folders_are_refused(self): + for name in ("output", "input", "temp"): + inside = os.path.join(self.comfy, name, "pm") + self.assertEqual(self.problem(inside), "data-dir-served", name) + self.assertEqual( + self.problem(os.path.join(self.comfy, "output")), "data-dir-served" + ) + + def test_public_user_folder_is_refused_but_system_users_are_fine(self): + public = os.path.join(self.comfy, "user", "default", "pm") + self.assertEqual(self.problem(public), "data-dir-served") + self.assertIsNone(self.problem(os.path.join(self.comfy, "user", "__pm"))) + + def test_gallery_source_is_refused(self): + source = os.path.join(self.base, "gallery") + os.makedirs(source) + self.assertEqual( + self.problem(os.path.join(source, "pm"), gallery_roots=[source]), + "data-dir-served", + ) + self.assertEqual( + self.problem( + os.path.join(source, "pm"), + environ={ll.EXTRA_GALLERY_ROOTS_ENV: source}, + ), + "data-dir-served", + ) + + def test_symlink_into_output_is_refused(self): + link = os.path.join(self.base, "link-to-output") + try: + os.symlink(os.path.join(self.comfy, "output"), link) + except (OSError, NotImplementedError): + self.skipTest("symlinks are not available") + self.assertEqual(self.problem(os.path.join(link, "pm")), "data-dir-served") + + def test_case_variant_spelling_of_output_is_refused_where_the_volume_folds_case( + self, + ): + upper = os.path.join(self.comfy, "OUTPUT") + if not os.path.isdir(upper): + self.skipTest("this volume is case-sensitive") + self.assertEqual(self.problem(os.path.join(upper, "pm")), "data-dir-served") + + def test_invalid_environment_folder_is_ignored_with_a_problem(self): + inside = os.path.join(self.comfy, "output", "pm") + location = self.resolve(environ={ll.DATA_DIR_ENV: inside}) + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + self.assertEqual(location.problem, "data-dir-served") + self.assertFalse(os.path.exists(inside)) + relative = self.resolve(environ={ll.DATA_DIR_ENV: "relative/dir"}) + self.assertEqual(relative.problem, "data-dir-not-absolute") + + +class TestSameFolder(_Case): + def test_same_folder_follows_links_and_case(self): + self.assertTrue(ll.same_folder(self.ext, self.ext + os.sep)) + self.assertFalse(ll.same_folder(self.ext, self.default_dir)) + self.assertFalse(ll.same_folder(self.ext, os.path.join(self.base, "missing"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_library_location_wiring.py b/tests/test_library_location_wiring.py new file mode 100644 index 00000000..d2fa31af --- /dev/null +++ b/tests/test_library_location_wiring.py @@ -0,0 +1,169 @@ +"""Every data path asks the one resolver (#249 S6b). + +The database, the diagnostics database, config.json, the logs and the LoRA +image cache all follow ``utils.library_location``: a folder named by +``PROMPT_MANAGER_DATA_DIR`` (and, after a live move, the switched location). +""" + +import os +import pathlib +import sys +import tempfile +import types +import unittest +import unittest.mock as mock +from unittest.mock import MagicMock + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules.setdefault("server", _mock_server) + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils import library_location as ll # noqa: E402 + + +class _EnvCase(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + base = os.path.realpath(self._tmp.name) + self.data = os.path.join(base, "pm-data") + patcher = mock.patch.dict( + os.environ, + {ll.DATA_DIR_ENV: self.data, ll.CONFIG_PATH_ENV: ""}, + ) + patcher.start() + self.addCleanup(patcher.stop) + # This test's own ComfyUI folders (the shared stub serves the system + # temp folder as "output", which would rightly refuse self.data). + comfy = os.path.join(base, "ComfyUI") + fp = types.ModuleType("folder_paths") + for name in ("output", "input", "temp", "user"): + setattr(fp, f"get_{name}_directory", lambda n=name: os.path.join(comfy, n)) + modules = mock.patch.dict(sys.modules, {"folder_paths": fp}) + modules.start() + self.addCleanup(modules.stop) + ll.clear_override() + self.addCleanup(ll.clear_override) + + +class TestDatabasePath(_EnvCase): + def test_database_follows_the_data_folder(self): + from database.operations import _resolve_db_path + + with mock.patch.dict(sys.modules, {"py.config": None}): + path = _resolve_db_path(None) + self.assertEqual(path, os.path.join(self.data, "prompts.db")) + self.assertTrue(os.path.isdir(self.data)) + if os.name == "posix": + self.assertEqual(os.stat(self.data).st_mode & 0o777, 0o700) + + def test_diagnostics_use_the_same_resolver(self): + from database.operations import _resolve_db_path + from utils.diagnostics import GalleryDiagnostics + + with mock.patch.dict(sys.modules, {"py.config": None}): + self.assertEqual(GalleryDiagnostics._resolve_db_path(), _resolve_db_path()) + + def test_switched_location_wins(self): + from database.operations import _resolve_db_path + + moved = ll.LibraryLocation( + ll.SOURCE_DEFAULT, + os.path.join(self.data, "moved", "prompts.db"), + os.path.join(self.data, "moved", "config.json"), + os.path.join(self.data, "moved"), + ) + ll.switch_to(moved) + self.assertEqual(_resolve_db_path(None), moved.db_path) + + +class TestConfigPath(_EnvCase): + def test_config_json_follows_the_data_folder(self): + from py.config import PromptManagerConfig + + self.assertEqual( + PromptManagerConfig.get_config_path(), + os.path.join(self.data, "config.json"), + ) + + def test_explicit_config_path_still_wins(self): + from py.config import PromptManagerConfig + + custom = os.path.join(self.data, "elsewhere.json") + with mock.patch.dict(os.environ, {ll.CONFIG_PATH_ENV: custom}): + self.assertEqual(PromptManagerConfig.get_config_path(), custom) + + def test_keep_choice_round_trips(self): + from py.config import PromptManagerConfig + + original = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER + self.addCleanup( + setattr, PromptManagerConfig, "KEEP_LIBRARY_IN_NODE_FOLDER", original + ) + PromptManagerConfig.update_config({"library": {"keep_in_node_folder": True}}) + self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + self.assertEqual( + PromptManagerConfig.get_config()["library"], {"keep_in_node_folder": True} + ) + PromptManagerConfig.update_config({"library": {"keep_in_node_folder": "yes"}}) + self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + + +class TestLogsAndCache(_EnvCase): + def test_logs_go_to_the_data_folder(self): + from utils import logging_config + from utils.logging_config import PromptManagerLogger + + manager = object.__new__(PromptManagerLogger) + self.addCleanup(self._close, manager) + manager.__init__() + self.assertEqual(manager.log_dir, pathlib.Path(self.data, "logs")) + self.assertTrue(manager.log_dir.is_dir()) + self.assertIs(logging_config.library_location, ll) + + @staticmethod + def _close(manager): + import logging + + for handler in logging.getLogger("prompt_manager").handlers[:]: + if getattr(handler, "_prompt_manager_owned", False): + handler.close() + try: + from utils.logging_config import get_logger_manager + + get_logger_manager()._setup_loggers() + except Exception: + pass + + def test_lora_image_cache_follows_the_data_folder(self): + from py import lora_utils + + cache = lora_utils.get_lora_image_cache_dir() + self.assertEqual(cache, pathlib.Path(self.data, "data", "lora_images")) + + +class TestNodeFolderUnchangedForLegacyInstalls(unittest.TestCase): + """A node folder that holds the library keeps every path where it was.""" + + def test_legacy_paths(self): + with tempfile.TemporaryDirectory() as tmp: + ext = os.path.join(tmp, "ext") + os.makedirs(ext) + open(os.path.join(ext, "prompts.db"), "wb").close() + fp = types.SimpleNamespace( + get_user_directory=lambda: os.path.join(tmp, "u") + ) + location = ll.resolve_library( + None, environ={}, extension_root=ext, folder_paths=fp + ) + self.assertEqual(location.db_path, os.path.join(ext, "prompts.db")) + self.assertEqual( + ll.logs_dir(environ={}, extension_root=ext, folder_paths=fp), + os.path.join(ext, "logs"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 094d81c1..7d1aa769 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -69,8 +69,8 @@ class GalleryDiagnostics: Args: db_path: Path to the SQLite database file to diagnose. - If None, resolves from config.json or defaults to - prompts.db in the extension root. + If None, the library's database (see + ``utils.library_location``). """ if db_path is None: db_path = self._resolve_db_path() @@ -79,18 +79,12 @@ class GalleryDiagnostics: @staticmethod def _resolve_db_path() -> str: - """Resolve database path from config, matching PromptDatabase logic.""" - extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + """The library's database: the same resolver as PromptDatabase.""" try: - 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 = os.path.normpath(os.path.join(extension_root, db_path)) - return db_path + from ..database.operations import resolve_library_location + except ImportError: + from database.operations import resolve_library_location + return resolve_library_location().db_path def run_full_diagnostic(self) -> Dict[str, Any]: """Run a complete diagnostic check. diff --git a/utils/library_location.py b/utils/library_location.py new file mode 100644 index 00000000..f5b1797b --- /dev/null +++ b/utils/library_location.py @@ -0,0 +1,418 @@ +"""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); + ``problem`` is a code for something the user should know (an invalid or + empty ``PROMPT_MANAGER_DATA_DIR``, a move that did not finish). + """ + + source: str + db_path: str + config_path: str + data_dir: str + target_dir: Optional[str] = None + problem: Optional[str] = None + + @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 + + +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 + 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, env_dir, None) + problem = "move-unfinished" if state == MARKER_PENDING else "data-dir-empty" + return _Decision(extension_root, SOURCE_NODE_FOLDER, env_dir, problem) + if default is None: + return _Decision(extension_root, SOURCE_NODE_FOLDER, None, problem) + if _marker_state(default) == MARKER_MOVED and os.path.isfile( + os.path.join(default, DATABASE_NAME) + ): + return _Decision(default, SOURCE_DEFAULT, target, problem) + if legacy: + if _marker_state(default) == MARKER_PENDING: + problem = problem or "move-unfinished" + if legacy_db: + return _Decision(extension_root, SOURCE_NODE_FOLDER, target, 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, target, problem, default) + return _Decision(default, SOURCE_DEFAULT, target, 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) + config_path = environ.get(CONFIG_PATH_ENV, "").strip() or os.path.join( + decision.data_dir, CONFIG_NAME + ) + 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, + ) + 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, + ) + + +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/``.""" + return os.path.join(config_location(**kwargs)[1], "data", name) diff --git a/utils/logging_config.py b/utils/logging_config.py index d9a62866..bb147a2e 100644 --- a/utils/logging_config.py +++ b/utils/logging_config.py @@ -18,6 +18,8 @@ from datetime import datetime from typing import Dict, Any, Optional, List from pathlib import Path +from . import library_location + class _OwnedRotatingFileHandler(logging.handlers.RotatingFileHandler): """Keep a retired owned sink closed when an in-flight record resumes.""" @@ -65,8 +67,12 @@ class PromptManagerLogger: return self._initialized = True - self.log_dir = Path(__file__).parent.parent / "logs" - self.log_dir.mkdir(exist_ok=True) + # The library's data folder (#249 S6b): the node folder for a library + # that was not moved, else ComfyUI's user/__promptmanager. + self.log_dir = Path( + library_location.logs_dir(extension_root=str(Path(__file__).parent.parent)) + ) + library_location.ensure_private_dir(str(self.log_dir)) # Configuration self.config = {