from __future__ import annotations import asyncio import json from collections.abc import Mapping from pathlib import Path from typing import Any from asgiref.sync import sync_to_async from django.utils import timezone from abxpkg.binary_service import BinaryEvent, BinaryRequestEvent from abxpkg.config import load_derived_cache from abxbus import BaseEvent, EventBus from abx_dl.services.base import BaseService class ArchiveBoxBinaryService(BaseService): """Preserve ArchiveBox's legacy Binary Process rows around abxpkg requests.""" LISTENS_TO = [BinaryRequestEvent, BinaryEvent] EMITS: list[type[BaseEvent]] = [] def __init__(self, bus: EventBus): super().__init__(bus) self.process_ids_by_request_id: dict[str, str] = {} self._missing_finalize_tasks: set[asyncio.Task] = set() self.bus.on(BinaryRequestEvent, self.on_BinaryRequestEvent__project_process) self.bus.on(BinaryRequestEvent, self.on_BinaryRequestEvent__schedule_missing_finalize) self.bus.on(BinaryEvent, self.on_BinaryEvent__finalize_process) async def on_BinaryRequestEvent__project_process(self, request: BinaryRequestEvent) -> None: from archivebox.machine.models import Machine, Process, _canonical_binary_name machine = await sync_to_async(Machine.current, thread_sensitive=True)() binary_name = _canonical_binary_name(request.name) if not binary_name: return binary = await self._get_or_create_binary(machine, binary_name, request) started_at = timezone.now() output_dir = self._process_output_dir(binary, request) await sync_to_async(output_dir.mkdir, thread_sensitive=True)(parents=True, exist_ok=True) process = await Process.objects.acreate( machine=machine, iface=None, process_type=Process.TypeChoices.BINARY, worker_type="", pwd=str(output_dir), cmd=self._process_cmd(request), env={}, timeout=int(request.event_timeout or request.install_timeout or 600), pid=None, url=None, started_at=started_at, ended_at=None, stdout="", stderr="", exit_code=None, status=Process.StatusChoices.RUNNING, retry_at=None, binary=binary, ) self.process_ids_by_request_id[request.event_id] = str(process.id) async def on_BinaryEvent__finalize_process(self, event: BinaryEvent) -> None: from archivebox.machine.models import Binary, Process, _canonical_binary_name request = await self.bus.find( BinaryRequestEvent, past=True, future=False, where=lambda candidate: self.bus.event_is_child_of(event, candidate), ) request = request if isinstance(request, BinaryRequestEvent) else None process_id = self.process_ids_by_request_id.pop(request.event_id, "") if request is not None else "" if not process_id: return process = await Process.objects.filter(id=process_id).select_related("binary").afirst() if process is None: return binary_name = _canonical_binary_name(event.name) binary = process.binary if binary is not None and binary_name: binary.abspath = event.abspath if event.version: binary.version = str(event.version) if event.sha256: binary.sha256 = str(event.sha256) binary.binproviders = event.binproviders or binary.binproviders binary.binprovider = event.binprovider or binary.binprovider binary.status = Binary.StatusChoices.INSTALLED binary.retry_at = None await binary.asave( update_fields=["abspath", "version", "sha256", "binproviders", "binprovider", "status", "retry_at", "modified_at"], ) process.ended_at = timezone.now() process.stdout = json.dumps(self._binary_event_json(event, binary)) + "\n" process.stderr = "" process.exit_code = 0 process.status = Process.StatusChoices.EXITED await process.asave(update_fields=["ended_at", "stdout", "stderr", "exit_code", "status", "modified_at"]) if binary is not None: await sync_to_async(self._write_binary_index, thread_sensitive=True)(binary, process, Path(process.pwd)) async def _get_or_create_binary(self, machine, binary_name: str, request: BinaryRequestEvent): from archivebox.machine.models import Binary binary_id = str(request.extra_context.get("binary_id") or "") if binary_id: binary = await Binary.objects.filter(id=binary_id).afirst() if binary is not None: return binary binary = await Binary.objects.filter(machine=machine, name=binary_name).order_by("-modified_at").afirst() if binary is not None: return binary return await Binary.objects.acreate( machine=machine, name=binary_name, binproviders=_binproviders_to_str(request.binproviders), overrides=_persisted_overrides_for_request(request), status=Binary.StatusChoices.QUEUED, ) def _process_cmd(self, request: BinaryRequestEvent) -> list[str]: cmd = [ "abxpkg", "install", f"--name={request.name}", f"--binproviders={_binproviders_to_str(request.binproviders)}", ] if request.overrides: cmd.append(f"--overrides={json.dumps(request.overrides, sort_keys=True)}") return cmd def _binary_event_json(self, event: BinaryEvent, binary) -> dict[str, Any]: if binary is not None: data = binary.to_json() else: data = {"type": "Binary", "name": event.name} data.update( { "type": "Binary", "name": event.name, "binproviders": event.binproviders, "binprovider": event.binprovider, "abspath": event.abspath, "version": str(event.version or ""), "sha256": event.sha256 or "", "status": "installed", }, ) return data async def _finalize_missing_process(self, request: BinaryRequestEvent) -> None: from archivebox.machine.models import Process process_id = self.process_ids_by_request_id.pop(request.event_id, "") if not process_id: return process = await Process.objects.filter(id=process_id).afirst() if process is None or process.status == Process.StatusChoices.EXITED: return process.ended_at = timezone.now() process.stderr = f"Binary request did not resolve: {request.name}" process.exit_code = 1 process.status = Process.StatusChoices.EXITED await process.asave(update_fields=["ended_at", "stderr", "exit_code", "status", "modified_at"]) async def _finalize_request_when_done(self, request: BinaryRequestEvent) -> None: try: await request.wait(timeout=request.event_timeout) except TimeoutError: await self._finalize_missing_process(request) return binary_event = await self.bus.find( BinaryEvent, child_of=request, past=True, future=False, name=request.name, where=lambda candidate: bool(candidate.abspath), ) if not isinstance(binary_event, BinaryEvent): await self._finalize_missing_process(request) def _schedule_missing_finalize(self, request: BinaryRequestEvent) -> None: task = asyncio.create_task(self._finalize_request_when_done(request)) self._missing_finalize_tasks.add(task) task.add_done_callback(lambda done: self._missing_finalize_tasks.discard(done) or (None if done.cancelled() else done.exception())) async def flush_missing_finalizers(self) -> None: if self._missing_finalize_tasks: await asyncio.gather(*tuple(self._missing_finalize_tasks), return_exceptions=False) async def on_BinaryRequestEvent__schedule_missing_finalize(self, request: BinaryRequestEvent) -> None: self._schedule_missing_finalize(request) def _process_output_dir(self, binary, request: BinaryRequestEvent) -> Path: raw_output_dir = str(request.extra_context.get("output_dir") or "").strip() if raw_output_dir: output_dir = Path(raw_output_dir).expanduser() if output_dir.name == str(binary.id): return output_dir.parent return output_dir return binary.output_dir.parent def _write_binary_index(self, binary, process, output_dir: Path) -> None: output_dir.mkdir(parents=True, exist_ok=True) index_path = output_dir / "index.jsonl" with index_path.open("w", encoding="utf-8") as f: f.write(json.dumps(binary.to_json()) + "\n") f.write(json.dumps(process.to_json()) + "\n") def project_abxpkg_derived_cache_to_db(lib_dir: Path | str | None) -> None: """Mirror abxpkg subprocess-resolved binaries into ArchiveBox's DB. Hook shebangs resolve binaries through the abxpkg CLI in subprocesses, so those resolutions cannot emit in-process BinaryRequestEvent/BinaryEvent objects on the ArchiveBox runner bus. abxpkg's generic cross-process projection point is ``LIB_DIR/env/derived.env``; ArchiveBox imports those resolved records here after hook execution. """ if lib_dir is None: return lib_path = Path(lib_dir).expanduser() derived_env_paths = sorted(lib_path.rglob("derived.env")) if lib_path.is_dir() else [] if not derived_env_paths: return from archivebox.machine.models import Binary, Machine, Process, _canonical_binary_name machine = Machine.current() for derived_env_path in derived_env_paths: for record in load_derived_cache(derived_env_path).values(): if not isinstance(record, Mapping): continue binary_name = _canonical_binary_name(str(record.get("bin_name") or "")) if not binary_name: continue abspath = str(record.get("abspath") or "").strip() if not abspath: continue binary_path = Path(abspath).expanduser().resolve(strict=False) if not binary_path.exists(): continue version = str(record.get("loaded_version") or "") sha256 = str(record.get("loaded_sha256") or "") provider_name = str(record.get("provider_name") or "") resolved_provider_name = str(record.get("resolved_provider_name") or provider_name) installed_abspath = str(binary_path) binary, _created = Binary.objects.get_or_create( machine=machine, name=binary_name, defaults={ "status": Binary.StatusChoices.QUEUED, "binproviders": provider_name or resolved_provider_name or "env", }, ) previous_projection = ( binary.status, binary.abspath, binary.version, binary.sha256, binary.binprovider, ) binary.abspath = installed_abspath binary.version = version binary.sha256 = sha256 binary.binproviders = provider_name or resolved_provider_name or binary.binproviders or "env" binary.binprovider = resolved_provider_name or provider_name or binary.binprovider binary.status = Binary.StatusChoices.INSTALLED binary.retry_at = None binary.save( update_fields=[ "abspath", "version", "sha256", "binproviders", "binprovider", "status", "retry_at", "modified_at", ], ) current_projection = ( binary.status, binary.abspath, binary.version, binary.sha256, binary.binprovider, ) if current_projection == previous_projection: continue output_dir = binary.output_dir.parent output_dir.mkdir(parents=True, exist_ok=True) now = timezone.now() process = Process.objects.create( machine=machine, iface=None, process_type=Process.TypeChoices.BINARY, worker_type="", pwd=str(output_dir), cmd=[ "abxpkg", "run", "--script", f"--name={binary_name}", f"--binproviders={binary.binproviders}", ], env={}, timeout=0, pid=None, url=None, started_at=now, ended_at=now, stdout=json.dumps( { "type": "Binary", "name": binary_name, "binproviders": binary.binproviders, "binprovider": binary.binprovider, "abspath": binary.abspath, "version": binary.version, "sha256": binary.sha256, "status": "installed", }, ) + "\n", stderr="", exit_code=0, status=Process.StatusChoices.EXITED, retry_at=None, binary=binary, ) index_path = output_dir / "index.jsonl" with index_path.open("w", encoding="utf-8") as f: f.write(json.dumps(binary.to_json()) + "\n") f.write(json.dumps(process.to_json()) + "\n") def _provider_names(binproviders: str | list[str] | None) -> list[str]: if isinstance(binproviders, str): raw_names = [part.strip() for part in binproviders.split(",")] elif binproviders: raw_names = [str(part).strip() for part in binproviders] else: raw_names = ["env"] names: list[str] = [] for name in raw_names: if name and name not in names: names.append(name) return names or ["env"] def _binproviders_to_str(binproviders: str | list[str] | None) -> str: return ",".join(_provider_names(binproviders)) def _persisted_overrides_for_request(request: BinaryRequestEvent | None) -> dict[str, Any]: if request is None: return {} return dict(request.overrides or {})