diff --git a/py/api/admin.py b/py/api/admin.py index dc4e9bcb..5b530ef9 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -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( diff --git a/py/api/job_routes.py b/py/api/job_routes.py index f8556601..b0c6b013 100644 --- a/py/api/job_routes.py +++ b/py/api/job_routes.py @@ -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) diff --git a/tests/test_scan_scope.py b/tests/test_scan_scope.py index 60447149..fbd48b9e 100644 --- a/tests/test_scan_scope.py +++ b/tests/test_scan_scope.py @@ -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."""