nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(scan): scan one selected source

Browse Source 
POST /prompt_manager/jobs/scan accepts roots: indexes into the effective
sources, strict non-negative ints (no bools or floats), deduplicated, at
most 64, non-empty when present. An index past the sources is a 400
"unknown-source" and starts nothing. The job walks only the selected
sources, resolved when it was requested, and keeps each file's listing
root_index. Its scope records files, roots (null for every source) and
output_folders, the number of folders walked, which the log's first line
counts. The legacy /scan route still scans every source, and attaching to
a scoped job keeps its payload keys.
This commit is contained in:
Vito Sansevero
2026-10-23 08:09:06 -07:00
parent b3d2e271fb
commit fa98f02fa8
3 changed files with 275 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/admin.py
+34 -18
View File
@@ -4,6 +4,7 @@ import asyncio
import contextlib
from contextlib import closing
import datetime
import functools
import hashlib
import json
import math
@@ -1344,28 +1345,34 @@ class AdminRoutesMixin:
}
)
def _start_scan_job(self, files=DEFAULT_SCAN_SCOPE):
def _start_scan_job(self, files=DEFAULT_SCAN_SCOPE, roots=None, output_dirs=None):
"""Register the scan job synchronously; its runner starts on the loop.
``files`` is the scope (``SCAN_SCOPES``). Raises JobBusy while another
record-changing job runs.
``files`` is the scope (``SCAN_SCOPES``). ``roots`` selects sources by
their index in ``output_dirs`` (the effective sources the caller
resolved); None scans every source, resolved when the job starts.
Raises JobBusy while another record-changing job runs.
"""
if files not in SCAN_SCOPES:
raise ValueError(f"Unknown scan scope: {files!r}")
if roots is not None and (
output_dirs is None or not roots or max(roots) >= len(output_dirs)
):
raise ValueError("A scoped scan needs the sources its roots index")
job = self._jobs.start(
SCAN_JOB_KIND,
SCAN_JOB_TITLE,
record_changing=True,
scope={"files": files},
runner=self._run_scan_job,
scope={"files": files, "roots": None if roots is None else list(roots)},
runner=functools.partial(self._run_scan_job, output_dirs=output_dirs),
groups=SCAN_GROUP_NAMES,
)
self._scan_job = job
return job
async def _run_scan_job(self, job):
async def _run_scan_job(self, job, output_dirs=None):
final = None
async for payload in self._scan_images_events(job):
async for payload in self._scan_images_events(job, output_dirs):
job.publish(payload)
final = payload.get("type")
state = SCAN_FINAL_STATES.get(final, "finished")
@@ -1400,16 +1407,25 @@ class AdminRoutesMixin:
)
return response
async def _scan_images_events(self, job):
"""Yield progress payloads while scanning every configured output directory.
async def _scan_images_events(self, job, output_dirs=None):
"""Yield progress payloads while scanning the output directories.
The payloads are the legacy SSE events, unchanged; ``job`` also gets
its status line, totals and one outcome per processed file.
Every effective source is walked, or only those ``job.scope["roots"]``
selects (indexes into ``output_dirs``). The payloads are the legacy SSE
events, unchanged; ``job`` also gets its status line, totals and one
outcome per processed file. Rows keep each source's listing index.
"""
try:
self.logger.info("Starting image scan operation")
output_dirs = self._get_all_output_dirs()
if not output_dirs:
if output_dirs is None:
output_dirs = self._get_all_output_dirs()
selected = job.scope.get("roots")
walked = (
output_dirs
if selected is None
else [output_dirs[index] for index in selected]
)
if not walked:
self.logger.error("No output directories found")
message = (
"No output directories found. "
@@ -1419,18 +1435,18 @@ class AdminRoutesMixin:
yield {"type": "error", "message": message}
return
dir_names = [self._public_path(d) for d in output_dirs]
folders = _plural(len(output_dirs), "output folder")
job.scope["output_folders"] = len(output_dirs)
dir_names = [self._public_path(d) for d in walked]
folders = _plural(len(walked), "output folder")
job.scope["output_folders"] = len(walked)
scope = SCAN_SCOPE_LOG[job.scope.get("files", "all")]
job.log_line(f"Scan started · {folders} · {scope}")
job.update(status_line=f"Looking for media files in {folders}")
yield self._scan_progress(
0, f"Scanning {len(output_dirs)} directory(ies) for media files..."
0, f"Scanning {len(walked)} directory(ies) for media files..."
)
media_files = await self._run_in_executor(
self._collect_output_media_sync, output_dirs
self._collect_output_media_sync, walked
)
counts = dict.fromkeys(SCAN_COUNT_KEYS, 0)
job.update(
py/api/job_routes.py
+64 -17
View File
@@ -18,29 +18,63 @@ def _not_found():
return web.json_response({"success": False, "error": "Job not found"}, status=404)
def _bad_request(message):
return web.json_response({"success": False, "error": message}, status=400)
MAX_SCAN_ROOTS = 64
_SCOPE_ERROR = 'The scan scope must be "new" or "all"'
_ROOTS_ERROR = (
f"roots must be a non-empty list of up to {MAX_SCAN_ROOTS} source indexes"
)
async def _scan_scope(request):
"""The requested scan scope ("new" or "all"), or None when the body is invalid.
def _bad_request(message, code=None):
body = {"success": False, "error": message}
if code is not None:
body["code"] = code
return web.json_response(body, status=400)
An empty body (also a chunked one with no bytes), or one without
``scope``, asks for the default.
def _scan_roots(body):
"""``(roots, error)``: the sorted, deduplicated source indexes, or None.
Absent ``roots`` means every source. When present it is a non-empty list
of at most MAX_SCAN_ROOTS non-negative ints (never bools or floats).
"""
if "roots" not in body:
return None, None
raw = body["roots"]
if not isinstance(raw, list) or not raw or len(raw) > MAX_SCAN_ROOTS:
return None, _ROOTS_ERROR
for value in raw:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
return None, _ROOTS_ERROR
return sorted(set(raw)), None
async def _scan_request(request):
"""``({files, roots}, None)`` for a scan start, or ``(None, message)``.
``files`` is the scope ("new" or "all"); ``roots`` the selected source
indexes or None for every source. An empty body (also a chunked one with
no bytes), or one without ``scope``, asks for the default scope.
"""
default = {"files": DEFAULT_SCAN_SCOPE, "roots": None}
if not request.can_read_body:
return DEFAULT_SCAN_SCOPE
return default, None
text = await request.text()
if not text.strip(): # e.g. a chunked POST that sent no bytes
return DEFAULT_SCAN_SCOPE
return default, None
try:
body = json.loads(text)
except ValueError:
return None
return None, _SCOPE_ERROR
if not isinstance(body, dict):
return None
scope = body.get("scope", DEFAULT_SCAN_SCOPE)
return scope if scope in SCAN_SCOPES else None
return None, _SCOPE_ERROR
files = body.get("scope", DEFAULT_SCAN_SCOPE)
if files not in SCAN_SCOPES:
return None, _SCOPE_ERROR
roots, error = _scan_roots(body)
if error:
return None, error
return {"files": files, "roots": roots}, None
def busy_response(busy: JobBusy):
@@ -132,13 +166,26 @@ class JobRoutesMixin:
"""Start the metadata scan as a job (202), or 409 naming the busy job.
The body's ``scope`` is "new" (the default: new and unresolved files)
or "all"; anything else is a 400.
or "all". ``roots``, when given, selects sources by their index in
the effective sources (``_get_all_output_dirs``); an index past them
is a 400 ``unknown-source`` and starts nothing. The job walks the
sources resolved here, so a later source change cannot move it.
"""
files = await _scan_scope(request)
if files is None:
return _bad_request('The scan scope must be "new" or "all"')
scan, error = await _scan_request(request)
if error:
return _bad_request(error)
output_dirs = None
if scan["roots"] is not None:
output_dirs = await self._run_in_executor(self._get_all_output_dirs)
if scan["roots"][-1] >= len(output_dirs):
return _bad_request(
"That source is no longer configured. Reload the sources.",
code="unknown-source",
)
try:
job = self._start_scan_job(files)
job = self._start_scan_job(
scan["files"], roots=scan["roots"], output_dirs=output_dirs
)
except JobBusy as busy:
return busy_response(busy)
return web.json_response({"success": True, "job": job.summary()}, status=202)
tests/test_scan_scope.py
+177 -5
View File
@@ -260,7 +260,7 @@ class TestScanScope(ScanScopeTestCase):
async def test_empty_read_body_asks_for_the_default_scope(self):
"""A chunked POST with no bytes can read a body: still the default."""
from py.api.job_routes import _scan_scope
from py.api.job_routes import _scan_request
class Request:
can_read_body = True
@@ -271,10 +271,16 @@ class TestScanScope(ScanScopeTestCase):
async def text(self):
return self._text
self.assertEqual(await _scan_scope(Request("")), "new")
self.assertEqual(await _scan_scope(Request(" \n")), "new")
self.assertEqual(await _scan_scope(Request('{"scope": "all"}')), "all")
self.assertIsNone(await _scan_scope(Request("{")))
default = ({"files": "new", "roots": None}, None)
self.assertEqual(await _scan_request(Request("")), default)
self.assertEqual(await _scan_request(Request(" \n")), default)
self.assertEqual(
await _scan_request(Request('{"scope": "all", "roots": [2, 0, 2]}')),
({"files": "all", "roots": [0, 2]}, None),
)
request, error = await _scan_request(Request("{"))
self.assertIsNone(request)
self.assertIsInstance(error, str)
async def test_invalid_scope_is_refused(self):
for body in ({"scope": "everything"}, {"scope": 1}, ["new"]):
@@ -292,6 +298,172 @@ class TestScanScope(ScanScopeTestCase):
self.assertIsNone(self.api._scan_job)
class TestPerSourceScope(ScanScopeTestCase):
"""``roots``: scan only the selected effective sources (#249 S4 Task 10)."""
async def asyncSetUp(self):
await super().asyncSetUp()
self.a = self.output_dir / "a"
self.b = self.output_dir / "b"
self.a.mkdir()
self.b.mkdir()
_write_png(self.a / "in_a.png", prompt_text="a prompt in a")
_write_png(self.b / "in_b.png", prompt_text="a prompt in b")
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.a), self._root(self.b)]
async def _refused(self, body):
resp = await self.client.request("POST", "/prompt_manager/jobs/scan", json=body)
data = await resp.json()
self.assertEqual(resp.status, 400, data)
self.assertFalse(data["success"])
self.assertNotIn(self.tmpdir, json.dumps(data))
self.assertIsNone(self.api._scan_job)
self.assertEqual(self.api._jobs.list(), [])
return data
async def test_scoped_scan_only_walks_selected_root(self):
walked = []
original = self.api._collect_output_media_sync
def spy(output_dirs):
walked.append([os.path.basename(str(d)) for d in output_dirs])
return original(output_dirs)
self.api._collect_output_media_sync = spy
reads = self._spy_reads()
job = await self._scan({"scope": "all", "roots": [1]})
self.assertEqual(walked, [["b"]])
self.assertEqual(reads, ["in_b.png"])
self.assertEqual(
job["scope"], {"files": "all", "roots": [1], "output_folders": 1}
)
self.assertEqual(
job["outcomes"]["added"],
[
{
"path": "in_b.png",
"root": "output/b",
"root_index": 1,
"detail": job["outcomes"]["added"][0]["detail"],
}
],
)
self.assertEqual(job["counters"]["added"], 1)
self.assertEqual(self.api._scan_job.last_event["directories"], ["output/b"])
async def test_duplicate_roots_are_one_selection(self):
job = await self._scan({"roots": [1, 0, 1]})
self.assertEqual(job["scope"]["roots"], [0, 1])
self.assertEqual(job["scope"]["output_folders"], 2)
self.assertEqual(job["counters"]["added"], 2)
async def test_scope_line_counts_selected_folders(self):
job = await self._scan({"roots": [0]})
self.assertTrue(
job["log"][0].endswith(
"Scan started · 1 output folder · new and unresolved files"
),
job["log"][0],
)
self.assertEqual(job["scope"]["output_folders"], 1)
unscoped = await self._scan({"scope": "all"})
self.assertEqual(
unscoped["scope"], {"files": "all", "roots": None, "output_folders": 2}
)
self.assertIn("Scan started · 2 output folders · all files", unscoped["log"][0])
async def test_invalid_root_index_is_400_without_job(self):
data = await self._refused({"roots": [2]})
self.assertEqual(data["code"], "unknown-source")
data = await self._refused({"roots": [0, 64]})
self.assertEqual(data["code"], "unknown-source")
async def test_bool_and_float_roots_rejected(self):
for roots in (
[True],
[False],
[0.0],
[1.5],
["0"],
[None],
[-1],
[0, None],
[],
None,
0,
"0",
{"0": 0},
list(range(65)),
):
with self.subTest(roots=roots):
await self._refused({"roots": roots})
async def test_legacy_scan_unscoped_unchanged(self):
from py.api.admin import SCAN_COUNT_KEYS
resp = await self.client.request("POST", "/prompt_manager/scan", data="{}")
payloads = _sse_payloads(await resp.text())
final = payloads[-1]
self.assertEqual(set(final), set(SCAN_COUNT_KEYS) | {"type", "directories"})
self.assertEqual(final["directories"], ["output/a", "output/b"])
self.assertEqual(final["added"], 2)
self.assertEqual(
self.api._scan_job.scope,
{"files": "all", "roots": None, "output_folders": 2},
)
async def test_legacy_scan_attaches_to_a_scoped_job_as_is(self):
import threading
from py.api.admin import SCAN_COUNT_KEYS
gate = threading.Event()
original = self.api._collect_output_media_sync
def gated(output_dirs):
gate.wait(5)
return original(output_dirs)
self.api._collect_output_media_sync = gated
resp = await self.client.request(
"POST", "/prompt_manager/jobs/scan", json={"roots": [0]}
)
self.assertEqual(resp.status, 202)
started = (await resp.json())["job"]
legacy = asyncio.ensure_future(
self.client.request("POST", "/prompt_manager/scan", data="{}")
)
await asyncio.sleep(0.05)
gate.set()
payloads = _sse_payloads(await (await legacy).text())
await asyncio.wait_for(self.api._scan_job.task, 10)
self.assertEqual(self.api._scan_job.id, started["id"])
self.assertEqual(len(self.api._jobs.list()), 1)
final = payloads[-1]
self.assertEqual(set(final), set(SCAN_COUNT_KEYS) | {"type", "directories"})
self.assertEqual(final["directories"], ["output/a"])
self.assertEqual(final["added"], 1)
async def test_scoped_scan_uses_the_sources_it_was_started_with(self):
"""Sources saved after the start never change what the job walks."""
resolved = []
original = self.api._get_all_output_dirs
def counting():
resolved.append(1)
return original()
self.api._get_all_output_dirs = counting
job = await self._scan({"roots": [1]})
self.assertEqual(len(resolved), 1) # the request's own resolution only
self.assertEqual(job["outcomes"]["added"][0]["root_index"], 1)
class TestKnownResolvedPaths(ScanScopeTestCase):
"""``PromptDatabase.get_known_resolved_paths``: one streamed query."""