AI/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-10-24 02:49:04 +03:00
Code Activity

fix(assets): write the prune and offline marking in short batches so saves aren't locked out (#16696)

Browse Source 
* fix(assets): batch the prune's and the offline marking's writes

The startup prune, POST /api/assets/prune and the fast scan's marking step
each held the SQLite write lock for their whole loop, so foreground output
registration failed with "database is locked" during a large one. They now
write in short batches, wait while a prompt runs between batches, and the
prune endpoint runs off the event loop.

* fix(assets): start the queued scan after a standalone prune, and recheck listing rows after a pause

A prompt that ends while POST /api/assets/prune runs queues its output rescan;
the prune now starts it when it finishes, as a scan does. The output-listing
rescan takes its batch gate before reading the live rows, so a pause during the
walk makes the marking re-stat what it retires. A cancel that arrives after the
last batch no longer reports a finished prune as cancelled.

* refactor(assets): drop the pause rechecks and the cancellable standalone prune

Batching the writes is what keeps the lock short; the layers on top of it
guarded edge cases that heal on the next scan. Batches now just commit, sleep
about as long as they held the lock, and between batches honour the scan's
pause/cancel checkpoint. The standalone prune is batched but not pausable, so
it needs no cancel status or pending-scan handling, and the API contract is
unchanged apart from running off the event loop.

* fix(assets): start the scan queued behind a standalone prune; skip the last batch's yield

POST /api/assets/prune now runs off the event loop, so a prompt can finish
while it runs and queue its output rescan; the prune starts it when it ends,
as a scan does. The batch loop checks for a stop before every batch and no
longer sleeps after the last one.

* test(assets): compare the set-mark paths in their stored, absolute form

create_content stores os.path.abspath(path), which carries a drive letter on
Windows, so the expected list must be built the same way.

* fix(assets): a seed request during an API prune waits for it instead of 409

The prune now runs off the event loop, so POST /api/assets/seed can arrive
while it holds the seeder; start() fails and the route answered 409, which a
client reads as "a scan is already coming". A prune emits no scan events, so
the refresh was lost. The route now waits the prune out and starts the scan,
as it effectively did when the prune blocked the loop.

* fix(assets): a cancel or shutdown stops a standalone prune between batches

The API prune runs on a worker thread that interpreter exit joins, so a
shutdown that only flagged it left Ctrl-C waiting for the whole prune. It now
stops at the next batch once cancelled, and shutdown waits for that. A seed
request also retries start() once after any failure, covering a prune that
ends between the failed start and the check.

* fix(assets): report a cancelled API prune as cancelled, not completed

A cancel now stops a standalone prune between batches, so its response can
carry a partial count; say so with status "cancelled" rather than presenting
it as a finished prune.

* fix(assets): a cancelled standalone prune leaves a queued scan queued

Shutdown cancels the prune; starting the scan a prompt had queued from the
prune's finalizer would run it on into teardown after shutdown returned. It
now stays queued for the next scan's finalizer.

* test(assets): assert the cancelled prune's outcome in the test thread

pytest.raises inside the worker thread only produced a warning when the
exception was missing, so the test could not fail on it.

* fix(assets): wait for a prune on the loop, and close shutdown gaps around it

A seed request during an API prune now polls on the event loop instead of
holding an executor thread for the prune's length, and retries while a prune
holds the seeder. Shutdown marks the seeder so a prune that has not started
yet does not, both of its waits share one deadline, and the prune's idle flag
is set even if its cleanup raises.
This commit is contained in:
Simon Pinfold
2026-10-03 01:20:37 -07:00
committed by GitHub
parent 42664106ef
commit e9027f2b30
9 changed files with 979 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
app/assets/api/routes.py
View File

