diff --git a/app/assets/scanner_changes.py b/app/assets/scanner_changes.py index c5350d6fa..43ee8a47a 100644 --- a/app/assets/scanner_changes.py +++ b/app/assets/scanner_changes.py @@ -21,7 +21,7 @@ from app.assets.database.queries.records import ( mark_content_missing, unset_content_missing, ) -from app.assets.helpers import to_stored_hash +from app.assets.helpers import sql_path_under_prefix, to_stored_hash from app.assets.services.path_utils import compute_loader_path, get_name_and_tags_from_asset_path from app.assets.services.snapshot_hash import snapshot_hash @@ -193,7 +193,13 @@ def drain_pending_verifications(session: Session, limit: int | None = None) -> i def live_contents_under_prefixes(session: Session, prefixes: list[str]) -> list[AssetContent]: - contents = session.scalars( - sa.select(AssetContent).where(AssetContent.is_missing.is_(False)) + if not prefixes: + return [] + return list( + session.scalars( + sa.select(AssetContent).where( + AssetContent.is_missing.is_(False), + sa.or_(*(sql_path_under_prefix(AssetContent.path, prefix) for prefix in prefixes)), + ) + ) ) - return [content for content in contents if is_path_under_prefixes(content.path, prefixes)] diff --git a/main.py b/main.py index d9f1d2544..3c6e41d35 100644 --- a/main.py +++ b/main.py @@ -457,13 +457,22 @@ def setup_database(asset_manager): init_db() asset_manager.startup() except Exception as e: - if "database is locked" in str(e) or "Could not acquire lock on database" in str(e): + if "database is locked" in str(e): logging.error( "Database is locked. Another ComfyUI process is already using this database.\n" "To resolve this, specify a separate database file for this instance:\n" " --database-url sqlite:///path/to/another.db" ) sys.exit(1) + if "Could not acquire lock on database" in str(e): + logging.error( + "Database is locked. Another ComfyUI process is already using this database.\n" + "To resolve this, specify a separate database file for this instance:\n" + " --database-url sqlite:///path/to/another.db" + ) + if args.enable_assets: + sys.exit(1) + return if args.enable_assets: logging.error( f"Failed to initialize database: {e}\n" diff --git a/tests-unit/app_test/test_db_init_locking.py b/tests-unit/app_test/test_db_init_locking.py index 3ccd13254..b9cfa8c42 100644 --- a/tests-unit/app_test/test_db_init_locking.py +++ b/tests-unit/app_test/test_db_init_locking.py @@ -102,7 +102,7 @@ def test_setup_database_routes_file_lock_to_lock_guidance(monkeypatch, caplog): def _raise_file_lock(): raise RuntimeError( - "Could not acquire lock on database '/some/path.db'. " + "Could not acquire lock on database 'x.db'. " "Another ComfyUI process may already be using it. " "Use --database-url to specify a separate database file." ) @@ -110,9 +110,46 @@ def test_setup_database_routes_file_lock_to_lock_guidance(monkeypatch, caplog): monkeypatch.setattr(main, "init_db", _raise_file_lock) monkeypatch.setattr(main.args, "enable_assets", False) + with caplog.at_level(logging.ERROR): + result = main.setup_database(None) + + assert result is None + assert "Database is locked. Another ComfyUI process is already using this database." in caplog.text + assert "Failed to initialize database." not in caplog.text + + +def test_setup_database_exits_for_file_lock_when_assets_are_enabled(monkeypatch, caplog): + monkeypatch.setattr(main, "dependencies_available", lambda: True) + + def _raise_file_lock(): + raise RuntimeError( + "Could not acquire lock on database 'x.db'. " + "Another ComfyUI process may already be using it. " + "Use --database-url to specify a separate database file." + ) + + monkeypatch.setattr(main, "init_db", _raise_file_lock) + monkeypatch.setattr(main.args, "enable_assets", True) + + with caplog.at_level(logging.ERROR), pytest.raises(SystemExit) as error: + main.setup_database(None) + + assert error.value.code == 1 + assert "Database is locked. Another ComfyUI process is already using this database." in caplog.text + assert "The --enable-assets flag requires a working database connection." not in caplog.text + + +def test_setup_database_exits_for_driver_lock_when_assets_are_disabled(monkeypatch, caplog): + monkeypatch.setattr(main, "dependencies_available", lambda: True) + + def _raise_driver_lock(): + raise Exception("sqlite3.OperationalError: database is locked") + + monkeypatch.setattr(main, "init_db", _raise_driver_lock) + monkeypatch.setattr(main.args, "enable_assets", False) + with caplog.at_level(logging.ERROR), pytest.raises(SystemExit) as error: main.setup_database(None) assert error.value.code == 1 assert "Database is locked. Another ComfyUI process is already using this database." in caplog.text - assert "Failed to initialize database." not in caplog.text diff --git a/tests-unit/assets_test/services/test_enrichment_predicate.py b/tests-unit/assets_test/services/test_enrichment_predicate.py index 8ad337d0d..a04905434 100644 --- a/tests-unit/assets_test/services/test_enrichment_predicate.py +++ b/tests-unit/assets_test/services/test_enrichment_predicate.py @@ -9,9 +9,12 @@ from unittest.mock import patch import sqlalchemy as sa from sqlalchemy.orm import Session -from app.assets.database.queries import create_content, create_record +from app.assets.database.queries import create_content, create_record, mark_content_missing from app.assets.scanner import get_unenriched_assets_for_roots -from app.assets.scanner_changes import is_path_under_prefixes +from app.assets.scanner_changes import ( + is_path_under_prefixes, + live_contents_under_prefixes, +) from .path_prefix_cases import prefix_case_paths @@ -205,3 +208,100 @@ def test_prefix_holding_metacharacters_matches_only_literal_children( assert is_path_under_prefixes(decoy_path, [root]) is False assert returned == {inside_path} + + +# --- live_contents_under_prefixes direct pins --- +# These pins exercise live_contents_under_prefixes directly, not the enrichment predicate. + + +def test_live_contents_under_prefixes_returns_empty_for_empty_prefixes( + session: Session, temp_dir: Path +) -> None: + _ = create_content(session, str(temp_dir / "seed.safetensors")) + + returned = {content.path for content in live_contents_under_prefixes(session, [])} + + assert returned == set() + + +def test_live_contents_under_prefixes_matches_directory_and_exact_path_prefixes( + session: Session, temp_dir: Path +) -> None: + content = create_content(session, str(temp_dir / "root" / "model.safetensors")) + + under_directory = {candidate.path for candidate in live_contents_under_prefixes(session, [str(temp_dir / "root")])} + exact_path = {candidate.path for candidate in live_contents_under_prefixes(session, [content.path])} + + assert under_directory == {content.path} + assert exact_path == {content.path} + + +def test_live_contents_under_prefixes_respects_sibling_boundary( + session: Session, temp_dir: Path +) -> None: + sibling_path = str(temp_dir / "a" / "bc" / "model.safetensors") + _ = create_content(session, sibling_path) + + returned = {content.path for content in live_contents_under_prefixes(session, [str(temp_dir / "a" / "b")])} + + assert returned == set() + + +def test_live_contents_under_prefixes_is_case_sensitive( + session: Session, temp_dir: Path +) -> None: + case_different_path = str(temp_dir / "data" / "TEMP" / "model.safetensors") + _ = create_content(session, case_different_path) + + returned = {content.path for content in live_contents_under_prefixes(session, [str(temp_dir / "data" / "temp")])} + + assert returned == set() + + +def test_live_contents_under_prefixes_treats_metacharacters_literally( + session: Session, temp_dir: Path +) -> None: + prefix = temp_dir / "a_b%c*d?e[f" + literal_child = create_content(session, str(prefix / "child.safetensors")) + decoy_path = str(temp_dir / "aXbYc*d?e[f" / "keep.safetensors") + _ = create_content(session, decoy_path) + + returned = {content.path for content in live_contents_under_prefixes(session, [str(prefix)])} + + assert returned == {literal_child.path} + + +def test_live_contents_under_prefixes_matches_any_prefix( + session: Session, temp_dir: Path +) -> None: + content = create_content(session, str(temp_dir / "second" / "model.safetensors")) + prefixes = [str(temp_dir / "first"), str(temp_dir / "second")] + + returned = {candidate.path for candidate in live_contents_under_prefixes(session, prefixes)} + + assert returned == {content.path} + + +def test_live_contents_under_prefixes_excludes_missing_content( + session: Session, temp_dir: Path +) -> None: + content = create_content(session, str(temp_dir / "root" / "missing.safetensors")) + mark_content_missing(session, content.id) + + returned = {candidate.path for candidate in live_contents_under_prefixes(session, [str(temp_dir / "root")])} + + assert returned == set() + + +def test_live_contents_under_prefixes_equals_python_predicate_for_corpus( + session: Session, temp_dir: Path +) -> None: + root = str(temp_dir / "root") + corpus = prefix_case_paths(root) + stored_paths = {create_content(session, path).path for path in corpus} + prefixes = [root] + + returned = {content.path for content in live_contents_under_prefixes(session, prefixes)} + expected = {path for path in stored_paths if is_path_under_prefixes(path, prefixes)} + + assert returned == expected