From 4d66996569640aa7aa8e7ef0865d4bd24bc4d92e Mon Sep 17 00:00:00 2001 From: Nick Sweeting Date: Mon, 6 Apr 2026 23:47:38 -0700 Subject: [PATCH] small fixes --- archivebox/cli/archivebox_version.py | 94 ++++++- archivebox/core/settings.py | 17 +- archivebox/hooks.py | 21 +- archivebox/machine/models.py | 8 +- archivebox/services/runner.py | 244 +++++++++++++----- archivebox/services/snapshot_service.py | 107 +++++--- .../tests/test_archive_result_service.py | 4 +- archivebox/tests/test_hooks.py | 69 +++++ archivebox/tests/test_machine_models.py | 16 ++ archivebox/tests/test_runner.py | 138 ++++++++++ archivebox/tests/test_schedule_e2e.py | 1 + 11 files changed, 605 insertions(+), 114 deletions(-) diff --git a/archivebox/cli/archivebox_version.py b/archivebox/cli/archivebox_version.py index 8d61c64c..cd088079 100755 --- a/archivebox/cli/archivebox_version.py +++ b/archivebox/cli/archivebox_version.py @@ -13,6 +13,79 @@ import rich_click as click from archivebox.misc.util import docstring, enforce_types +def _format_binary_abspath( + abspath: str, + *, + pwd: Path, + lib_dir: Path, + personas_dir: Path, + home: Path, +) -> str: + path = Path(abspath).expanduser() + try: + normalized = path.resolve(strict=False) + except Exception: + normalized = path + + machine = platform.machine().lower() + system = platform.system().lower() + arch_scope = f"{machine}-{system}" + + candidate_bases: tuple[tuple[Path, str], ...] = ( + (pwd, "./"), + (lib_dir, "LIB_DIR/"), + (Path(os.environ.get("LIB_DIR", "")), "LIB_DIR/") if os.environ.get("LIB_DIR") else (Path(), ""), + (personas_dir, "PERSONAS_DIR/"), + (Path(os.environ.get("PERSONAS_DIR", "")), "PERSONAS_DIR/") if os.environ.get("PERSONAS_DIR") else (Path(), ""), + (home / ".config" / "abx" / "lib" / arch_scope, "LIB_DIR/"), + (home / ".config" / "abx" / "lib", "LIB_DIR/"), + (home / ".config" / "abx" / "personas", "PERSONAS_DIR/"), + (home, "~/"), + ) + + for base, prefix in candidate_bases: + if not prefix: + continue + for candidate in (base, base.resolve(strict=False)): + try: + relative = normalized.relative_to(candidate) + except ValueError: + continue + + relative_str = relative.as_posix() + if prefix == "./": + return "." if not relative_str else f"./{relative_str}" + if prefix == "LIB_DIR/": + return "LIB_DIR" if not relative_str else f"LIB_DIR/{relative_str}" + if prefix == "PERSONAS_DIR/": + return "PERSONAS_DIR" if not relative_str else f"PERSONAS_DIR/{relative_str}" + return "~" if not relative_str else f"~/{relative_str}" + + return normalized.as_posix() + + +def _render_binary_abspath(abspath: str): + from rich.text import Text + + if abspath.startswith("LIB_DIR/"): + return Text.assemble(("LIB_DIR", "bright_blue"), (abspath.removeprefix("LIB_DIR"), "green")) + if abspath == "LIB_DIR": + return Text("LIB_DIR", style="bright_blue") + if abspath.startswith("PERSONAS_DIR/"): + return Text.assemble(("PERSONAS_DIR", "medium_purple"), (abspath.removeprefix("PERSONAS_DIR"), "green")) + if abspath == "PERSONAS_DIR": + return Text("PERSONAS_DIR", style="medium_purple") + if abspath.startswith("~/"): + return Text.assemble(("~", "cyan"), (abspath.removeprefix("~"), "green")) + if abspath == "~": + return Text("~", style="cyan") + if abspath.startswith("./"): + return Text.assemble((".", "cyan"), (abspath.removeprefix("."), "green")) + if abspath == ".": + return Text(".", style="cyan") + return Text(abspath, style="green") + + @enforce_types def version( quiet: bool = False, @@ -30,7 +103,7 @@ def version( from rich.panel import Panel from rich.console import Console - from archivebox.config import CONSTANTS, DATA_DIR + from archivebox.config import CONSTANTS from archivebox.config.version import get_COMMIT_HASH, get_BUILD_TIME from archivebox.config.permissions import ARCHIVEBOX_USER, ARCHIVEBOX_GROUP, RUNNING_AS_UID, RUNNING_AS_GID, IN_DOCKER from archivebox.config.paths import get_data_locations, get_code_locations @@ -141,23 +214,36 @@ def version( prnt("", "[grey53]No binaries detected. Run [green]archivebox install[/green] to detect dependencies.[/grey53]") else: any_available = False + compact_paths = console.is_terminal for name in all_binary_names: if requested_names and name not in requested_names: continue installed = db_binaries.get(name) if installed and installed.is_valid: - display_path = installed.abspath.replace(str(DATA_DIR), ".").replace(str(Path("~").expanduser()), "~") + display_name = Path(name).expanduser().name if ("/" in name or name.startswith("~")) else name + display_path = ( + _format_binary_abspath( + installed.abspath, + pwd=Path.cwd(), + lib_dir=STORAGE_CONFIG.LIB_DIR, + personas_dir=Path.home() / ".config" / "abx" / "personas", + home=Path.home(), + ) + if compact_paths + else installed.abspath + ) + rendered_path = _render_binary_abspath(display_path) if compact_paths else display_path version_str = (installed.version or "unknown")[:15] provider = (installed.binprovider or "env")[:8] prnt( "", "[green]√[/green]", "", - name.ljust(18), + display_name.ljust(18), version_str.ljust(16), provider.ljust(8), - display_path, + rendered_path, overflow="ignore", crop=False, ) diff --git a/archivebox/core/settings.py b/archivebox/core/settings.py index 966909c1..d3f8ef02 100644 --- a/archivebox/core/settings.py +++ b/archivebox/core/settings.py @@ -557,12 +557,17 @@ if DEBUG_TOOLBAR: MIDDLEWARE = [*MIDDLEWARE, "debug_toolbar.middleware.DebugToolbarMiddleware"] if DEBUG: - INSTALLED_APPS += ["django_autotyping"] - AUTOTYPING = { - "STUBS_GENERATION": { - "LOCAL_STUBS_DIR": PACKAGE_DIR / "typings", - }, - } + try: + import django_autotyping # noqa + except ImportError: + pass + else: + INSTALLED_APPS += ["django_autotyping"] + AUTOTYPING = { + "STUBS_GENERATION": { + "LOCAL_STUBS_DIR": PACKAGE_DIR / "typings", + }, + } # https://github.com/bensi94/Django-Requests-Tracker (improved version of django-debug-toolbar) # Must delete archivebox/templates/admin to use because it relies on some things we override diff --git a/archivebox/hooks.py b/archivebox/hooks.py index 9817e268..56a33b9b 100644 --- a/archivebox/hooks.py +++ b/archivebox/hooks.py @@ -863,7 +863,7 @@ def get_config_defaults_from_plugins() -> dict[str, Any]: return defaults -def get_plugin_special_config(plugin_name: str, config: dict[str, Any]) -> dict[str, Any]: +def get_plugin_special_config(plugin_name: str, config: dict[str, Any], _visited: set[str] | None = None) -> dict[str, Any]: """ Extract special config keys for a plugin following naming conventions. @@ -945,6 +945,25 @@ def get_plugin_special_config(plugin_name: str, config: dict[str, Any]) -> dict[ # Handle string values from config file ("true"/"false") enabled = enabled.lower() not in ("false", "0", "no", "") + plugin_configs = discover_plugin_configs() + plugin_name_lower = plugin_name.lower() + + if enabled: + visited = _visited or set() + if plugin_name_lower not in visited: + next_visited = visited | {plugin_name_lower} + schema = plugin_configs.get(plugin_name_lower, {}) + required_plugins = schema.get("required_plugins", []) + if isinstance(required_plugins, list): + for required_plugin in required_plugins: + required_plugin_name = str(required_plugin).strip() + if not required_plugin_name: + continue + required_config = get_plugin_special_config(required_plugin_name, config, _visited=next_visited) + if not required_config["enabled"]: + enabled = False + break + # 2. Timeout: PLUGINNAME_TIMEOUT (fallback to TIMEOUT, default 300) timeout_key = f"{plugin_upper}_TIMEOUT" timeout = config.get(timeout_key) or config.get("TIMEOUT", 300) diff --git a/archivebox/machine/models.py b/archivebox/machine/models.py index 1ffae6c9..10444b36 100755 --- a/archivebox/machine/models.py +++ b/archivebox/machine/models.py @@ -172,8 +172,12 @@ class Machine(ModelWithHealthStats): global _CURRENT_MACHINE if _CURRENT_MACHINE: if timezone.now() < _CURRENT_MACHINE.modified_at + timedelta(seconds=MACHINE_RECHECK_INTERVAL): - return cls._sanitize_config(_CURRENT_MACHINE) - _CURRENT_MACHINE = None + if not cls.objects.filter(id=_CURRENT_MACHINE.id).exists(): + _CURRENT_MACHINE = None + else: + return cls._sanitize_config(_CURRENT_MACHINE) + else: + _CURRENT_MACHINE = None _CURRENT_MACHINE, _ = cls.objects.update_or_create( guid=get_host_guid(), defaults={"hostname": socket.gethostname(), **get_os_info(), **get_vm_info(), "stats": get_host_stats()}, diff --git a/archivebox/services/runner.py b/archivebox/services/runner.py index 352ca05a..b142905a 100644 --- a/archivebox/services/runner.py +++ b/archivebox/services/runner.py @@ -16,10 +16,10 @@ from asgiref.sync import sync_to_async from django.utils import timezone from rich.console import Console -from abx_dl.events import BinaryRequestEvent +from abx_dl.events import BinaryRequestEvent, MachineEvent from abx_dl.heartbeat import CrawlHeartbeat from abx_dl.limits import CrawlLimitState -from abx_dl.models import Plugin, discover_plugins, filter_plugins +from abx_dl.models import Plugin, Snapshot as AbxSnapshot, discover_plugins, filter_plugins from abx_dl.orchestrator import ( create_bus, download, @@ -30,6 +30,7 @@ from abx_dl.orchestrator import ( from .archive_result_service import ArchiveResultService from .binary_service import BinaryService from .crawl_service import CrawlService +from .machine_service import MachineService from .process_service import ProcessService from .snapshot_service import SnapshotService from .tag_service import TagService @@ -46,6 +47,33 @@ def _count_selected_hooks(plugins: dict[str, Plugin], selected_plugins: list[str return sum(1 for plugin in selected.values() for hook in plugin.hooks if "CrawlSetup" in hook.name or "Snapshot" in hook.name) +def _normalize_runtime_config(config: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in json.loads(json.dumps(config, default=str)).items() if value is not None} + + +async def _emit_machine_config( + bus, + *, + config: dict[str, Any], + derived_config: dict[str, Any], +) -> None: + user_config = _normalize_runtime_config(config) + derived_machine_config = _normalize_runtime_config(derived_config) + await bus.emit( + MachineEvent( + config=user_config, + config_type="user", + ), + ) + if derived_machine_config: + await bus.emit( + MachineEvent( + config=derived_machine_config, + config_type="derived", + ), + ) + + def ensure_background_runner(*, allow_under_pytest: bool = False) -> bool: if os.environ.get("PYTEST_CURRENT_TEST") and not allow_under_pytest: return False @@ -99,6 +127,7 @@ class CrawlRunner: BinaryService(self.bus) TagService(self.bus) CrawlService(self.bus, crawl_id=str(crawl.id)) + MachineService(self.bus) self.process_discovered_snapshots_inline = process_discovered_snapshots_inline async def ignore_snapshot(_snapshot_id: str) -> None: @@ -118,11 +147,12 @@ class CrawlRunner: self.base_config: dict[str, Any] = {} self.derived_config: dict[str, Any] = {} self.primary_url = "" + self.crawl_output_dir = "" self._live_stream = None async def run(self) -> None: heartbeat = CrawlHeartbeat( - Path(self.crawl.output_dir), + Path(self.crawl_output_dir), runtime="archivebox", crawl_id=str(self.crawl.id), ) @@ -131,17 +161,35 @@ class CrawlRunner: live_ui = self._create_live_ui() with live_ui if live_ui is not None else nullcontext(): await heartbeat.start() - setup_abx_services( + if snapshot_ids: + root_snapshot = await sync_to_async(self.load_snapshot_payload, thread_sensitive=True)(snapshot_ids[0]) + setup_abx_services( + self.bus, + plugins=self.plugins, + url=root_snapshot["url"], + snapshot=AbxSnapshot( + id=root_snapshot["id"], + url=root_snapshot["url"], + depth=int(root_snapshot["depth"]), + crawl_id=str(self.crawl.id), + ), + output_dir=Path(root_snapshot["output_dir"]), + install_enabled=False, + crawl_setup_enabled=False, + crawl_start_enabled=False, + snapshot_cleanup_enabled=False, + crawl_cleanup_enabled=False, + persist_derived=False, + auto_install=True, + emit_jsonl=False, + ) + await _emit_machine_config( self.bus, - plugins=self.plugins, - config_overrides={ + config={ **self.base_config, "ABX_RUNTIME": "archivebox", }, - derived_config_overrides=self.derived_config, - persist_derived=False, - auto_install=True, - emit_jsonl=False, + derived_config=self.derived_config, ) if snapshot_ids: root_snapshot_id = snapshot_ids[0] @@ -186,6 +234,7 @@ class CrawlRunner: def load_run_state(self) -> list[str]: from archivebox.config.configset import get_config + from archivebox.hooks import discover_hooks from archivebox.machine.models import Machine, NetworkInterface, Process self.primary_url = self.crawl.get_urls_list()[0] if self.crawl.get_urls_list() else "" @@ -198,10 +247,18 @@ class CrawlRunner: self.persona = self.crawl.resolve_persona() self.base_config = get_config(crawl=self.crawl) self.derived_config = dict(Machine.current().config) + self.crawl_output_dir = str(self.crawl.output_dir) self.base_config["ABX_RUNTIME"] = "archivebox" if self.selected_plugins is None: - raw_plugins = self.base_config["PLUGINS"].strip() - self.selected_plugins = [name.strip() for name in raw_plugins.split(",") if name.strip()] if raw_plugins else None + raw_plugins = str(self.base_config.get("PLUGINS") or "").strip() + if raw_plugins: + self.selected_plugins = [name.strip() for name in raw_plugins.split(",") if name.strip()] + else: + runtime_events = ("CrawlSetup", "CrawlCleanup", "Snapshot", "SnapshotCleanup") + runtime_plugins = { + hook.parent.name for event_name in runtime_events for hook in discover_hooks(event_name, config=self.base_config) + } + self.selected_plugins = sorted(runtime_plugins) or None if self.persona: self.base_config.update( self.persona.prepare_runtime_for_crawl( @@ -275,7 +332,7 @@ class CrawlRunner: ) live_ui.print_intro( url=self.primary_url or "crawl", - output_dir=Path(self.crawl.output_dir), + output_dir=Path(self.crawl_output_dir), plugins_label=plugins_label, ) return live_ui @@ -287,7 +344,7 @@ class CrawlRunner: snapshot = Snapshot.objects.select_related("crawl").get(id=snapshot_id) config = get_config(crawl=self.crawl, snapshot=snapshot) config.update(self.base_config) - config["CRAWL_DIR"] = str(self.crawl.output_dir) + config["CRAWL_DIR"] = self.crawl_output_dir config["SNAP_DIR"] = str(snapshot.output_dir) extra_context: dict[str, Any] = {} if config.get("EXTRA_CONTEXT"): @@ -310,8 +367,62 @@ class CrawlRunner: "status": snapshot.status, "output_dir": str(snapshot.output_dir), "config": config, + "_snapshot": snapshot, } + async def enqueue_discovered_snapshots_from_outputs(self, snapshot_payload: dict[str, Any]) -> None: + from archivebox.core.models import Snapshot + from archivebox.hooks import collect_urls_from_plugins + + if int(snapshot_payload["depth"]) >= self.crawl.max_depth: + return + if CrawlLimitState.from_config(snapshot_payload["config"]).get_stop_reason() == "max_size": + return + + discovered_urls = await sync_to_async(collect_urls_from_plugins, thread_sensitive=True)(Path(snapshot_payload["output_dir"])) + if not discovered_urls: + return + + parent_snapshot = snapshot_payload.get("_snapshot") + if parent_snapshot is None: + parent_snapshot = await sync_to_async( + lambda: Snapshot.objects.select_related("crawl", "crawl__created_by").filter(id=snapshot_payload["id"]).first(), + thread_sensitive=True, + )() + if parent_snapshot is None: + return + + for record in discovered_urls: + url = str(record.get("url") or "").strip() + if not url: + continue + passes_filters = await sync_to_async(self.crawl.url_passes_filters, thread_sensitive=True)(url, snapshot=parent_snapshot) + if not passes_filters: + continue + child_snapshot = await sync_to_async(Snapshot.from_json, thread_sensitive=True)( + { + "url": url, + "depth": parent_snapshot.depth + 1, + "title": str(record.get("title") or "").strip(), + "tags": str(record.get("tags") or "").strip(), + "parent_snapshot_id": str(parent_snapshot.id), + "crawl_id": str(self.crawl.id), + }, + overrides={ + "crawl": self.crawl, + "snapshot": parent_snapshot, + "created_by_id": self.crawl.created_by_id, + }, + queue_for_extraction=False, + ) + if child_snapshot is None or child_snapshot.status == child_snapshot.StatusChoices.SEALED: + continue + child_snapshot.status = child_snapshot.StatusChoices.QUEUED + child_snapshot.retry_at = timezone.now() + await child_snapshot.asave(update_fields=["status", "retry_at", "modified_at"]) + if self.process_discovered_snapshots_inline: + await self.enqueue_snapshot(str(child_snapshot.id)) + async def run_crawl_setup(self, snapshot_id: str) -> None: snapshot = await sync_to_async(self.load_snapshot_payload, thread_sensitive=True)(snapshot_id) await download( @@ -319,8 +430,8 @@ class CrawlRunner: plugins=self.plugins, output_dir=Path(snapshot["output_dir"]), selected_plugins=self.selected_plugins, - config_overrides=snapshot["config"], - derived_config_overrides=self.derived_config, + config_overrides=_normalize_runtime_config(snapshot["config"]), + derived_config_overrides=_normalize_runtime_config(self.derived_config), bus=self.bus, emit_jsonl=False, install_enabled=True, @@ -328,11 +439,11 @@ class CrawlRunner: crawl_start_enabled=False, snapshot_cleanup_enabled=False, crawl_cleanup_enabled=False, - machine_service=None, - binary_service=None, - process_service=None, - archive_result_service=None, - tag_service=None, + MachineService=None, + BinaryService=None, + ProcessService=None, + ArchiveResultService=None, + TagService=None, ) async def run_crawl_cleanup(self, snapshot_id: str) -> None: @@ -343,19 +454,19 @@ class CrawlRunner: output_dir=Path(snapshot["output_dir"]), plugins=self.plugins, selected_plugins=self.selected_plugins, - config_overrides=snapshot["config"], - derived_config_overrides=self.derived_config, + config_overrides=_normalize_runtime_config(snapshot["config"]), + derived_config_overrides=_normalize_runtime_config(self.derived_config), emit_jsonl=False, install_enabled=False, crawl_setup_enabled=False, crawl_start_enabled=False, snapshot_cleanup_enabled=False, crawl_cleanup_enabled=True, - machine_service=None, - binary_service=None, - process_service=None, - archive_result_service=None, - tag_service=None, + MachineService=None, + BinaryService=None, + ProcessService=None, + ArchiveResultService=None, + TagService=None, ) async def run_snapshot(self, snapshot_id: str) -> None: @@ -372,8 +483,8 @@ class CrawlRunner: plugins=self.plugins, output_dir=Path(snapshot["output_dir"]), selected_plugins=self.selected_plugins, - config_overrides=snapshot["config"], - derived_config_overrides=self.derived_config, + config_overrides=_normalize_runtime_config(snapshot["config"]), + derived_config_overrides=_normalize_runtime_config(self.derived_config), bus=self.bus, emit_jsonl=False, install_enabled=False, @@ -381,12 +492,13 @@ class CrawlRunner: crawl_start_enabled=True, snapshot_cleanup_enabled=True, crawl_cleanup_enabled=False, - machine_service=None, - binary_service=None, - process_service=None, - archive_result_service=None, - tag_service=None, + MachineService=None, + BinaryService=None, + ProcessService=None, + ArchiveResultService=None, + TagService=None, ) + await self.enqueue_discovered_snapshots_from_outputs(snapshot) finally: current_task = asyncio.current_task() if current_task is not None and self.snapshot_tasks.get(snapshot_id) is current_task: @@ -431,22 +543,28 @@ async def _run_binary(binary_id: str) -> None: plugins = discover_plugins() config = get_config() machine = await sync_to_async(Machine.current, thread_sensitive=True)() - derived_config = dict(machine.config) + derived_config = _normalize_runtime_config(dict(machine.config)) config["ABX_RUNTIME"] = "archivebox" + config = _normalize_runtime_config(config) bus = create_bus(name=_bus_name("ArchiveBox_binary", str(binary.id)), total_timeout=1800.0) ProcessService(bus) BinaryService(bus) TagService(bus) ArchiveResultService(bus) + MachineService(bus) setup_abx_services( bus, plugins=plugins, - config_overrides=config, - derived_config_overrides=derived_config, + install_enabled=False, + crawl_setup_enabled=False, + crawl_start_enabled=False, + snapshot_cleanup_enabled=False, + crawl_cleanup_enabled=False, persist_derived=False, auto_install=True, emit_jsonl=False, ) + await _emit_machine_config(bus, config=config, derived_config=derived_config) try: await bus.emit( @@ -476,22 +594,16 @@ async def _run_install(plugin_names: list[str] | None = None) -> None: plugins = discover_plugins() config = get_config() machine = await sync_to_async(Machine.current, thread_sensitive=True)() - derived_config = dict(machine.config) + derived_config = _normalize_runtime_config(dict(machine.config)) config["ABX_RUNTIME"] = "archivebox" + config = _normalize_runtime_config(config) bus = create_bus(name="ArchiveBox_install", total_timeout=3600.0) ProcessService(bus) BinaryService(bus) TagService(bus) ArchiveResultService(bus) - setup_abx_services( - bus, - plugins=plugins, - config_overrides=config, - derived_config_overrides=derived_config, - persist_derived=False, - auto_install=True, - emit_jsonl=False, - ) + MachineService(bus) + await _emit_machine_config(bus, config=config, derived_config=derived_config) live_stream = None try: @@ -549,7 +661,7 @@ async def _run_install(plugin_names: list[str] | None = None) -> None: plugins_label=plugins_label, ) with live_ui if live_ui is not None else nullcontext(): - results = await abx_install_plugins( + await abx_install_plugins( plugin_names=plugin_names, plugins=plugins, output_dir=output_dir, @@ -557,12 +669,10 @@ async def _run_install(plugin_names: list[str] | None = None) -> None: derived_config_overrides=derived_config, emit_jsonl=False, bus=bus, - machine_service=None, - binary_service=None, - process_service=None, + MachineService=None, ) if live_ui is not None: - live_ui.print_summary(results, output_dir=output_dir) + live_ui.print_summary(output_dir=output_dir) finally: await bus.stop() try: @@ -713,19 +823,6 @@ def run_pending_crawls(*, daemon: bool = False, crawl_id: str | None = None) -> if schedule.is_due(now): schedule.enqueue(queued_at=now) - if crawl_id is None: - binary = ( - Binary.objects.filter(retry_at__lte=timezone.now()) - .exclude(status=Binary.StatusChoices.INSTALLED) - .order_by("retry_at", "created_at") - .first() - ) - if binary is not None: - if not binary.claim_processing_lock(lock_seconds=60): - continue - run_binary(str(binary.id)) - continue - queued_crawls = Crawl.objects.filter( retry_at__lte=timezone.now(), status=Crawl.StatusChoices.QUEUED, @@ -759,6 +856,21 @@ def run_pending_crawls(*, daemon: bool = False, crawl_id: str | None = None) -> ) continue + if crawl_id is None: + # Standalone binary backlog should not starve queued crawls or snapshots. + # Crawl.run() already claims and installs crawl-declared Binary rows as needed. + binary = ( + Binary.objects.filter(retry_at__lte=timezone.now()) + .exclude(status=Binary.StatusChoices.INSTALLED) + .order_by("retry_at", "created_at") + .first() + ) + if binary is not None: + if not binary.claim_processing_lock(lock_seconds=60): + continue + run_binary(str(binary.id)) + continue + pending = Crawl.objects.filter( retry_at__lte=timezone.now(), status=Crawl.StatusChoices.STARTED, diff --git a/archivebox/services/snapshot_service.py b/archivebox/services/snapshot_service.py index f84632ba..a82b2d74 100644 --- a/archivebox/services/snapshot_service.py +++ b/archivebox/services/snapshot_service.py @@ -1,5 +1,7 @@ from __future__ import annotations +from pathlib import Path + from asgiref.sync import sync_to_async from django.utils import timezone @@ -19,6 +21,43 @@ class SnapshotService(BaseService): self.bus.on(SnapshotEvent, self.on_SnapshotEvent) self.bus.on(SnapshotCompletedEvent, self.on_SnapshotCompletedEvent) + async def _upsert_discovered_snapshot(self, parent_snapshot, *, url: str, depth: int, title: str = "", tags: str = "") -> str | None: + from archivebox.core.models import Snapshot + + crawl = parent_snapshot.crawl + if depth > crawl.max_depth: + return None + stop_reason = await sync_to_async(self._crawl_limit_stop_reason, thread_sensitive=True)(crawl) + if stop_reason == "max_size": + return None + passes_filters = await sync_to_async(crawl.url_passes_filters, thread_sensitive=True)(url, snapshot=parent_snapshot) + if not passes_filters: + return None + + snapshot = await sync_to_async(Snapshot.from_json, thread_sensitive=True)( + { + "url": url, + "depth": depth, + "title": title, + "tags": tags, + "parent_snapshot_id": str(parent_snapshot.id), + "crawl_id": str(crawl.id), + }, + overrides={ + "crawl": crawl, + "snapshot": parent_snapshot, + "created_by_id": crawl.created_by_id, + }, + queue_for_extraction=False, + ) + if snapshot is None or snapshot.status == Snapshot.StatusChoices.SEALED: + return None + + snapshot.status = Snapshot.StatusChoices.QUEUED + snapshot.retry_at = timezone.now() + await snapshot.asave(update_fields=["status", "retry_at", "modified_at"]) + return str(snapshot.id) + async def on_SnapshotEvent(self, event: SnapshotEvent) -> None: from archivebox.core.models import Snapshot from archivebox.crawls.models import Crawl @@ -33,36 +72,25 @@ class SnapshotService(BaseService): await snapshot.asave(update_fields=["status", "retry_at", "modified_at"]) snapshot_id = str(snapshot.id) elif event.depth > 0: - if event.depth <= crawl.max_depth and self._crawl_limit_stop_reason(crawl) != "max_size": - parent_event = await self.bus.find( - SnapshotEvent, - past=True, - future=False, - where=lambda candidate: candidate.depth == event.depth - 1 and self.bus.event_is_child_of(event, candidate), + parent_event = await self.bus.find( + SnapshotEvent, + past=True, + future=False, + where=lambda candidate: candidate.depth == event.depth - 1 and self.bus.event_is_child_of(event, candidate), + ) + parent_snapshot = None + if parent_event is not None: + parent_snapshot = ( + await Snapshot.objects.select_related("crawl", "crawl__created_by") + .filter(id=parent_event.snapshot_id, crawl=crawl) + .afirst() + ) + if parent_snapshot is not None: + snapshot_id = await self._upsert_discovered_snapshot( + parent_snapshot, + url=event.url, + depth=event.depth, ) - parent_snapshot = None - if parent_event is not None: - parent_snapshot = await Snapshot.objects.filter(id=parent_event.snapshot_id, crawl=crawl).afirst() - if parent_snapshot is not None and self._url_passes_filters(crawl, parent_snapshot, event.url): - snapshot = await sync_to_async(Snapshot.from_json, thread_sensitive=True)( - { - "url": event.url, - "depth": event.depth, - "parent_snapshot_id": str(parent_snapshot.id), - "crawl_id": str(crawl.id), - }, - overrides={ - "crawl": crawl, - "snapshot": parent_snapshot, - "created_by_id": crawl.created_by_id, - }, - queue_for_extraction=False, - ) - if snapshot is not None and snapshot.status != Snapshot.StatusChoices.SEALED: - snapshot.retry_at = None - snapshot.status = Snapshot.StatusChoices.QUEUED - await snapshot.asave(update_fields=["status", "retry_at", "modified_at"]) - snapshot_id = str(snapshot.id) if snapshot_id: snapshot = await Snapshot.objects.filter(id=snapshot_id).select_related("crawl", "crawl__created_by").afirst() @@ -74,14 +102,15 @@ class SnapshotService(BaseService): async def on_SnapshotCompletedEvent(self, event: SnapshotCompletedEvent) -> None: from archivebox.core.models import Snapshot - snapshot = await Snapshot.objects.select_related("crawl").filter(id=event.snapshot_id).afirst() + snapshot = await Snapshot.objects.select_related("crawl", "crawl__created_by").filter(id=event.snapshot_id).afirst() snapshot_id: str | None = None if snapshot is not None: snapshot.status = Snapshot.StatusChoices.SEALED snapshot.retry_at = None snapshot.downloaded_at = snapshot.downloaded_at or timezone.now() await snapshot.asave(update_fields=["status", "retry_at", "downloaded_at", "modified_at"]) - if snapshot.crawl_id and self._crawl_limit_stop_reason(snapshot.crawl) == "max_size": + stop_reason = await sync_to_async(self._crawl_limit_stop_reason, thread_sensitive=True)(snapshot.crawl) + if snapshot.crawl_id and stop_reason == "max_size": await ( Snapshot.objects.filter( crawl_id=snapshot.crawl_id, @@ -101,9 +130,21 @@ class SnapshotService(BaseService): await sync_to_async(snapshot.write_index_jsonl, thread_sensitive=True)() await sync_to_async(snapshot.write_json_details, thread_sensitive=True)() await sync_to_async(snapshot.write_html_details, thread_sensitive=True)() + stop_reason = await sync_to_async(self._crawl_limit_stop_reason, thread_sensitive=True)(snapshot.crawl) + if snapshot.depth < snapshot.crawl.max_depth and stop_reason != "max_size": + from archivebox.hooks import collect_urls_from_plugins - def _url_passes_filters(self, crawl, parent_snapshot, url: str) -> bool: - return crawl.url_passes_filters(url, snapshot=parent_snapshot) + discovered_urls = await sync_to_async(collect_urls_from_plugins, thread_sensitive=True)(Path(snapshot.output_dir)) + for record in discovered_urls: + discovered_snapshot_id = await self._upsert_discovered_snapshot( + snapshot, + url=str(record.get("url") or "").strip(), + depth=snapshot.depth + 1, + title=str(record.get("title") or "").strip(), + tags=str(record.get("tags") or "").strip(), + ) + if discovered_snapshot_id: + await self.schedule_snapshot(discovered_snapshot_id) def _crawl_limit_stop_reason(self, crawl) -> str: config = dict(crawl.config or {}) diff --git a/archivebox/tests/test_archive_result_service.py b/archivebox/tests/test_archive_result_service.py index bae88398..14ad8168 100644 --- a/archivebox/tests/test_archive_result_service.py +++ b/archivebox/tests/test_archive_result_service.py @@ -366,7 +366,7 @@ def test_process_started_hydrates_binary_and_iface_from_existing_binary_records( output_dir.mkdir() bus = create_bus(name="test_process_started_binary_hydration") - DlProcessService(bus, emit_jsonl=False, stderr_is_tty=False) + DlProcessService(bus, emit_jsonl=False, interactive_tty=False) ArchiveBoxProcessService(bus) async def run_test() -> None: @@ -435,7 +435,7 @@ def test_process_started_uses_node_binary_for_js_hooks_without_plugin_binary(mon output_dir.mkdir() bus = create_bus(name="test_process_started_node_fallback") - DlProcessService(bus, emit_jsonl=False, stderr_is_tty=False) + DlProcessService(bus, emit_jsonl=False, interactive_tty=False) ArchiveBoxProcessService(bus) async def run_test() -> None: diff --git a/archivebox/tests/test_hooks.py b/archivebox/tests/test_hooks.py index d0bd8f83..e00126e6 100755 --- a/archivebox/tests/test_hooks.py +++ b/archivebox/tests/test_hooks.py @@ -238,6 +238,75 @@ class TestHookDiscovery(unittest.TestCase): self.assertEqual(hooks_module.normalize_hook_event_name("SnapshotCleanupEvent"), "SnapshotCleanup") self.assertEqual(hooks_module.normalize_hook_event_name("CrawlCleanupEvent"), "CrawlCleanup") + def test_discover_hooks_skips_plugins_with_disabled_required_dependencies(self): + """Plugins whose required_plugins are disabled should not run.""" + from archivebox import hooks as hooks_module + + chrome_dir = self.plugins_dir / "chrome" + chrome_dir.mkdir(exist_ok=True) + (chrome_dir / "config.json").write_text( + json.dumps( + { + "type": "object", + "required_plugins": [], + "properties": { + "CHROME_ENABLED": { + "type": "boolean", + "default": True, + "x-aliases": ["USE_CHROME"], + }, + }, + }, + ), + ) + (chrome_dir / "on_Snapshot__20_chrome.js").write_text("// chrome hook") + + accessibility_dir = self.plugins_dir / "accessibility" + accessibility_dir.mkdir(exist_ok=True) + (accessibility_dir / "config.json").write_text( + json.dumps( + { + "type": "object", + "required_plugins": ["chrome"], + "properties": { + "ACCESSIBILITY_ENABLED": { + "type": "boolean", + "default": True, + }, + }, + }, + ), + ) + (accessibility_dir / "on_Snapshot__10_accessibility.js").write_text("// accessibility hook") + + wget_dir = self.plugins_dir / "wget" + (wget_dir / "config.json").write_text( + json.dumps( + { + "type": "object", + "required_plugins": [], + "properties": { + "WGET_ENABLED": { + "type": "boolean", + "default": True, + "x-aliases": ["SAVE_WGET"], + }, + }, + }, + ), + ) + + with ( + patch.object(hooks_module, "BUILTIN_PLUGINS_DIR", self.plugins_dir), + patch.object(hooks_module, "USER_PLUGINS_DIR", self.test_dir / "user_plugins"), + ): + hooks = hooks_module.discover_hooks("Snapshot", config={"CHROME_ENABLED": False, "WGET_ENABLED": True}) + + hook_names = [hook.parent.name for hook in hooks] + self.assertIn("wget", hook_names) + self.assertNotIn("chrome", hook_names) + self.assertNotIn("accessibility", hook_names) + def test_get_plugins_includes_non_snapshot_plugin_dirs(self): """get_plugins() should include binary-only plugins with standardized metadata.""" env_dir = self.plugins_dir / "env" diff --git a/archivebox/tests/test_machine_models.py b/archivebox/tests/test_machine_models.py index b50edcf1..773668ff 100644 --- a/archivebox/tests/test_machine_models.py +++ b/archivebox/tests/test_machine_models.py @@ -76,6 +76,22 @@ class TestMachineModel(TestCase): # Should have fetched/updated the machine (same GUID) self.assertEqual(machine1.guid, machine2.guid) + def test_machine_current_recreates_stale_cached_row(self): + """Machine.current() should recreate the cached machine if the row was deleted.""" + import archivebox.machine.models as models + + machine1 = Machine.current() + machine1_id = machine1.id + machine1_guid = machine1.guid + + machine1.delete() + models._CURRENT_MACHINE = machine1 + + machine2 = Machine.current() + + self.assertNotEqual(machine1_id, machine2.id) + self.assertEqual(machine1_guid, machine2.guid) + def test_machine_from_jsonl_update(self): """Machine.from_json() should update machine config.""" Machine.current() # Ensure machine exists diff --git a/archivebox/tests/test_runner.py b/archivebox/tests/test_runner.py index 89f8f232..d832d07a 100644 --- a/archivebox/tests/test_runner.py +++ b/archivebox/tests/test_runner.py @@ -70,6 +70,7 @@ def test_run_snapshot_reuses_crawl_bus_for_all_snapshots(monkeypatch): monkeypatch.setattr(runner_module, "CrawlService", _DummyService) monkeypatch.setattr(runner_module, "SnapshotService", _DummyService) monkeypatch.setattr(runner_module, "ArchiveResultService", _DummyService) + monkeypatch.setattr(runner_module, "_emit_machine_config", lambda *args, **kwargs: asyncio.sleep(0)) monkeypatch.setattr(runner_module, "setup_abx_services", lambda *args, **kwargs: None) download_calls = [] @@ -284,6 +285,65 @@ def test_load_run_state_uses_machine_config_as_derived_config(monkeypatch): assert crawl_runner.derived_config == machine.config +def test_load_run_state_uses_enabled_plugins_when_plugins_key_missing(monkeypatch): + from archivebox.machine.models import Machine, NetworkInterface, Process + from archivebox.services import runner as runner_module + from archivebox.config import configset as configset_module + from archivebox import hooks as hooks_module + from archivebox.base_models.models import get_or_create_system_user_pk + from archivebox.crawls.models import Crawl + from pathlib import Path + + machine = Machine.objects.create( + guid="test-guid-runner-missing-plugins", + hostname="runner-host-missing-plugins", + hw_in_docker=False, + hw_in_vm=False, + hw_manufacturer="Test", + hw_product="Test Product", + hw_uuid="test-hw-runner-missing-plugins", + os_arch="arm64", + os_family="darwin", + os_platform="macOS", + os_release="14.0", + os_kernel="Darwin", + stats={}, + config={}, + ) + crawl = Crawl.objects.create( + urls="https://example.com", + created_by_id=get_or_create_system_user_pk(), + ) + proc = SimpleNamespace(iface_id=str(machine.id), machine_id=str(machine.id), iface=None, machine=machine, save=lambda **kwargs: None) + + monkeypatch.setattr( + NetworkInterface, + "current", + classmethod(lambda cls, refresh=False: SimpleNamespace(id=machine.id, machine=machine)), + ) + monkeypatch.setattr(Process, "current", classmethod(lambda cls: proc)) + monkeypatch.setattr(Machine, "current", classmethod(lambda cls: machine)) + monkeypatch.setattr(configset_module, "get_config", lambda **kwargs: {"CHROME_BINARY": "", "TIMEOUT": 60}) + monkeypatch.setattr( + hooks_module, + "discover_hooks", + lambda event_name, config=None: ( + [ + Path(f"/tmp/{event_name.lower()}/wget/on_{event_name}__test.py"), + Path(f"/tmp/{event_name.lower()}/favicon/on_{event_name}__test.py"), + ] + if event_name in {"CrawlSetup", "Snapshot"} + else [] + ), + ) + + crawl_runner = runner_module.CrawlRunner(crawl) + snapshot_ids = crawl_runner.load_run_state() + + assert crawl_runner.selected_plugins == ["favicon", "wget"] + assert len(snapshot_ids) == 1 + + def test_run_snapshot_skips_descendant_when_max_size_already_reached(monkeypatch, tmp_path): from archivebox.base_models.models import get_or_create_system_user_pk from archivebox.crawls.models import Crawl @@ -456,8 +516,19 @@ def test_crawl_runner_does_not_seal_unfinished_crawl(monkeypatch): status=Snapshot.StatusChoices.STARTED, ) + monkeypatch.setattr(runner_module, "_emit_machine_config", lambda *args, **kwargs: asyncio.sleep(0)) monkeypatch.setattr(runner_module, "setup_abx_services", lambda *args, **kwargs: None) monkeypatch.setattr(runner_module.CrawlRunner, "load_run_state", lambda self: [str(snapshot.id)]) + monkeypatch.setattr( + runner_module.CrawlRunner, + "load_snapshot_payload", + lambda self, _snapshot_id: { + "id": str(snapshot.id), + "url": snapshot.url, + "depth": snapshot.depth, + "output_dir": str(snapshot.output_dir), + }, + ) monkeypatch.setattr(runner_module.CrawlRunner, "_create_live_ui", lambda self: None) monkeypatch.setattr(runner_module.CrawlRunner, "run_crawl_setup", lambda self, snapshot_id: asyncio.sleep(0)) monkeypatch.setattr(runner_module.CrawlRunner, "enqueue_snapshot", lambda self, snapshot_id: asyncio.sleep(0)) @@ -497,8 +568,19 @@ def test_crawl_runner_calls_load_and_finalize_run_state(monkeypatch): monkeypatch.setattr(runner_module, "CrawlService", _DummyService) monkeypatch.setattr(runner_module, "SnapshotService", _DummyService) monkeypatch.setattr(runner_module, "ArchiveResultService", _DummyService) + monkeypatch.setattr(runner_module, "_emit_machine_config", lambda *args, **kwargs: asyncio.sleep(0)) monkeypatch.setattr(runner_module, "setup_abx_services", lambda *args, **kwargs: None) monkeypatch.setattr(runner_module.CrawlRunner, "load_run_state", lambda self: [str(snapshot.id)]) + monkeypatch.setattr( + runner_module.CrawlRunner, + "load_snapshot_payload", + lambda self, _snapshot_id: { + "id": str(snapshot.id), + "url": snapshot.url, + "depth": snapshot.depth, + "output_dir": str(snapshot.output_dir), + }, + ) monkeypatch.setattr(runner_module.CrawlRunner, "_create_live_ui", lambda self: None) monkeypatch.setattr(runner_module.CrawlRunner, "run_crawl_setup", lambda self, snapshot_id: asyncio.sleep(0)) monkeypatch.setattr(runner_module.CrawlRunner, "enqueue_snapshot", lambda self, snapshot_id: asyncio.sleep(0)) @@ -588,8 +670,19 @@ def test_crawl_runner_calls_crawl_cleanup_after_snapshot_phase(monkeypatch): status=Snapshot.StatusChoices.STARTED, ) + monkeypatch.setattr(runner_module, "_emit_machine_config", lambda *args, **kwargs: asyncio.sleep(0)) monkeypatch.setattr(runner_module, "setup_abx_services", lambda *args, **kwargs: None) monkeypatch.setattr(runner_module.CrawlRunner, "load_run_state", lambda self: [str(snapshot.id)]) + monkeypatch.setattr( + runner_module.CrawlRunner, + "load_snapshot_payload", + lambda self, _snapshot_id: { + "id": str(snapshot.id), + "url": snapshot.url, + "depth": snapshot.depth, + "output_dir": str(snapshot.output_dir), + }, + ) monkeypatch.setattr(runner_module.CrawlRunner, "_create_live_ui", lambda self: None) monkeypatch.setattr(runner_module.CrawlRunner, "run_crawl_setup", lambda self, snapshot_id: asyncio.sleep(0)) monkeypatch.setattr(runner_module.CrawlRunner, "enqueue_snapshot", lambda self, snapshot_id: asyncio.sleep(0)) @@ -746,3 +839,48 @@ def test_run_pending_crawls_prioritizes_new_queued_crawl_before_snapshot_backlog runner_module.run_pending_crawls(daemon=False) assert run_calls == [(str(newer_crawl.id), None, False)] + + +def test_run_pending_crawls_prioritizes_queued_crawl_before_unrelated_binary_backlog(monkeypatch): + from archivebox.base_models.models import get_or_create_system_user_pk + from archivebox.crawls.models import Crawl + from archivebox.machine.models import Binary, Machine + from archivebox.services import runner as runner_module + + queued_crawl = Crawl.objects.create( + urls="https://scheduled.example.com", + created_by_id=get_or_create_system_user_pk(), + status=Crawl.StatusChoices.QUEUED, + retry_at=runner_module.timezone.now(), + ) + unrelated_binary = Binary.objects.create( + machine=Machine.current(), + name="papers-dl", + status=Binary.StatusChoices.QUEUED, + retry_at=runner_module.timezone.now(), + ) + + monkeypatch.setattr(type(queued_crawl), "claim_processing_lock", lambda self, lock_seconds=60: True) + monkeypatch.setattr(type(unrelated_binary), "claim_processing_lock", lambda self, lock_seconds=60: True) + + run_calls: list[tuple[str, list[str] | None, bool]] = [] + binary_calls: list[str] = [] + + class _StopScheduling(Exception): + pass + + def fake_run_crawl(crawl_id, snapshot_ids=None, selected_plugins=None, process_discovered_snapshots_inline=True): + run_calls.append((crawl_id, snapshot_ids, process_discovered_snapshots_inline)) + raise _StopScheduling + + def fake_run_binary(binary_id): + binary_calls.append(binary_id) + + monkeypatch.setattr(runner_module, "run_crawl", fake_run_crawl) + monkeypatch.setattr(runner_module, "run_binary", fake_run_binary) + + with pytest.raises(_StopScheduling): + runner_module.run_pending_crawls(daemon=False) + + assert run_calls == [(str(queued_crawl.id), None, False)] + assert binary_calls == [] diff --git a/archivebox/tests/test_schedule_e2e.py b/archivebox/tests/test_schedule_e2e.py index 19b18db9..7b4b6c1b 100644 --- a/archivebox/tests/test_schedule_e2e.py +++ b/archivebox/tests/test_schedule_e2e.py @@ -35,6 +35,7 @@ def build_test_env(port: int, **extra: str) -> dict[str, str]: env.pop("DATA_DIR", None) env.update( { + "PLUGINS": "wget", "LISTEN_HOST": f"archivebox.localhost:{port}", "ALLOWED_HOSTS": "*", "CSRF_TRUSTED_ORIGINS": f"http://admin.archivebox.localhost:{port}",