@@ -43,7 +43,7 @@ from app.assets.database.queries.records import (
get_preview_file_paths_by_ids,
list_records_page,
)
from app.assets.seeder import ScanInProgressError, asset_seeder
from app.assets.seeder import PruneCancelledError, ScanInProgressError, asset_seeder
from app.assets.services import (
DependencyMissingError,
HashMismatchError,
@@ -1030,6 +1030,9 @@ async def get_tags_refine(request: web.Request) -> web.Response:
return web.json_response(payload.model_dump(mode="json", exclude_none=True), status=200)
_PRUNE_POLL_SECONDS = 0.25
@ROUTES.post("/api/assets/seed")
@_require_assets_feature_enabled
async def seed_assets(request: web.Request) -> web.Response:
@@ -1059,6 +1062,18 @@ async def seed_assets(request: web.Request) -> web.Response:
started = asset_seeder.start(
roots=valid_roots, compute_hashes=mode.hashing_enabled()
)
# A prune from POST /api/assets/prune is not a scan and emits no scan events, so
# wait it out rather than answer 409, which a client takes as "a scan is coming".
# Polled on the loop: no executor thread is held for the prune's length. The
# retry also covers a prune that ended, or another that began, after start() failed.
for _ in range(2):
if started:
break
while asset_seeder.standalone_prune_running():
await asyncio.sleep(_PRUNE_POLL_SECONDS)
started = asset_seeder.start(
roots=valid_roots, compute_hashes=mode.hashing_enabled()
)
if not started:
return web.json_response({"status": "already_running"}, status=409)
@@ -1125,17 +1140,21 @@ async def mark_missing_assets(request: web.Request) -> web.Response:
Returns:
200 OK with count of marked assets
200 OK with status "cancelled" and the count marked before a cancel stopped it
409 Conflict if a scan is currently running
500 Internal Server Error with PRUNE_FAILED if the marking failed, so a
prune that did not run is never reported as a completed one
"""
try:
marked = asset_seeder.mark_missing_outside_prefixes()
# Off the event loop: a large prune takes seconds to minutes.
marked = await asyncio.to_thread(asset_seeder.mark_missing_outside_prefixes)
except ScanInProgressError:
return web.json_response(
{"status": "scan_running", "marked": 0},
status=409,
)
except PruneCancelledError as cancelled:
return web.json_response({"status": "cancelled", "marked": cancelled.marked}, status=200)
if marked is None:
return _build_error_response(
500,

2
app/assets/database/queries/__init__.py
View File

@@ -16,6 +16,7 @@ from app.assets.database.queries.records import (
is_live_path_conflict,
list_records_page,
mark_content_missing,
mark_contents_missing,
rename_record,
unset_content_missing,
update_record_access_time,
@@ -31,6 +32,7 @@ __all__ = [
"is_live_path_conflict",
"list_records_page",
"mark_content_missing",
"mark_contents_missing",
"rename_record",
"unset_content_missing",
"update_record_access_time",

27
app/assets/database/queries/records.py
View File

@@ -325,6 +325,33 @@ def mark_content_missing(session: Session, content_id: str) -> None:
session.flush()
def mark_contents_missing(session: Session, content_ids: Sequence[str]) -> list[str]:
"""mark_content_missing for many rows in a few statements; returns the ids it marked.
A row that is gone or already missing is skipped."""
if not content_ids:
return []
# "= 0", not "IS 0": SQLite only uses the partial live-path index for "= 0". This
# lookup is by primary key either way; the form matches the other live-row lookups.
live = list(session.scalars(sa.select(AssetContent.id).where(AssetContent.id.in_(content_ids), AssetContent.is_missing == sa.false())))
if not live:
return []
session.execute(sa.update(AssetContent).where(AssetContent.id.in_(live)).values(is_missing=True))
ensure_tag(session, "missing")
unlinked = sa.select(Asset.id, sa.literal("missing"), sa.literal("automatic"), sa.literal(get_utc_now(), sa.DateTime())).where(
Asset.content_id.in_(live),
~sa.exists().where(AssetTag.asset_id == Asset.id, AssetTag.tag_name == "missing"),
)
try:
with session.begin_nested():
session.execute(sa.insert(AssetTag).from_select(["asset_id", "tag_name", "origin", "added_at"], unlinked))
except IntegrityError:
# A concurrent writer linked one of them first; settle each link the race-safe way.
for record_id in session.scalars(sa.select(Asset.id).where(Asset.content_id.in_(live))):
ensure_tag_link(session, asset_id=record_id, tag_name="missing", origin="automatic")
session.flush()
return live
def unset_content_missing(session: Session, content_id: str) -> None:
content = session.get(AssetContent, content_id)
if content is None:

138
app/assets/scanner.py
View File

@@ -11,6 +11,7 @@ so the rows it never reached are selected again when the scan resumes.
import enum
import logging
import os
import time
from dataclasses import dataclass
from itertools import islice
from pathlib import Path
@@ -25,7 +26,7 @@ from app.assets.event_log import emit, error_kind, error_type
from app.assets.database.queries import (
create_content_reporting_insert,
is_live_path_conflict,
mark_content_missing,
mark_contents_missing,
create_record,
)
from app.assets.database.models import Asset, AssetContent
@@ -134,6 +135,46 @@ class _ReferenceObservation(NamedTuple):
stat_result: os.stat_result | None
# Before each write batch: blocks while the scan is paused, returns True to stop.
ShouldStop = Callable[[], bool]
def _never_stop() -> bool:
return False
# The marking steps write in batches of WRITE_BATCH_ROWS, each its own transaction, so
# the write lock is never held for a whole pass and a foreground write (an output being
# registered) waits at most one batch. After each commit the thread sleeps about as
# long as it held the lock: a writer in SQLite's busy handler polls with growing sleeps
# (up to 100 ms), so taking the lock straight back would keep winning it.
WRITE_BATCH_ROWS = 256
WRITE_YIELD_MIN_SECONDS = 0.02
WRITE_YIELD_MAX_SECONDS = 0.1
def _write_in_batches(
items: list,
write_batch: Callable[[Session, list], int],
should_stop: ShouldStop,
committed: list[int],
) -> None:
"""Apply ``write_batch`` to ``items`` a batch per transaction, appending the count
each batch reports to ``committed`` once it commits. A failure leaves the batches
before it committed, and ``committed`` says how much they wrote."""
for start in range(0, len(items), WRITE_BATCH_ROWS):
if should_stop():
break
opened = time.perf_counter()
with create_write_session() as session:
count = write_batch(session, items[start : start + WRITE_BATCH_ROWS])
session.commit()
committed.append(count)
held = time.perf_counter() - opened
if start + WRITE_BATCH_ROWS < len(items):
time.sleep(min(max(held, WRITE_YIELD_MIN_SECONDS), WRITE_YIELD_MAX_SECONDS))
def _log_scan_error(phase: str, error: OSError) -> None:
error_type = (
"permission_denied" if isinstance(error, PermissionError) else "os_error"
@@ -241,7 +282,9 @@ def apply_reference_observations(
session: Session, observations: list[_ReferenceObservation]
) -> int:
"""Apply the observations; returns how many rows were marked missing."""
marked = 0
# One query loads the rows, so the session.get calls below never go to the database.
list(session.scalars(sa.select(AssetContent).where(AssetContent.id.in_([o.content_id for o in observations]))))
gone: list[str] = []
for observation in observations:
content = session.get(AssetContent, observation.content_id)
# Skip a row another writer changed since it was observed; the next scan sees it afresh.
@@ -253,8 +296,7 @@ def apply_reference_observations(
):
continue
if observation.stat_result is None:
mark_content_missing(session, content.id)
marked += 1
gone.append(content.id)
continue
detect_content_change(
session,
@@ -262,35 +304,38 @@ def apply_reference_observations(
observation.stat_result,
hashing_is_enabled=mode.hashing_enabled(),
)
return marked
return len(mark_contents_missing(session, gone))
def _sync_prefixes_in_write_txn(
prefixes: list[str], progress: _ScanProgress | None
) -> tuple[set[str], int]:
"""Returns the surviving paths and how many rows were marked missing."""
def _sync_prefixes(
prefixes: list[str],
progress: _ScanProgress | None,
should_stop: ShouldStop,
marked: list[int],
) -> set[str]:
"""Returns the surviving paths; ``marked`` gets each committed batch's count of rows
marked missing, which a caller still has if a later batch raises."""
with create_session() as session:
observations, survivors = observe_references_on_filesystem(
session, prefixes, progress
)
marked = 0
if observations:
with create_write_session() as session:
marked = apply_reference_observations(session, observations)
session.commit()
return survivors, marked
_write_in_batches(observations, apply_reference_observations, should_stop, marked)
return survivors
def sync_root_safely(
root: RootType, progress: _ScanProgress | None = None
root: RootType,
progress: _ScanProgress | None = None,
should_stop: ShouldStop = _never_stop,
) -> set[str]:
"""Sync a single root's references with the filesystem.
Returns survivors (existing paths) or empty set on failure.
"""
marked: list[int] = []
try:
survivors, marked = _sync_prefixes_in_write_txn(
get_scan_prefixes_for_root(root), progress
survivors = _sync_prefixes(
get_scan_prefixes_for_root(root), progress, should_stop, marked
)
except Exception as exc:
logging.exception("fast DB scan failed for %s: %s", root, exc)
@@ -300,18 +345,19 @@ def sync_root_safely(
error_type=error_type(exc),
error_kind=error_kind(exc),
)
return set()
survivors = set()
if progress is not None:
progress.missing_marked += marked
progress.missing_marked += sum(marked)
return survivors
def sync_temp_references_safely(
progress: _ScanProgress | None = None,
should_stop: ShouldStop = _never_stop,
) -> None:
"""Retire temp references whose file is gone; temp is never scanned, so nothing else stats them."""
try:
_sync_prefixes_in_write_txn(get_temp_prefixes(), progress)
_sync_prefixes(get_temp_prefixes(), progress, should_stop, [])
except Exception as exc:
logging.exception("temp reference sync failed: %s", exc)
emit(
@@ -322,19 +368,27 @@ def sync_temp_references_safely(
)
def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int | None:
def mark_missing_outside_prefixes_safely(
prefixes: list[str], should_stop: ShouldStop = _never_stop
) -> int | None:
"""Mark references as missing when outside the given prefixes.
This is a non-destructive soft-delete. Returns the count marked, or None when
the operation fails.
the operation fails; batches committed before a failure stay committed.
"""
marked: list[int] = []
try:
with create_session() as sess:
count = mark_contents_missing_outside_prefixes(sess, prefixes)
sess.commit()
return count
content_ids = content_ids_outside_prefixes(sess, prefixes)
_write_in_batches(
content_ids,
lambda session, batch: len(mark_contents_missing(session, batch)),
should_stop,
marked,
)
return sum(marked)
except Exception as exc:
logging.exception("marking missing assets failed: %s", exc)
logging.exception("marking missing assets failed after marking %d: %s", sum(marked), exc)
emit(
"scanner.mark_missing_failed",
error_type=error_type(exc),
@@ -343,19 +397,16 @@ def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int | None:
return None
def mark_contents_missing_outside_prefixes(
session: Session, prefixes: list[str]
) -> int:
contents = session.scalars(
sa.select(AssetContent)
.where(AssetContent.is_missing.is_(False))
def content_ids_outside_prefixes(session: Session, prefixes: list[str]) -> list[str]:
"""The live rows outside every prefix. A read: the marking re-checks each is still
live inside its own write transaction."""
is_owned = path_prefix_matcher(prefixes)
rows = session.execute(
sa.select(AssetContent.id, AssetContent.path)
.where(AssetContent.is_missing == sa.false())
.execution_options(yield_per=500)
)
is_owned = path_prefix_matcher(prefixes)
missing = [content for content in contents if not is_owned(content.path)]
for content in missing:
mark_content_missing(session, content.id)
return len(missing)
return [content_id for content_id, path in rows if not is_owned(path)]
def collect_paths_for_roots(
@@ -509,15 +560,13 @@ def mark_unlisted_references_missing_safely(
root: RootType,
observations: list[_ReferenceObservation],
progress: _ScanProgress | None = None,
should_stop: ShouldStop = _never_stop,
) -> None:
"""Retire rows whose file the listing lacks, through the same guarded write
sync_root applies to a row whose file has vanished."""
if not observations:
return
marked: list[int] = []
try:
with create_write_session() as session:
marked = apply_reference_observations(session, observations)
session.commit()
_write_in_batches(observations, apply_reference_observations, should_stop, marked)
except Exception as exc:
logging.exception("fast DB scan failed for %s: %s", root, exc)
emit(
@@ -526,9 +575,8 @@ def mark_unlisted_references_missing_safely(
error_type=error_type(exc),
error_kind=error_kind(exc),
)
return
if progress is not None:
progress.missing_marked += marked
progress.missing_marked += sum(marked)
def list_output_for_rescan() -> ListingWalk:

120
app/assets/seeder.py
View File

@@ -43,6 +43,15 @@ class ScanInProgressError(Exception):
"""Raised when an operation cannot proceed because a scan is running."""
class PruneCancelledError(Exception):
"""A standalone prune stopped by a cancel. The batches before it stay committed;
``marked`` counts them."""
def __init__(self, marked: int) -> None:
super().__init__(f"prune cancelled after marking {marked}")
self.marked = marked
class State(Enum):
"""Seeder state machine states."""
@@ -172,6 +181,11 @@ class _AssetSeeder:
self._cancel_event = threading.Event()
self._run_gate = threading.Event()
self._run_gate.set() # Start unpaused (set = running, clear = paused)
# Clear while a standalone prune holds the seeder (set = no prune running).
self._prune_idle = threading.Event()
self._prune_idle.set()
# Set by shutdown(): a standalone prune that has not started by then does not.
self._shutting_down = False
self._roots: tuple[RootType, ...] = ()
self._phase: ScanPhase = ScanPhase.FULL
self._compute_hashes: bool = False
@@ -447,8 +461,14 @@ class _AssetSeeder:
Returns:
True if the scan thread joined cleanly; False on timeout.
"""
with self._lock:
self._shutting_down = True
self.cancel()
# A standalone prune stops at its next batch once cancelled. One deadline
# covers both waits.
deadline = time.monotonic() + timeout
joined = self.wait(timeout=timeout)
joined = self.wait_for_standalone_prune(max(0.0, deadline - time.monotonic())) and joined
if not joined:
logging.warning(
"Asset seeder thread did not exit within %ss",
@@ -480,13 +500,18 @@ class _AssetSeeder:
Raises:
ScanInProgressError: If a scan is currently running
PruneCancelledError: If a cancel stopped it part way
"""
with self._lock:
if self._state != State.IDLE:
raise ScanInProgressError(
"Cannot mark missing assets while scan is running"
)
if self._shutting_down:
raise PruneCancelledError(0)
self._state = State.RUNNING
self._cancel_event.clear()
self._prune_idle.clear()
try:
if not dependencies_available():
@@ -496,7 +521,16 @@ class _AssetSeeder:
return 0
all_prefixes = get_owned_prefixes()
marked = mark_missing_outside_prefixes_safely(all_prefixes)
# Not pausable (the API waits on it), but a cancel or shutdown stops it
# between batches: it runs on a worker thread that exit would wait for.
stopped = False
def should_stop() -> bool:
nonlocal stopped
stopped = self._cancel_event.is_set()
return stopped
marked = mark_missing_outside_prefixes_safely(all_prefixes, should_stop)
if marked is None:
return None
emit(
@@ -504,12 +538,32 @@ class _AssetSeeder:
count=marked,
stage=_ScanStage.MARK_MISSING.value,
)
if stopped:
logging.info("Marking missing assets cancelled after marking %d", marked)
raise PruneCancelledError(marked)
if marked > 0:
logging.info("Marked %d references as missing", marked)
return marked
finally:
# The API runs this off the event loop, so a prompt can finish meanwhile
# and queue its output rescan; start it now. Not after a cancel: shutdown
# cancels, and a scan started here would run on into teardown. It stays
# queued for the next scan to start.
with self._lock:
self._reset_to_idle()
try:
if self._cancel_event.is_set():
self._reset_to_idle()
else:
self._finish_and_start_pending()
finally:
self._prune_idle.set()
def standalone_prune_running(self) -> bool:
return not self._prune_idle.is_set()
def wait_for_standalone_prune(self, timeout: float | None = None) -> bool:
"""Block until no standalone prune holds the seeder. True unless it timed out."""
return self._prune_idle.wait(timeout)
def _reset_to_idle(self) -> None:
"""Reset state to IDLE, preserving last progress. Caller must hold _lock."""
@@ -655,11 +709,13 @@ class _AssetSeeder:
if self._prune_first:
all_prefixes = get_owned_prefixes()
marked = mark_missing_outside_prefixes_safely(all_prefixes)
marked = mark_missing_outside_prefixes_safely(
all_prefixes, lambda: self._check_pause_and_cancel(_ScanStage.PRUNING)
)
marked_count = 0 if marked is None else marked
if marked is None:
self._add_error(
"Marking missing assets failed; scan continued without pruning"
"Marking missing assets failed; scan continued with the prune incomplete"
)
else:
emit(
@@ -671,7 +727,9 @@ class _AssetSeeder:
logging.info(
"Marked %d refs as missing before scan", marked_count
)
sync_temp_references_safely(scan_state)
sync_temp_references_safely(
scan_state, lambda: self._check_pause_and_cancel(_ScanStage.PRUNING)
)
if self._check_pause_and_cancel(_ScanStage.PRUNING):
logging.info("Asset scan cancelled after pruning phase")
@@ -801,23 +859,28 @@ class _AssetSeeder:
)
finally:
with self._lock:
start_paused = self._state is State.PAUSED
self._reset_to_idle()
pending = self._pending_scan
if pending is not None:
self._pending_scan = None
if not self.start(
roots=pending["roots"],
phase=pending["phase"],
prune_first=False,
compute_hashes=pending["compute_hashes"],
_start_paused=start_paused,
):
logging.warning(
"Pending scan could not start (roots=%s, phase=%s)",
pending["roots"],
pending["phase"].value,
)
self._finish_and_start_pending()
def _finish_and_start_pending(self) -> None:
"""Reset to IDLE, then start the scan queued while this run held the seeder,
paused if this run was. Caller must hold _lock."""
start_paused = self._state is State.PAUSED
self._reset_to_idle()
pending = self._pending_scan
if pending is not None:
self._pending_scan = None
if not self.start(
roots=pending["roots"],
phase=pending["phase"],
prune_first=False,
compute_hashes=pending["compute_hashes"],
_start_paused=start_paused,
):
logging.warning(
"Pending scan could not start (roots=%s, phase=%s)",
pending["roots"],
pending["phase"].value,
)
@staticmethod
def _emit_marked_missing(root: RootType, marked: int) -> None:
@@ -854,7 +917,11 @@ class _AssetSeeder:
existing_paths.update(live_references)
else:
marked_before = scan_state.missing_marked
existing_paths.update(sync_root_safely(r, scan_state))
existing_paths.update(
sync_root_safely(
r, scan_state, lambda: self._check_pause_and_cancel(_ScanStage.FAST_SCAN)
)
)
self._emit_marked_missing(r, scan_state.missing_marked - marked_before)
logging.debug(
"Fast scan: sync_root phase took %.3fs (%d existing paths)",
@@ -877,7 +944,12 @@ class _AssetSeeder:
scan_state.dirs_listed += walk.dirs_listed
vanished, unlisted = unlisted_references(live_references, walk.listings, scan_state)
marked_before = scan_state.missing_marked
mark_unlisted_references_missing_safely("output", vanished, scan_state)
mark_unlisted_references_missing_safely(
"output",
vanished,
scan_state,
lambda: self._check_pause_and_cancel(_ScanStage.FAST_SCAN),
)
self._emit_marked_missing("output", scan_state.missing_marked - marked_before)
logging.debug(
"Fast scan: output listing: %d dirs listed, %d rows retired, "

1
openapi.yaml
View File

@@ -2402,6 +2402,7 @@ paths:
description: Number of assets marked as missing
type: integer
status:
description: '"completed", or "cancelled" when a cancel stopped the prune part way; the assets it marked before that stay marked'
type: string
type: object
description: Prune result

722
tests-unit/assets_test/services/test_batched_marking.py Normal file
View File

@@ -0,0 +1,722 @@
"""The prune and the offline marking write in short batches rather than one long
transaction, so a foreground write gets the lock between batches. These tests pin the
batch boundaries, what a failure or a stop leaves behind, the pause between batches,
what a batch does with rows another writer changed, and the set-based mark."""
import asyncio
import json
import os
import sqlite3
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import patch
import pytest
import sqlalchemy as sa
from aiohttp.test_utils import make_mocked_request
from sqlalchemy import event
from sqlalchemy.orm import Session as SASession, sessionmaker
from sqlalchemy.pool import StaticPool
from app.assets import scanner, seeder as seeder_module
from app.assets.api import routes
from app.assets.database.models import Asset, AssetContent, AssetTag, Base, Tag
from app.assets.database.queries.records import (
create_content,
create_record,
ensure_tag,
ensure_tag_link,
mark_content_missing,
mark_contents_missing,
)
from app.assets.scanner_admission import _WATCH_LIST
from app.assets.seeder import State
@pytest.fixture
def db_engine():
"""One in-memory database every thread shares, for the tests that prune on a worker."""
engine = sa.create_engine("sqlite://", poolclass=StaticPool, connect_args={"check_same_thread": False})
Base.metadata.create_all(engine)
return engine
@pytest.fixture(autouse=True)
def no_yield(monkeypatch):
monkeypatch.setattr(scanner, "WRITE_YIELD_MIN_SECONDS", 0.0)
monkeypatch.setattr(scanner, "WRITE_YIELD_MAX_SECONDS", 0.0)
@pytest.fixture
def catalog(db_engine):
"""Routes the scanner's sessions to the test engine and counts write transactions."""
opened: list[int] = []
write_factory = sessionmaker(bind=db_engine)
@contextmanager
def _create_session():
with SASession(db_engine) as sess:
yield sess
def _create_write_session():
opened.append(1)
return write_factory()
_WATCH_LIST.clear()
with patch("app.assets.scanner.create_session", _create_session), \
patch("app.assets.seeder.create_session", _create_session), \
patch("app.assets.scanner.create_write_session", _create_write_session):
yield opened
_WATCH_LIST.clear()
def _rows(session, directory: Path, count: int, *, files: bool = False) -> list[str]:
ids = []
for i in range(count):
path = directory / f"f{i:05d}.png"
if files:
path.write_bytes(f"bytes-{i}".encode())
stat = path.stat() if files else None
content = create_content(
session,
path=str(path),
size_bytes=stat.st_size if stat else 1,
mtime_ns=stat.st_mtime_ns if stat else 1,
)
create_record(session, content_id=content.id, name=path.name, tags=["output"])
ids.append(content.id)
session.commit()
return ids
def _live_ids(session) -> set[str]:
session.expire_all()
return set(session.scalars(sa.select(AssetContent.id).where(AssetContent.is_missing == sa.false())))
def _prune(owned: list[str]) -> int | None:
return scanner.mark_missing_outside_prefixes_safely(owned)
@pytest.mark.parametrize("count, batches", [(0, 0), (1, 1), (256, 1), (257, 2), (600, 3)])
def test_prune_commits_one_batch_per_256_rows(session, catalog, temp_dir, count, batches):
_rows(session, temp_dir, count)
assert _prune([]) == count
assert len(catalog) == batches
assert _live_ids(session) == set()
def test_rows_under_an_owned_prefix_are_not_pruned(session, catalog, temp_dir):
owned, gone = temp_dir / "owned", temp_dir / "gone"
owned.mkdir()
gone.mkdir()
kept = set(_rows(session, owned, 5))
_rows(session, gone, 5)
assert _prune([str(owned)]) == 5
assert _live_ids(session) == kept
def test_a_row_re_registered_between_batches_is_left_to_its_new_owner(session, catalog, temp_dir):
"""register_executed_output retires the live row at a path and inserts a new one.
Doing that between two batches must leave exactly the new row live."""
ids = _rows(session, temp_dir, 300)
last = session.get(AssetContent, ids[-1])
replacement: list[str] = []
def between_batches() -> bool:
if len(catalog) == 1:
mark_content_missing(session, last.id)
new = create_content(session, path=last.path, size_bytes=2, mtime_ns=2)
session.commit()
replacement.append(new.id)
return False
marked = scanner.mark_missing_outside_prefixes_safely([], between_batches)
# The replacement was never a candidate; the retired row is not counted twice.
assert marked == len(ids) - 1
assert _live_ids(session) == set(replacement)
def test_a_foreground_write_gets_the_lock_between_batches(tmp_path, monkeypatch):
"""On a real file database with the production write-session setup, another
connection can take the write lock at every point between two batches."""
db = tmp_path / "catalog.db"
read_engine = sa.create_engine(f"sqlite:///{db}")
write_engine = sa.create_engine(f"sqlite:///{db}")
@event.listens_for(write_engine, "connect")
def _connect(dbapi_connection, _record):
dbapi_connection.isolation_level = None
@event.listens_for(write_engine, "begin")
def _begin(connection):
connection.exec_driver_sql("BEGIN IMMEDIATE")
with read_engine.connect() as conn:
conn.exec_driver_sql("PRAGMA journal_mode=WAL")
Base.metadata.create_all(read_engine)
with SASession(read_engine) as session:
_rows(session, tmp_path, 600)
foreground: list[bool] = []
def between_batches() -> bool:
other = sqlite3.connect(db, timeout=0, isolation_level=None)
try:
other.execute("BEGIN IMMEDIATE")
other.execute("INSERT INTO tags (name) VALUES (?)", (f"fg-{len(foreground)}",))
other.execute("COMMIT")
foreground.append(True)
except sqlite3.OperationalError:
foreground.append(False)
finally:
other.close()
return False
monkeypatch.setattr(scanner, "create_session", lambda: SASession(read_engine))
monkeypatch.setattr(scanner, "create_write_session", sessionmaker(bind=write_engine))
assert scanner.mark_missing_outside_prefixes_safely([], between_batches) == 600
assert foreground == [True, True, True]
def test_a_failed_batch_keeps_the_batches_before_it(session, catalog, temp_dir):
_rows(session, temp_dir, 600)
real = scanner.mark_contents_missing
def fail_in_the_second_batch(sess, ids):
if len(catalog) == 2:
raise RuntimeError("disk I/O error")
return real(sess, ids)
with patch("app.assets.scanner.mark_contents_missing", fail_in_the_second_batch):
assert _prune([]) is None
assert len(_live_ids(session)) == 600 - scanner.WRITE_BATCH_ROWS
def test_sync_root_counts_the_batches_committed_before_a_failure(session, catalog, temp_dir, monkeypatch):
output = temp_dir / "output"
output.mkdir()
monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(output))
_rows(session, output, 300)
real = scanner.mark_contents_missing
def fail_in_the_second_batch(sess, ids):
if len(catalog) == 2:
raise RuntimeError("disk I/O error")
return real(sess, ids)
progress = seeder_module._ScanState()
with patch("app.assets.scanner.mark_contents_missing", fail_in_the_second_batch):
assert scanner.sync_root_safely("output", progress) == set()
assert progress.missing_marked == scanner.WRITE_BATCH_ROWS
def _gone_observations(session, temp_dir: Path, count: int) -> list[scanner._ReferenceObservation]:
ids = _rows(session, temp_dir, count, files=True)
observations = []
for content_id in ids:
content = session.get(AssetContent, content_id)
os.remove(content.path)
observations.append(
scanner._ReferenceObservation(content.id, content.size_bytes, content.mtime_ns, None)
)
return observations
def test_stop_between_batches_leaves_the_rest_live(session, catalog, temp_dir):
observations = _gone_observations(session, temp_dir, 600)
committed: list[int] = []
scanner._write_in_batches(observations, scanner.apply_reference_observations, lambda: bool(committed), committed)
assert committed == [scanner.WRITE_BATCH_ROWS]
assert len(_live_ids(session)) == 600 - scanner.WRITE_BATCH_ROWS
def test_a_scan_waits_between_batches_while_a_prompt_runs(session, catalog, temp_dir, monkeypatch):
output = temp_dir / "output"
output.mkdir()
monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(output))
_rows(session, output, 600)
instance = seeder_module._AssetSeeder()
instance._state = State.RUNNING
instance._scan_state = seeder_module._ScanState()
instance._run_gate.set()
real = scanner.mark_contents_missing
first_batch = threading.Event()
def mark(sess, ids):
if not first_batch.is_set():
assert instance.pause() # a prompt starts during the first batch
first_batch.set()
return real(sess, ids)
def sync():
scanner.sync_root_safely(
"output",
instance._scan_state,
lambda: instance._check_pause_and_cancel(seeder_module._ScanStage.FAST_SCAN),
)
with patch("app.assets.scanner.mark_contents_missing", mark):
worker = threading.Thread(target=sync)
worker.start()
assert first_batch.wait(5)
time.sleep(0.2)
assert len(catalog) == 1 # no batch while paused
assert instance.resume()
worker.join(5)
assert len(catalog) == 3
assert instance._scan_state.missing_marked == 600
def test_offline_rows_retired_across_batches_recover_when_the_drive_returns(
session, catalog, temp_dir, monkeypatch
):
"""#16646's hashing-off recovery with the marking split over several batches: files
that come back while the marking is part way through are retired by the batches
that follow, then recovered by the walk that follows in the same scan, every record
keeping its id."""
output = temp_dir / "output"
(temp_dir / "input").mkdir()
output.mkdir()
monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(output))
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir / "input"))
monkeypatch.setattr(scanner, "get_comfy_models_folders", lambda: [])
ids = _rows(session, output, 700, files=True)
records = {record.id: record.content_id for record in session.scalars(sa.select(Asset))}
parked = temp_dir / "parked"
output.rename(parked)
output.mkdir()
instance = seeder_module._AssetSeeder()
instance._scan_state = seeder_module._ScanState()
instance._phase = seeder_module.ScanPhase.FAST
instance._run_gate.set()
real = scanner.mark_contents_missing
def mark(sess, content_ids):
marked = real(sess, content_ids)
if len(catalog) == 1:
# The drive comes back after the first batch.
for name in os.listdir(parked):
os.rename(parked / name, output / name)
return marked
with patch("app.assets.scanner.mark_contents_missing", mark):
instance._run_fast_phase(("input", "output"))
session.expire_all()
assert instance._scan_state.missing_marked == 700
assert instance._scan_state.recovered == 700
assert _live_ids(session) == set(ids)
assert {record.id: record.content_id for record in session.scalars(sa.select(Asset))} == records
live_paths = session.scalars(sa.select(AssetContent.path).where(AssetContent.is_missing == sa.false())).all()
assert len(live_paths) == len(set(live_paths)) == 700
def _link_state(session) -> tuple[dict[str, bool], set[tuple[str, str, str]]]:
session.expire_all()
contents = {c.path: c.is_missing for c in session.scalars(sa.select(AssetContent))}
links = {
(record.name, link.tag_name, link.origin)
for record in session.scalars(sa.select(Asset))
for link in session.scalars(sa.select(AssetTag).where(AssetTag.asset_id == record.id))
}
return contents, links
def _equivalence_fixture(session) -> list[str]:
"""Rows covering each case the mark handles: no record, one, two, a record already
tagged missing by hand, an already-missing row, and an id that does not exist."""
none = create_content(session, path="/c/none.png")
one = create_content(session, path="/c/one.png")
create_record(session, content_id=one.id, name="one", tags=["output"])
two = create_content(session, path="/c/two.png")
create_record(session, content_id=two.id, name="two-a")
create_record(session, content_id=two.id, name="two-b", tags=["input"])
tagged = create_content(session, path="/c/tagged.png")
record = create_record(session, content_id=tagged.id, name="tagged")
ensure_tag(session, "missing")
ensure_tag_link(session, asset_id=record.id, tag_name="missing", origin="manual")
already = create_content(session, path="/c/already.png")
create_record(session, content_id=already.id, name="already")
mark_content_missing(session, already.id)
session.commit()
return [none.id, one.id, two.id, tagged.id, already.id, "no-such-id"]
def test_the_set_mark_matches_marking_row_by_row(session):
engine = sa.create_engine("sqlite:///:memory:")
Base.metadata.create_all(engine)
with SASession(engine) as per_row:
# Same fixture in both catalogs, then the ids made to match by path.
ids = _equivalence_fixture(session)
_equivalence_fixture(per_row)
by_path = {c.path: c.id for c in per_row.scalars(sa.select(AssetContent))}
paths = {c.id: c.path for c in session.scalars(sa.select(AssetContent))}
marked = mark_contents_missing(session, ids)
session.commit()
for content_id in ids:
path = paths.get(content_id)
content = per_row.get(AssetContent, by_path[path]) if path else None
if content is not None and not content.is_missing:
mark_content_missing(per_row, content.id)
per_row.commit()
# create_content stores os.path.abspath(path), so compare in that form (a drive letter on Windows).
expected = sorted(os.path.abspath(f"/c/{name}.png") for name in ("none", "one", "tagged", "two"))
assert sorted(paths[i] for i in marked) == expected
assert _link_state(session) == _link_state(per_row)
assert session.get(Tag, "missing") is not None
def test_the_set_mark_settles_a_link_race_row_by_row(session, monkeypatch):
content = create_content(session, path="/c/raced.png")
record = create_record(session, content_id=content.id, name="raced")
session.commit()
real_execute = session.execute
def insert_loses_the_race(statement, *args, **kwargs):
if isinstance(statement, sa.Insert) and statement.table is AssetTag.__table__:
raise sa.exc.IntegrityError("INSERT", {}, Exception("UNIQUE constraint failed"))
return real_execute(statement, *args, **kwargs)
monkeypatch.setattr(session, "execute", insert_loses_the_race)
assert mark_contents_missing(session, [content.id]) == [content.id]
session.commit()
link = session.get(AssetTag, (record.id, "missing"))
assert link is not None and link.origin == "automatic"
def test_the_batched_writes_run_no_table_scan_inside_their_transactions(session, db_engine, temp_dir, monkeypatch):
"""Each statement a batch runs while it holds the write lock is looked up by key or
index, so the lock window grows with the batch, not the catalog. (The prune's
candidate read is a full read by design, and runs before any write transaction.)"""
in_write: list[bool] = []
class _Tracked(SASession):
def __enter__(self):
in_write.append(True)
return super().__enter__()
def __exit__(self, *exc_info):
in_write.clear()
return super().__exit__(*exc_info)
monkeypatch.setattr(scanner, "create_write_session", sessionmaker(bind=db_engine, class_=_Tracked))
monkeypatch.setattr(scanner, "create_session", lambda: SASession(db_engine))
# A size change splits the row, and the new record's tags come from its root.
monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(temp_dir))
gone_dir, changed_dir, pruned_dir = (temp_dir / name for name in ("gone", "changed", "pruned"))
for directory in (gone_dir, changed_dir, pruned_dir):
directory.mkdir()
observations = _gone_observations(session, gone_dir, 50)
for content_id in _rows(session, changed_dir, 2, files=True):
content = session.get(AssetContent, content_id)
if content.path.endswith("0.png"):
Path(content.path).write_bytes(b"a different size")
os.utime(content.path, ns=(10**18, 10**18))
observations.append(
scanner._ReferenceObservation(content.id, content.size_bytes, content.mtime_ns, os.stat(content.path))
)
statements: list[tuple[str, object]] = []
def capture(conn, cursor, statement, parameters, context, executemany):
if in_write and statement.lstrip().upper().startswith(("SELECT", "UPDATE", "INSERT", "DELETE")):
statements.append((statement, parameters))
event.listen(db_engine, "before_cursor_execute", capture)
try:
scanner._write_in_batches(observations, scanner.apply_reference_observations, lambda: False, [])
_rows(session, pruned_dir, 50)
assert scanner.mark_missing_outside_prefixes_safely([str(gone_dir), str(changed_dir)]) == 50
finally:
event.remove(db_engine, "before_cursor_execute", capture)
kinds = {statement.split()[0].upper() for statement, _ in statements}
assert {"SELECT", "UPDATE", "INSERT"} <= kinds
with db_engine.connect() as conn:
for statement, parameters in statements:
plan = conn.exec_driver_sql(f"EXPLAIN QUERY PLAN {statement}", parameters).all()
scans = [row[-1] for row in plan if row[-1].startswith("SCAN")]
assert not scans, (statement, plan)
@pytest.mark.asyncio
async def test_the_prune_endpoint_keeps_the_event_loop_serving(monkeypatch):
def slow_prune() -> int:
time.sleep(0.5)
return 3
monkeypatch.setattr(routes.asset_seeder, "mark_missing_outside_prefixes", slow_prune)
ticks = 0
async def ticker():
nonlocal ticks
while True:
await asyncio.sleep(0.01)
ticks += 1
ticking = asyncio.create_task(ticker())
response = await routes.mark_missing_assets.__wrapped__(make_mocked_request("POST", "/api/assets/prune"))
ticking.cancel()
assert json.loads(response.body) == {"status": "completed", "marked": 3}
assert ticks >= 20
def test_the_standalone_prune_starts_the_scan_queued_while_it_ran(session, catalog, temp_dir, monkeypatch):
"""The API runs the prune off the event loop, so a prompt can finish meanwhile and
queue its output rescan, which cannot start while the prune holds the seeder."""
_rows(session, temp_dir, 10)
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
started: list[dict] = []
def start(**kwargs) -> bool:
if instance._state is not State.IDLE:
return False
started.append(kwargs)
return True
monkeypatch.setattr(instance, "start", start)
real = scanner.mark_contents_missing
def mark(sess, ids):
assert instance.enqueue_scan(roots=("output",), phase=seeder_module.ScanPhase.FULL) is False
return real(sess, ids)
with patch("app.assets.scanner.mark_contents_missing", mark):
assert instance.mark_missing_outside_prefixes() == 10
assert [kwargs["roots"] for kwargs in started] == [("output",)]
assert instance._pending_scan is None
def _seeder_with_recorded_starts(monkeypatch) -> tuple[seeder_module._AssetSeeder, list[tuple]]:
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
started: list[tuple] = []
def start(roots=("models", "input", "output"), **kwargs) -> bool:
if instance._state is not State.IDLE:
return False
started.append(tuple(roots))
return True
monkeypatch.setattr(instance, "start", start)
monkeypatch.setattr(routes, "asset_seeder", instance)
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
return instance, started
@pytest.mark.asyncio
async def test_a_seed_request_during_an_api_prune_waits_for_it_then_starts(monkeypatch):
"""The prune runs off the event loop now, so a seed request can arrive while it
holds the seeder. A 409 there would tell the client a scan is coming when none is."""
instance, started = _seeder_with_recorded_starts(monkeypatch)
release = threading.Event()
pruning = threading.Event()
def blocking_prune(prefixes, should_stop=None):
pruning.set()
assert release.wait(5)
return 0
monkeypatch.setattr(seeder_module, "mark_missing_outside_prefixes_safely", blocking_prune)
monkeypatch.setattr(routes, "_PRUNE_POLL_SECONDS", 0.01)
prune = asyncio.create_task(asyncio.to_thread(instance.mark_missing_outside_prefixes))
assert await asyncio.to_thread(pruning.wait, 5)
def must_not_block_a_thread(timeout=None):
raise AssertionError("the seed route held an executor thread for the prune")
# The route waits on the loop; a blocking wait would hold an executor thread per request.
monkeypatch.setattr(instance, "wait_for_standalone_prune", must_not_block_a_thread)
seed = asyncio.create_task(routes.seed_assets.__wrapped__(make_mocked_request("POST", "/api/assets/seed")))
await asyncio.sleep(0.2)
assert not seed.done() # waiting out the prune, not answering 409
release.set()
response = await seed
await prune
assert response.status == 202
assert started == [("models", "input", "output")]
@pytest.mark.asyncio
async def test_a_seed_request_during_a_scan_still_gets_409(monkeypatch):
instance, started = _seeder_with_recorded_starts(monkeypatch)
instance._state = State.RUNNING
response = await routes.seed_assets.__wrapped__(make_mocked_request("POST", "/api/assets/seed"))
assert response.status == 409
assert started == []
def test_a_cancel_stops_a_standalone_prune_and_shutdown_waits_for_it(session, catalog, temp_dir, monkeypatch):
"""The API prune runs on a worker thread that interpreter exit joins, so shutdown's
cancel must stop it between batches rather than let it run to the end."""
_rows(session, temp_dir, 600)
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
real = scanner.mark_contents_missing
first_batch = threading.Event()
shut_down = threading.Event()
def mark(sess, ids):
marked = real(sess, ids)
first_batch.set()
assert shut_down.wait(5) # hold the first batch until shutdown has cancelled
return marked
result: list[object] = []
def prune() -> None:
try:
result.append(instance.mark_missing_outside_prefixes())
except seeder_module.PruneCancelledError as cancelled:
result.append(cancelled)
with patch("app.assets.scanner.mark_contents_missing", mark):
worker = threading.Thread(target=prune)
worker.start()
assert first_batch.wait(5)
threading.Timer(0.1, shut_down.set).start()
assert instance.shutdown(timeout=5)
worker.join(5)
assert len(catalog) == 1
assert len(result) == 1 and isinstance(result[0], seeder_module.PruneCancelledError)
assert result[0].marked == scanner.WRITE_BATCH_ROWS
assert not instance.standalone_prune_running()
@pytest.mark.asyncio
async def test_a_cancelled_api_prune_is_not_reported_as_completed(monkeypatch):
def cancelled_prune() -> int:
raise seeder_module.PruneCancelledError(256)
monkeypatch.setattr(routes.asset_seeder, "mark_missing_outside_prefixes", cancelled_prune)
response = await routes.mark_missing_assets.__wrapped__(make_mocked_request("POST", "/api/assets/prune"))
assert response.status == 200
assert json.loads(response.body) == {"status": "cancelled", "marked": 256}
def test_a_prune_that_finishes_before_a_late_cancel_reports_completed(session, catalog, temp_dir, monkeypatch):
"""The cancel only counts if it stopped a batch from running."""
_rows(session, temp_dir, 10)
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
real = scanner.mark_contents_missing
def mark(sess, ids):
marked = real(sess, ids)
instance.cancel() # arrives during the last (only) batch
return marked
with patch("app.assets.scanner.mark_contents_missing", mark):
assert instance.mark_missing_outside_prefixes() == 10
def test_shutdown_during_a_prune_does_not_start_the_scan_a_prompt_queued(session, catalog, temp_dir, monkeypatch):
"""A scan started after shutdown cancelled the prune would run on into teardown."""
_rows(session, temp_dir, 600)
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
started: list[tuple] = []
def start(roots=("models", "input", "output"), **kwargs) -> bool:
if instance._state is not State.IDLE:
return False
started.append(tuple(roots))
return True
monkeypatch.setattr(instance, "start", start)
real = scanner.mark_contents_missing
first_batch = threading.Event()
shut_down = threading.Event()
def mark(sess, ids):
marked = real(sess, ids)
# A prompt finishes and queues its output rescan; the prune holds the seeder.
assert instance.enqueue_scan(roots=("output",), phase=seeder_module.ScanPhase.FULL) is False
first_batch.set()
assert shut_down.wait(5)
return marked
outcome: list[object] = []
def prune() -> None:
try:
outcome.append(instance.mark_missing_outside_prefixes())
except seeder_module.PruneCancelledError as cancelled:
outcome.append(cancelled)
with patch("app.assets.scanner.mark_contents_missing", mark):
worker = threading.Thread(target=prune)
worker.start()
assert first_batch.wait(5)
threading.Timer(0.1, shut_down.set).start()
assert instance.shutdown(timeout=5)
worker.join(5)
# Asserted here, not in the worker: a failure there would only surface as a warning.
assert len(outcome) == 1 and isinstance(outcome[0], seeder_module.PruneCancelledError)
assert started == []
assert instance._state is State.IDLE
def test_shutdown_before_a_prune_starts_keeps_it_from_starting(session, catalog, temp_dir, monkeypatch):
"""The API hands the prune to a worker thread; a shutdown that lands before it takes
the seeder must still keep it from running into teardown."""
_rows(session, temp_dir, 10)
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
assert instance.shutdown(timeout=1)
with pytest.raises(seeder_module.PruneCancelledError) as cancelled:
instance.mark_missing_outside_prefixes()
assert cancelled.value.marked == 0
assert len(catalog) == 0
assert len(_live_ids(session)) == 10
assert not instance.standalone_prune_running()
def test_the_prune_flag_clears_even_if_its_cleanup_raises(session, catalog, temp_dir, monkeypatch):
_rows(session, temp_dir, 10)
instance = seeder_module._AssetSeeder()
monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
def start_fails():
raise RuntimeError("can't start new thread")
monkeypatch.setattr(instance, "_finish_and_start_pending", start_fails)
with pytest.raises(RuntimeError):
instance.mark_missing_outside_prefixes()
assert not instance.standalone_prune_running()
assert instance.wait_for_standalone_prune(0)

