mirror of
https://github.com/ArchiveBox/ArchiveBox.git
synced 2026-09-12 19:50:57 +05:00
480ad08cf0
* fix snapshot backfill orchestration regressions * address snapshot orchestration review * fix admin process projection fixture * fix noresult projection fixture
204 lines
8.4 KiB
Python
204 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""Run abx-dl snapshot hooks for existing ArchiveBox snapshots."""
|
|
|
|
__package__ = "archivebox.cli"
|
|
__command__ = "archivebox extract"
|
|
|
|
import sys
|
|
from collections import defaultdict
|
|
from contextlib import redirect_stdout
|
|
|
|
import rich_click as click
|
|
|
|
|
|
def _resolve_requests(records: list[dict], plugins: str) -> tuple[dict[str, set[str]], set[str]]:
|
|
"""Resolve CLI records to snapshot-level execution requests.
|
|
|
|
ArchiveResult input is accepted as a convenient reference to its parent
|
|
Snapshot and plugin. It is never reset or converted into queued work.
|
|
"""
|
|
from archivebox.api.v1_core import _uuid_ref_query
|
|
from archivebox.core.models import ArchiveResult, Snapshot
|
|
from archivebox.misc.jsonl import TYPE_ARCHIVERESULT
|
|
|
|
explicit_plugins = {name.strip() for name in plugins.split(",") if name.strip()}
|
|
requested: dict[str, set[str]] = defaultdict(set)
|
|
missing: set[str] = set()
|
|
|
|
for record in records:
|
|
record_type = record.get("type")
|
|
record_id = str(record.get("id") or "")
|
|
|
|
if record_type == TYPE_ARCHIVERESULT and record.get("snapshot_id"):
|
|
requested[str(record["snapshot_id"])].update(explicit_plugins or {str(record.get("plugin") or "")})
|
|
requested[str(record["snapshot_id"])].discard("")
|
|
continue
|
|
|
|
if record_type == TYPE_ARCHIVERESULT:
|
|
result = ArchiveResult.objects.filter(_uuid_ref_query("id", record_id)).only("snapshot_id", "plugin").first()
|
|
if result is not None:
|
|
requested[str(result.snapshot_id)].update(explicit_plugins or {result.plugin})
|
|
continue
|
|
|
|
snapshot_id = str(record.get("snapshot_id") or record_id or "")
|
|
snapshot = Snapshot.objects.filter(id=snapshot_id).only("id").first() if snapshot_id else None
|
|
if snapshot is None and record.get("url"):
|
|
snapshot = Snapshot.objects.filter(url=record["url"]).order_by("-created_at").only("id").first()
|
|
# Bare UUID CLI arguments are parsed as Snapshot records because the
|
|
# input layer cannot know which model owns them. Preserve the public
|
|
# convenience of passing an ArchiveResult ID by trying that reference
|
|
# only after the Snapshot lookup misses.
|
|
if snapshot is None and record_id and not record.get("url"):
|
|
result = ArchiveResult.objects.filter(_uuid_ref_query("id", record_id)).only("snapshot_id", "plugin").first()
|
|
if result is not None:
|
|
requested[str(result.snapshot_id)].update(explicit_plugins or {result.plugin})
|
|
continue
|
|
if snapshot is None:
|
|
missing.add(snapshot_id or str(record.get("url") or ""))
|
|
continue
|
|
requested[str(snapshot.id)].update(explicit_plugins)
|
|
|
|
return requested, missing
|
|
|
|
|
|
def _run_snapshot_requests(requested: dict[str, set[str]], *, wait: bool, show_progress: bool) -> int:
|
|
from django.utils import timezone
|
|
from rich import print as rprint
|
|
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.services.runner import run_crawl
|
|
|
|
snapshots = {str(snapshot.id): snapshot for snapshot in Snapshot.objects.filter(id__in=requested).select_related("crawl")}
|
|
if not snapshots:
|
|
rprint("[red]No snapshots to process[/red]", file=sys.stderr)
|
|
return 1
|
|
|
|
from archivebox.config.common import get_config
|
|
from archivebox.plugins.discovery import get_enabled_plugins
|
|
|
|
for snapshot_id, plugin_names in requested.items():
|
|
snapshot = snapshots.get(snapshot_id)
|
|
if snapshot is not None and not plugin_names:
|
|
plugin_names.update(get_enabled_plugins(config=get_config(crawl=snapshot.crawl, snapshot=snapshot)))
|
|
|
|
if wait and any("search_backend_sonic" in plugin_names for plugin_names in requested.values()):
|
|
from archivebox.core.takeover_util import ensure_daemon_stack
|
|
|
|
with redirect_stdout(sys.stderr):
|
|
ensure_daemon_stack(reason="Sonic snapshot indexing")
|
|
|
|
# Explicit extraction resumes open/paused snapshots at the snapshot level.
|
|
# Sealed snapshots stay sealed during targeted maintenance backfills.
|
|
if wait:
|
|
for snapshot in snapshots.values():
|
|
if snapshot.status != Snapshot.StatusChoices.SEALED:
|
|
snapshot.update_and_requeue(
|
|
status=Snapshot.StatusChoices.QUEUED,
|
|
retry_at=timezone.now(),
|
|
)
|
|
|
|
if not wait:
|
|
for snapshot_id, plugin_names in requested.items():
|
|
snapshot = snapshots.get(snapshot_id)
|
|
if snapshot is None:
|
|
continue
|
|
if plugin_names:
|
|
snapshot.schedule_plugin_run(plugin_names, when=timezone.now())
|
|
elif snapshot.status == Snapshot.StatusChoices.SEALED:
|
|
snapshot.update_and_requeue(retry_at=timezone.now())
|
|
else:
|
|
snapshot.update_and_requeue(
|
|
status=Snapshot.StatusChoices.QUEUED,
|
|
retry_at=timezone.now(),
|
|
)
|
|
if show_progress:
|
|
rprint(f"[blue]Queued {len(snapshots)} snapshots for extraction[/blue]", file=sys.stderr)
|
|
return 0
|
|
|
|
grouped: dict[tuple[str, tuple[str, ...]], list[str]] = defaultdict(list)
|
|
for snapshot_id, plugin_names in requested.items():
|
|
snapshot = snapshots.get(snapshot_id)
|
|
if snapshot is not None:
|
|
grouped[(str(snapshot.crawl_id), tuple(sorted(plugin_names)))].append(snapshot_id)
|
|
|
|
for (crawl_id, plugin_names), snapshot_ids in grouped.items():
|
|
run_crawl(
|
|
crawl_id,
|
|
snapshot_ids=sorted(snapshot_ids),
|
|
selected_plugins=list(plugin_names) or None,
|
|
show_progress=show_progress,
|
|
)
|
|
return 0
|
|
|
|
|
|
def run_plugins(
|
|
args: tuple,
|
|
records: list[dict] | None = None,
|
|
plugins: str = "",
|
|
wait: bool = True,
|
|
emit_results: bool = True,
|
|
show_progress: bool = True,
|
|
) -> int:
|
|
"""Execute selected plugins through the snapshot-level runner."""
|
|
from rich import print as rprint
|
|
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.misc.jsonl import read_args_or_stdin, write_record
|
|
|
|
if records is None:
|
|
records = list(read_args_or_stdin(args))
|
|
if not records:
|
|
rprint("[yellow]No snapshots provided. Pass snapshot IDs as arguments or via stdin.[/yellow]", file=sys.stderr)
|
|
return 1
|
|
|
|
requested, missing = _resolve_requests(records, plugins)
|
|
for value in sorted(missing):
|
|
rprint(f"[yellow]Snapshot or ArchiveResult not found: {value}[/yellow]", file=sys.stderr)
|
|
if not requested:
|
|
return 1
|
|
|
|
exit_code = _run_snapshot_requests(requested, wait=wait, show_progress=show_progress)
|
|
if exit_code or not emit_results:
|
|
return exit_code
|
|
|
|
is_tty = sys.stdout.isatty()
|
|
for snapshot in Snapshot.objects.filter(id__in=requested):
|
|
results = snapshot.archiveresult_set.all()
|
|
requested_plugins = requested[str(snapshot.id)]
|
|
if requested_plugins:
|
|
results = results.filter(plugin__in=requested_plugins)
|
|
for result in results:
|
|
if is_tty:
|
|
color = {"succeeded": "green", "failed": "red", "skipped": "yellow"}.get(result.status, "dim")
|
|
rprint(f" [{color}]{result.status}[/{color}] {result.plugin} → {result.output_str or ''}", file=sys.stderr)
|
|
else:
|
|
write_record(result.to_json())
|
|
return 0
|
|
|
|
|
|
def process_archiveresult_by_id(archiveresult_id: str) -> int:
|
|
"""Re-run the parent Snapshot plugin referenced by an ArchiveResult."""
|
|
return run_plugins((), records=[{"id": archiveresult_id}], wait=True)
|
|
|
|
|
|
@click.command()
|
|
@click.option("--plugins", "--plugin", "-p", default="", help="Comma-separated list of plugins to run")
|
|
@click.option("--wait/--no-wait", default=True, help="Wait for plugins to complete (default: wait)")
|
|
@click.argument("args", nargs=-1)
|
|
def main(plugins: str, wait: bool, args: tuple):
|
|
"""Run plugins on Snapshots; ArchiveResult IDs select their parent plugin."""
|
|
from archivebox.misc.jsonl import read_args_or_stdin
|
|
|
|
records = list(read_args_or_stdin(args))
|
|
if not records:
|
|
from rich import print as rprint
|
|
|
|
rprint("[yellow]No Snapshot IDs or ArchiveResult IDs provided. Pass as arguments or via stdin.[/yellow]", file=sys.stderr)
|
|
sys.exit(1)
|
|
sys.exit(run_plugins(args, records=records, plugins=plugins, wait=wait))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|