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 import Binary as AbxBinary from abxpkg import BinProvider, PROVIDER_CLASS_BY_NAME 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 ArchiveBoxDBBinaryCacheBackend: """ArchiveBox machine.Binary projection backend for abxpkg BinaryCacheService.""" async def get(self, request: BinaryRequestEvent) -> AbxBinary | None: from archivebox.config.common import get_config from archivebox.machine.models import Binary, Machine, _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 None persisted_overrides = _persisted_overrides_for_request(request) native_overrides = request.overrides or {} requested_provider_names = _provider_names(request.binproviders) await sync_to_async(get_config, thread_sensitive=True)() installed_qs = ( Binary.objects.filter(machine=machine, name=binary_name, status=Binary.StatusChoices.INSTALLED) .exclude(abspath="") .exclude(abspath__isnull=True) .order_by("-modified_at") ) async for installed in installed_qs: installed_path = Path(installed.abspath).expanduser().resolve(strict=False) if not await sync_to_async(installed_path.exists, thread_sensitive=True)(): await _mark_binary_queued(installed) continue if persisted_overrides and installed.overrides != persisted_overrides: await _mark_binary_queued(installed) continue provider_name = (installed.binprovider or installed.binproviders.split(",", 1)[0]).strip() if provider_name and provider_name not in requested_provider_names: await _mark_binary_queued(installed) continue provider = _provider_for_name(provider_name, installed.name, native_overrides) binary_env = BinProvider.build_exec_env(providers=[provider], base_env={}) if provider is not None else {} provider_names = _provider_names(installed.binproviders or request.binproviders or "env") return AbxBinary.model_validate( { "name": request.name, "description": request.description, "binproviders": _providers_for_names(provider_names), "overrides": native_overrides, "loaded_binprovider": provider, "loaded_abspath": installed.abspath, "loaded_version": installed.version or None, "loaded_sha256": installed.sha256 or None, "env": binary_env, }, ) existing = await Binary.objects.filter(machine=machine, name=binary_name).order_by("-modified_at").afirst() if existing is None: await Binary.objects.acreate( machine=machine, name=binary_name, binproviders=_binproviders_to_str(request.binproviders), overrides=persisted_overrides, status=Binary.StatusChoices.QUEUED, ) else: changed = False requested_binproviders = _binproviders_to_str(request.binproviders) if requested_binproviders and existing.binproviders != requested_binproviders: existing.binproviders = requested_binproviders changed = True if persisted_overrides and existing.overrides != persisted_overrides: existing.overrides = persisted_overrides changed = True if existing.status != Binary.StatusChoices.QUEUED: existing.status = Binary.StatusChoices.QUEUED existing.retry_at = None changed = True if changed: await existing.asave(update_fields=["binproviders", "overrides", "status", "retry_at", "modified_at"]) return None async def set(self, request: BinaryRequestEvent | None, binary: AbxBinary) -> None: from archivebox.machine.models import Binary, Machine, _canonical_binary_name machine = await sync_to_async(Machine.current, thread_sensitive=True)() binary_name = _canonical_binary_name(binary.name) if not binary_name: return request_context = request.extra_context if request is not None else {} binary_id = str(request_context.get("binary_id") or "") if binary_id: existing = await Binary.objects.filter(id=binary_id).afirst() else: existing = None if existing is None: existing, _created = await Binary.objects.aget_or_create( machine=machine, name=binary_name, defaults={"status": Binary.StatusChoices.QUEUED}, ) existing.abspath = str(binary.loaded_abspath or "") if binary.loaded_version: existing.version = str(binary.loaded_version) if binary.loaded_sha256: existing.sha256 = str(binary.loaded_sha256) existing.binproviders = _binproviders_to_str( request.binproviders if request is not None else [provider.name for provider in binary.binproviders], ) if binary.loaded_binprovider is not None: existing.binprovider = binary.loaded_binprovider.name existing.overrides = _persisted_overrides_for_request(request) if request is not None else binary.overrides existing.status = Binary.StatusChoices.INSTALLED existing.retry_at = None await existing.asave( update_fields=["abspath", "version", "sha256", "binproviders", "binprovider", "overrides", "status", "retry_at", "modified_at"], ) async def invalidate(self, request: BinaryRequestEvent, binary: AbxBinary, reason: str) -> None: from archivebox.machine.models import Binary, Machine, _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 installed = ( await Binary.objects.filter(machine=machine, name=binary_name, status=Binary.StatusChoices.INSTALLED) .exclude(abspath="") .exclude(abspath__isnull=True) .order_by("-modified_at") .afirst() ) if installed is None: return installed.status = Binary.StatusChoices.QUEUED installed.retry_at = None await installed.asave(update_fields=["status", "retry_at", "modified_at"]) 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 _providers_for_names(names: list[str]) -> list[BinProvider]: providers: list[BinProvider] = [] for name in names: provider_class = PROVIDER_CLASS_BY_NAME.get(name) if provider_class is not None: providers.append(provider_class()) return providers def _provider_for_name(provider_name: str, binary_name: str, overrides: dict[str, Any] | None) -> BinProvider | None: provider_class = PROVIDER_CLASS_BY_NAME.get(provider_name) if provider_class is None: return None provider = provider_class() provider_overrides = overrides.get(provider_name) if isinstance(overrides, dict) else None if isinstance(provider_overrides, dict): provider = provider.get_provider_with_overrides( overrides={binary_name: provider_overrides}, ) return provider async def _mark_binary_queued(binary) -> None: from archivebox.machine.models import Binary if binary.status == Binary.StatusChoices.QUEUED: return binary.status = Binary.StatusChoices.QUEUED binary.retry_at = None await binary.asave(update_fields=["status", "retry_at", "modified_at"]) def _persisted_overrides_for_request(request: BinaryRequestEvent | None) -> dict[str, Any]: if request is None: return {} return dict(request.overrides or {})