12
tests-unit/assets_test/services/test_scanner_b.py
View File

@@ -5,14 +5,13 @@ from pathlib import Path
from unittest.mock import patch
from sqlalchemy import select
from sqlalchemy.orm import Session
from sqlalchemy.orm import Session, sessionmaker
from app.assets.database.models import Asset, AssetContent, AssetTag
from app.assets.helpers import to_stored_hash
from app.assets.scanner import (
build_asset_specs,
enrich_asset,
mark_contents_missing_outside_prefixes,
mark_missing_outside_prefixes_safely,
seed_asset_specs,
apply_reference_observations,
@@ -151,7 +150,7 @@ def test_seed_creates_content_and_record(session, temp_dir: Path):
]
def test_prune_marks_missing_not_deletes(session, temp_dir: Path):
def test_prune_marks_missing_not_deletes(session, db_engine, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
file_path = input_root / "removed-from-registry.png"
@@ -161,8 +160,11 @@ def test_prune_marks_missing_not_deletes(session, temp_dir: Path):
seed_asset_specs(session, _build_seed_specs(input_root))
session.commit()
marked = mark_contents_missing_outside_prefixes(session, prefixes=[])
session.commit()
with patch("app.assets.scanner.create_session", sessionmaker(bind=db_engine)), \
patch("app.database.db.WriteSession", sessionmaker(bind=db_engine)), \
patch("app.assets.scanner.get_owned_prefixes", return_value=[]):
marked = mark_missing_outside_prefixes_safely([])
session.expire_all()
content = session.scalar(select(AssetContent))
record = session.scalar(select(Asset))

20
tests-unit/seeder_test/test_seeder.py
View File

@@ -93,7 +93,7 @@ def _configure_fast_phase(
specs: list[SeedAssetSpec],
) -> None:
monkeypatch.setattr(
seeder_module, "sync_root_safely", lambda _root, _progress: set()
seeder_module, "sync_root_safely", lambda _root, _progress, _should_stop=None: set()
)
monkeypatch.setattr(
seeder_module,
@@ -435,7 +435,7 @@ def test_idle_reset_survives_a_raising_cancellation_emit(
monkeypatch.setattr(scan_seeder, "_check_pause_and_cancel", cancel_at_pruning)
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ())
monkeypatch.setattr(
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 0
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 0
)
original_emit = seeder_module.emit
@@ -462,7 +462,7 @@ def test_scan_paused_after_its_last_phase_still_completes(
scan_seeder._phase = ScanPhase.ENRICH
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ())
monkeypatch.setattr(
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 0
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 0
)
def pause_while_finishing(roots) -> tuple[bool, int]:
@@ -507,10 +507,10 @@ def test_prune_before_scan_emits_marked_missing_with_pruning_stage(
scan_seeder._phase = ScanPhase.FAST
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ())
monkeypatch.setattr(
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 5
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 5
)
monkeypatch.setattr(
seeder_module, "sync_temp_references_safely", lambda _progress: None
seeder_module, "sync_temp_references_safely", lambda _progress, _should_stop=None: None
)
monkeypatch.setattr(scan_seeder, "_run_fast_phase", lambda roots: (0, 0, 0))
@@ -530,7 +530,7 @@ def test_standalone_mark_missing_emits_count_with_mark_missing_stage(
scan_seeder._state = State.IDLE
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ())
monkeypatch.setattr(
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 7
seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 7
)
with caplog.at_level(logging.INFO):
@@ -575,10 +575,10 @@ def test_scan_prune_failure_is_reported_and_the_scan_still_runs(
fast_phase_roots: list[tuple[str, ...]] = []
monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: [])
monkeypatch.setattr(
seeder_module, "mark_missing_outside_prefixes_safely", lambda _prefixes: None
seeder_module, "mark_missing_outside_prefixes_safely", lambda _prefixes, _should_stop=None: None
)
monkeypatch.setattr(
seeder_module, "sync_temp_references_safely", lambda _progress: None
seeder_module, "sync_temp_references_safely", lambda _progress, _should_stop=None: None
)
def run_fast_phase(roots: tuple[str, ...]) -> tuple[int, int, int]:
@@ -591,7 +591,7 @@ def test_scan_prune_failure_is_reported_and_the_scan_still_runs(
scan_seeder._run_scan()
assert scan_seeder._errors == [
"Marking missing assets failed; scan continued without pruning"
"Marking missing assets failed; scan continued with the prune incomplete"
]
assert fast_phase_roots == [("models", "input")]
assert events_named(caplog, "seeder.marked_missing") == []
@@ -604,7 +604,7 @@ def test_batch_insert_failure_emits_only_the_exception_type(
) -> None:
session = Mock()
monkeypatch.setattr(
seeder_module, "sync_root_safely", lambda _root, _progress: set()
seeder_module, "sync_root_safely", lambda _root, _progress, _should_stop=None: set()
)
monkeypatch.setattr(
seeder_module,