diff --git a/archivebox/api/v1_core.py b/archivebox/api/v1_core.py index 97d7d916..0ade551d 100644 --- a/archivebox/api/v1_core.py +++ b/archivebox/api/v1_core.py @@ -11,7 +11,7 @@ from uuid import UUID from typing import Union, Any, Annotated from datetime import datetime, time -from django.db import transaction +from django.db import IntegrityError from django.db.models import Model, Q from django.http import HttpRequest, HttpResponse from django.http.multipartparser import MultiPartParser, MultiPartParserError @@ -517,52 +517,53 @@ def create_archiveresult( normalized_status = ArchiveResult.normalize_status(status) parsed_output_json = _parse_archiveresult_output_json(output_json) hook = hook_name or ARCHIVERESULT_UPLOAD_HOOK_NAME - existing_result = ArchiveResult.objects.filter(snapshot=snapshot, plugin=plugin_name).first() - existing_output_files = dict(existing_result.output_files or {}) if existing_result else {} - output_files = _write_archiveresult_files( + result_lookup = { + "snapshot": snapshot, + "plugin": plugin_name, + "hook_name": hook, + } + uploaded_output_files = _write_archiveresult_files( request, snapshot, plugin_name, - existing_output_files=existing_output_files, allow_empty=True, ) - now = timezone.now() - - with transaction.atomic(): - Snapshot.objects.select_for_update().get(pk=snapshot.pk) - existing_result = ArchiveResult.objects.filter(snapshot=snapshot, plugin=plugin_name).first() - if existing_result: - output_files = { - **dict(existing_result.output_files or {}), - **output_files, - } - result = existing_result - else: - existing_result = None - result = ArchiveResult( - snapshot=snapshot, - plugin=plugin_name, - hook_name=hook, - ) - - if ( - existing_result - and normalized_status == ArchiveResult.StatusChoices.STARTED - and existing_result.status != ArchiveResult.StatusChoices.STARTED - ): - normalized_status = existing_result.status + result = ArchiveResult.objects.filter(**result_lookup).first() + for _attempt in range(3): + output_files = { + **(result.output_file_map() if result else {}), + **uploaded_output_files, + } output_size, output_mimetypes = _summarize_archiveresult_output_files(output_files) output_file_paths = list(output_files.keys()) - result.status = normalized_status - result.hook_name = hook - result.output_str = output_str or (output_file_paths[0] if output_file_paths else "") - result.output_json = parsed_output_json - result.output_files = output_files - result.output_size = output_size - result.output_mimetypes = output_mimetypes - result.start_ts = result.start_ts or now - result.end_ts = now - result.save() + result_status = normalized_status + if result and result_status == ArchiveResult.StatusChoices.STARTED and result.status != ArchiveResult.StatusChoices.STARTED: + result_status = result.status + now = timezone.now() + values = { + "status": result_status, + "output_str": output_str or (output_file_paths[0] if output_file_paths else ""), + "output_json": parsed_output_json, + "output_files": output_files, + "output_size": output_size, + "output_mimetypes": output_mimetypes, + "start_ts": result.start_ts or now if result else now, + "end_ts": now, + } + if result: + if result.safe_update(values): + break + continue + result, created = ArchiveResult.get_or_create_by_hook( + snapshot, + plugin_name, + hook, + defaults=values, + ) + if created: + break + else: + raise HttpError(409, "ArchiveResult changed while upload metadata was being updated") if result.status != ArchiveResult.StatusChoices.STARTED: _queue_archiveresult_snapshot_maintenance(snapshot) @@ -576,44 +577,46 @@ def patch_archiveresult( ): """Append or replace files on an existing ArchiveResult.""" result = ArchiveResult.objects.select_related("snapshot__crawl__created_by").get(_uuid_ref_query("id", archiveresult_id)) - output_files = _write_archiveresult_files( + uploaded_output_files = _write_archiveresult_files( request, result.snapshot, result.plugin, - existing_output_files=result.output_file_map(), ) - latest_result = ArchiveResult.objects.only("output_files", "status").get(pk=result.pk) - output_files = { - **latest_result.output_file_map(), - **output_files, - } - output_size, output_mimetypes = _summarize_archiveresult_output_files(output_files) - - update_fields = ["output_files", "output_size", "output_mimetypes", "end_ts", "modified_at"] - result.output_files = output_files - result.output_size = output_size - result.output_mimetypes = output_mimetypes - result.end_ts = timezone.now() output_str = _get_archiveresult_upload_form_value(request, "output_str") status = _get_archiveresult_upload_form_value(request, "status") output_json = _get_archiveresult_upload_form_value(request, "output_json") - if output_str: - result.output_str = output_str - update_fields.append("output_str") - if status: - normalized_status = ArchiveResult.normalize_status(status) - if normalized_status == ArchiveResult.StatusChoices.STARTED and latest_result.status != ArchiveResult.StatusChoices.STARTED: - normalized_status = latest_result.status - result.status = normalized_status - update_fields.append("status") - elif latest_result.status == ArchiveResult.StatusChoices.QUEUED and ArchiveResult.output_files_upload_complete(output_files): - result.status = ArchiveResult.StatusChoices.SUCCEEDED - update_fields.append("status") - if output_json: - result.output_json = _parse_archiveresult_output_json(output_json) - update_fields.append("output_json") + parsed_output_json = _parse_archiveresult_output_json(output_json) if output_json else None + + if not ArchiveResult.output_files_upload_complete(uploaded_output_files): + result.output_files = {**result.output_file_map(), **uploaded_output_files} + result.output_size, result.output_mimetypes = _summarize_archiveresult_output_files(result.output_files) + return result + + for _attempt in range(3): + output_files = {**result.output_file_map(), **uploaded_output_files} + output_size, output_mimetypes = _summarize_archiveresult_output_files(output_files) + values: dict[str, Any] = { + "output_files": output_files, + "output_size": output_size, + "output_mimetypes": output_mimetypes, + "end_ts": timezone.now(), + } + if output_str: + values["output_str"] = output_str + if status: + normalized_status = ArchiveResult.normalize_status(status) + if normalized_status == ArchiveResult.StatusChoices.STARTED and result.status != ArchiveResult.StatusChoices.STARTED: + normalized_status = result.status + values["status"] = normalized_status + elif result.status == ArchiveResult.StatusChoices.QUEUED: + values["status"] = ArchiveResult.StatusChoices.SUCCEEDED + if output_json: + values["output_json"] = parsed_output_json + if result.safe_update(values): + break + else: + raise HttpError(409, "ArchiveResult changed while upload metadata was being updated") - result.save(update_fields=update_fields) if result.status != ArchiveResult.StatusChoices.STARTED: _queue_archiveresult_snapshot_maintenance(result.snapshot) @@ -930,52 +933,47 @@ def create_snapshot(request: HttpRequest, data: SnapshotCreateSchema): created_by=request.user if isinstance(request.user, User) else None, ) - # Browser clients sync metadata and upload capture artifacts immediately - # after queueing a URL. The runner holds the crawl lifecycle lock for the - # whole crawl, so idempotently posting an already-created snapshot must not - # wait behind that potentially long-running archive job. - existing_snapshot = Snapshot.objects.filter(url=data.url, crawl=crawl).first() - if existing_snapshot is not None and (status is None or existing_snapshot.status == status): - if data.title is not None and existing_snapshot.title != data.title: - Snapshot.objects.filter(pk=existing_snapshot.pk).update(title=data.title, modified_at=timezone.now()) - existing_snapshot.title = data.title - if tags: - existing_snapshot.save_tags(tags) - setattr(request, "with_archiveresults", False) - return existing_snapshot + # Browser uploads must not wait behind the runner's crawl-wide lifecycle + # lock. The unique insert recovery and CAS update below are the request-side + # coordination boundary for this idempotent metadata sync. + snapshot = Snapshot.objects.filter(url=data.url, crawl=crawl).first() + if snapshot is None: + try: + snapshot = Snapshot.objects.create( + url=data.url, + crawl=crawl, + depth=data.depth, + title=data.title, + timestamp=str(timezone.now().timestamp()), + status=status or Snapshot.StatusChoices.QUEUED, + retry_at=timezone.now(), + ) + except IntegrityError: + snapshot = Snapshot.objects.filter(url=data.url, crawl=crawl).first() + if snapshot is None: + raise - with crawl_lifecycle_lock(str(crawl.id)): - snapshot_defaults = { - "depth": data.depth, - "title": data.title, - "timestamp": str(timezone.now().timestamp()), - "status": status or Snapshot.StatusChoices.QUEUED, - "retry_at": timezone.now(), - } - snapshot, _ = Snapshot.objects.get_or_create( - url=data.url, - crawl=crawl, - defaults=snapshot_defaults, + for _attempt in range(3): + updates: dict[str, Any] = {} + if data.title is not None and snapshot.title != data.title: + updates["title"] = data.title + if status is not None and snapshot.status != status: + updates["status"] = status + if not updates or snapshot.safe_update(updates, extra_filter={"modified_at": snapshot.modified_at}): + break + else: + raise HttpError(409, "Snapshot changed while metadata was being updated") + + if tags: + snapshot.save_tags( + tags, + created_by=request.user if isinstance(request.user, User) else None, ) - update_fields: list[str] = [] - if data.title is not None and snapshot.title != data.title: - snapshot.title = data.title - update_fields.append("title") - if status is not None and snapshot.status != status: - snapshot.status = status - update_fields.append("status") - if update_fields: - update_fields.append("modified_at") - snapshot.save(update_fields=update_fields) - - if tags: - snapshot.save_tags(tags) - - try: - snapshot.ensure_crawl_symlink() - except Exception: - pass + try: + snapshot.ensure_crawl_symlink() + except Exception: + pass setattr(request, "with_archiveresults", False) return snapshot @@ -1023,7 +1021,10 @@ def patch_snapshot(request: HttpRequest, snapshot_id: str, data: SnapshotUpdateS update_fields.append("retry_at") if tags is not None: - snapshot.save_tags(normalize_tag_list(tags)) + snapshot.save_tags( + normalize_tag_list(tags), + created_by=request.user if isinstance(request.user, User) else None, + ) if payload.get("status") == Snapshot.StatusChoices.SEALED: snapshot.cancel() @@ -1410,8 +1411,7 @@ def tags_add_to_snapshot(request: HttpRequest, data: TagSnapshotRequestSchema): else: raise HttpError(400, "Either tag_name or tag_id is required") - # Add the tag to the snapshot - snapshot.tags.add(tag.pk) + snapshot.add_tag_ids([tag.pk]) return { "success": True, @@ -1439,8 +1439,7 @@ def tags_remove_from_snapshot(request: HttpRequest, data: TagSnapshotRequestSche else: raise HttpError(400, "Either tag_name or tag_id is required") - # Remove the tag from the snapshot - snapshot.tags.remove(tag.pk) + snapshot.remove_tag_ids([tag.pk]) return { "success": True, diff --git a/archivebox/cli/archivebox_archiveresult.py b/archivebox/cli/archivebox_archiveresult.py index 63809f78..dc7cf81a 100644 --- a/archivebox/cli/archivebox_archiveresult.py +++ b/archivebox/cli/archivebox_archiveresult.py @@ -67,7 +67,7 @@ def create_archiveresults( Reads Snapshot records from stdin and emits ArchiveResult request JSONL. Pass-through: Non-Snapshot/ArchiveResult records are output unchanged. If --plugin is specified, only emits requests for that plugin. - Otherwise, emits one request for each enabled Snapshot plugin. + Otherwise, emits requests for all enabled snapshot hooks. Exit codes: 0: Success @@ -144,17 +144,18 @@ def create_archiveresults( created_count = 0 for snapshot in snapshots: - if plugin: + config = get_config(crawl=snapshot.crawl, snapshot=snapshot) + hooks = [ + hook + for hook in discover_hooks("Snapshot", filter_disabled=not plugin, config=config) + if not plugin or hook.parent.name == plugin + ] + for hook_path in hooks: + hook_name = hook_path.stem + plugin_name = hook_path.parent.name if not is_tty: - write_record(build_archiveresult_request(snapshot.id, plugin, status=status)) + write_record(build_archiveresult_request(snapshot.id, plugin_name, hook_name=hook_name, status=status)) created_count += 1 - else: - config = get_config(crawl=snapshot.crawl, snapshot=snapshot) - hooks = discover_hooks("Snapshot", filter_disabled=True, config=config) - for plugin_name in dict.fromkeys(hook_path.parent.name for hook_path in hooks): - if not is_tty: - write_record(build_archiveresult_request(snapshot.id, plugin_name, status=status)) - created_count += 1 rprint(f"[green]Created {created_count} archive result request records[/green]", file=sys.stderr) return 0 @@ -254,11 +255,9 @@ def update_archiveresults( # Apply updates from CLI flags if status: - if status == ArchiveResult.StatusChoices.QUEUED: - result.reset_for_retry() - else: - result.status = status - result.save(update_fields=["status", "modified_at"]) + result.status = status + + result.save() updated_count += 1 if not is_tty: diff --git a/archivebox/cli/archivebox_extract.py b/archivebox/cli/archivebox_extract.py index 87f23852..f5f915fd 100644 --- a/archivebox/cli/archivebox_extract.py +++ b/archivebox/cli/archivebox_extract.py @@ -70,7 +70,12 @@ def process_archiveresult_by_id(archiveresult_id: str) -> int: ) return 1 - run_crawl(str(snapshot.crawl_id), snapshot_ids=[str(snapshot.id)], selected_plugins=[archiveresult.plugin]) + run_crawl( + str(snapshot.crawl_id), + snapshot_ids=[str(snapshot.id)], + selected_plugins=[archiveresult.plugin], + selected_plugins_are_explicit=False, + ) archiveresult.refresh_from_db() if archiveresult.status == ArchiveResult.StatusChoices.SUCCEEDED: @@ -122,6 +127,7 @@ def run_plugins( from archivebox.core.models import Snapshot from archivebox.core.models import ArchiveResult from archivebox.services.runner import run_crawl + from abx_dl.models import discover_plugins is_tty = sys.stdout.isatty() @@ -142,6 +148,8 @@ def run_plugins( # Gather snapshot IDs and optional plugin constraints to process snapshot_ids = set() requested_plugins_by_snapshot: dict[str, set[str]] = defaultdict(set) + requested_hooks_by_snapshot: dict[str, set[tuple[str, str]]] = defaultdict(set) + plugin_level_requests_by_snapshot: dict[str, set[str]] = defaultdict(set) for record in records: record_type = record.get("type") @@ -163,7 +171,14 @@ def run_plugins( snapshot_ids.add(str(snapshot_id)) plugin_name = record.get("plugin") if plugin_name and not plugins_list: - requested_plugins_by_snapshot[str(snapshot_id)].add(str(plugin_name)) + snapshot_key = str(snapshot_id) + plugin_key = str(plugin_name) + requested_plugins_by_snapshot[snapshot_key].add(plugin_key) + hook_name = str(record.get("hook_name") or "") + if hook_name: + requested_hooks_by_snapshot[snapshot_key].add((plugin_key, hook_name)) + else: + plugin_level_requests_by_snapshot[snapshot_key].add(plugin_key) elif "id" in record: # Assume it's a snapshot ID @@ -192,21 +207,38 @@ def run_plugins( if snapshot_id in existing_snapshot_ids for plugin_name in plugin_names ) - queued_rows: set[tuple[str, str]] = set() - if preserve_queued and requested_pairs: - queued_rows = { - (str(snapshot_id), plugin_name) - for snapshot_id, plugin_name in ArchiveResult.objects.filter( - snapshot_id__in=existing_snapshot_ids, - plugin__in={plugin_name for _snapshot_id, plugin_name in requested_pairs}, - status=ArchiveResult.StatusChoices.QUEUED, - ).values_list("snapshot_id", "plugin") + plugins_by_name = discover_plugins(runtime="archivebox") + requested_rows: set[tuple[str, str, str]] = set() + for snapshot_id, plugin_name in requested_pairs: + exact_hook_names = { + hook_name + for requested_plugin, hook_name in requested_hooks_by_snapshot.get(snapshot_id, set()) + if requested_plugin == plugin_name } - rows_to_queue = requested_pairs - queued_rows + if exact_hook_names and plugin_name not in plugin_level_requests_by_snapshot.get(snapshot_id, set()): + requested_rows.update((snapshot_id, plugin_name, hook_name) for hook_name in exact_hook_names) + continue + plugin = plugins_by_name.get(plugin_name) + hooks = plugin.filter_hooks("Snapshot") if plugin is not None else [] + if hooks: + requested_rows.update((snapshot_id, plugin_name, hook.name) for hook in hooks) + else: + requested_rows.add((snapshot_id, plugin_name, "")) + + queued_rows: set[tuple[str, str, str]] = set() + if preserve_queued and requested_rows: + queued_rows = { + (str(snapshot_id), plugin_name, hook_name) + for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( + snapshot_id__in=existing_snapshot_ids, + plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in requested_rows}, + status=ArchiveResult.StatusChoices.QUEUED, + ).values_list("snapshot_id", "plugin", "hook_name") + } + rows_to_queue = requested_rows - queued_rows reset_fields = { "status": ArchiveResult.StatusChoices.QUEUED, - "hook_name": "", "output_str": "", "output_json": None, "output_files": {}, @@ -217,28 +249,28 @@ def run_plugins( "modified_at": timezone.now(), } if rows_to_queue and plugins_list: - rows_to_reset_by_plugin: dict[str, set[str]] = defaultdict(set) - for snapshot_id, plugin_name in rows_to_queue: - rows_to_reset_by_plugin[plugin_name].add(snapshot_id) - for plugin_name, plugin_snapshot_ids in rows_to_reset_by_plugin.items(): - ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name).update( + rows_to_reset_by_hook: dict[tuple[str, str], set[str]] = defaultdict(set) + for snapshot_id, plugin_name, hook_name in rows_to_queue: + rows_to_reset_by_hook[(plugin_name, hook_name)].add(snapshot_id) + for (plugin_name, hook_name), plugin_snapshot_ids in rows_to_reset_by_hook.items(): + ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name, hook_name=hook_name).update( **reset_fields, ) elif rows_to_queue and requested_plugins_by_snapshot: - snapshot_ids_by_plugin: dict[str, set[str]] = defaultdict(set) - for snapshot_id, plugin_name in rows_to_queue: - snapshot_ids_by_plugin[plugin_name].add(snapshot_id) - for plugin_name, plugin_snapshot_ids in snapshot_ids_by_plugin.items(): - ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name).update( + snapshot_ids_by_hook: dict[tuple[str, str], set[str]] = defaultdict(set) + for snapshot_id, plugin_name, hook_name in rows_to_queue: + snapshot_ids_by_hook[(plugin_name, hook_name)].add(snapshot_id) + for (plugin_name, hook_name), plugin_snapshot_ids in snapshot_ids_by_hook.items(): + ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name, hook_name=hook_name).update( **reset_fields, ) existing_rows = ( { - (str(snapshot_id), plugin_name) - for snapshot_id, plugin_name in ArchiveResult.objects.filter( + (str(snapshot_id), plugin_name, hook_name) + for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( snapshot_id__in=existing_snapshot_ids, - plugin__in={plugin_name for _snapshot_id, plugin_name in rows_to_queue}, - ).values_list("snapshot_id", "plugin") + plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in rows_to_queue}, + ).values_list("snapshot_id", "plugin", "hook_name") } if rows_to_queue else set() @@ -250,10 +282,10 @@ def run_plugins( ArchiveResult( snapshot_id=snapshot_id, plugin=plugin_name, - hook_name="", + hook_name=hook_name, status=ArchiveResult.StatusChoices.QUEUED, ) - for snapshot_id, plugin_name in sorted(missing_rows) + for snapshot_id, plugin_name, hook_name in sorted(missing_rows) ], batch_size=500, ) @@ -261,13 +293,13 @@ def run_plugins( processed_count = len(existing_snapshot_ids) queue_at = timezone.now() if existing_snapshot_ids: - if requested_pairs: - # Search indexing on a sealed Snapshot is the only targeted plugin + if requested_rows: + # Search indexing on a sealed Snapshot is the only targeted hook # allowed to bypass the normal lifecycle. Every other requested # plugin requeues its Snapshot through the unified lifecycle. - affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name in rows_to_queue} + affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in rows_to_queue} if preserve_queued and queued_rows: - queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name in queued_rows} + queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in queued_rows} affected_snapshot_ids.update( str(snapshot_id) for snapshot_id in Snapshot.objects.filter(id__in=queued_snapshot_ids) @@ -282,7 +314,7 @@ def run_plugins( ) ) requested_plugins_by_id: dict[str, set[str]] = defaultdict(set) - for snapshot_id, plugin_name in requested_pairs: + for snapshot_id, plugin_name, _hook_name in requested_rows: requested_plugins_by_id[snapshot_id].add(plugin_name) for snapshot in Snapshot.objects.filter(id__in=affected_snapshot_ids).only("id", "status", "modified_at"): # Guard the read-time status so we never bump retry_at on a @@ -318,7 +350,7 @@ def run_plugins( refresh=False, extra_filter={"status": snapshot.status}, ) - if existing_crawl_ids and not requested_pairs: + if existing_crawl_ids and not requested_rows: from archivebox.crawls.models import Crawl for crawl in Crawl.objects.filter(id__in=existing_crawl_ids).only("id", "status", "retry_at", "modified_at"): @@ -362,7 +394,7 @@ def run_plugins( selected_plugins = ( plugins_list or sorted( - {plugin for snapshot_id in crawl_snapshot_ids for plugin in requested_plugins_by_snapshot.get(str(snapshot_id), set())}, + {plugin for snapshot_id, plugin, _hook_name in requested_rows if snapshot_id in crawl_snapshot_ids}, ) or None ) @@ -371,6 +403,7 @@ def run_plugins( snapshot_ids=sorted(crawl_snapshot_ids), selected_plugins=selected_plugins, show_progress=show_progress, + selected_plugins_are_explicit=bool(plugins_list), ) if not emit_results: diff --git a/archivebox/cli/archivebox_run.py b/archivebox/cli/archivebox_run.py index 320e0e62..1275714a 100644 --- a/archivebox/cli/archivebox_run.py +++ b/archivebox/cli/archivebox_run.py @@ -156,18 +156,7 @@ def process_stdin_records() -> int: queued_count += 1 elif record_type == TYPE_ARCHIVERESULT: - if record_id: - # Existing archiveresult - re-queue - try: - archiveresult = ArchiveResult.objects.get(id=record_id) - except ArchiveResult.DoesNotExist: - archiveresult = None - else: - archiveresult = None - - snapshot_id = record.get("snapshot_id") - plugin_name = record.get("plugin") - snapshot = None + archiveresult = ArchiveResult.from_json(record) if archiveresult: if archiveresult.status in [ ArchiveResult.StatusChoices.FAILED, @@ -177,12 +166,10 @@ def process_stdin_records() -> int: ]: archiveresult.reset_for_retry() snapshot = archiveresult.snapshot - plugin_name = plugin_name or archiveresult.plugin - elif snapshot_id: - try: - snapshot = Snapshot.objects.get(id=snapshot_id) - except Snapshot.DoesNotExist: - snapshot = None + plugin_name = archiveresult.plugin + else: + snapshot = None + plugin_name = None if snapshot: snapshot.queue_for_extraction() diff --git a/archivebox/core/admin_snapshots.py b/archivebox/core/admin_snapshots.py index 45687c52..341daf78 100644 --- a/archivebox/core/admin_snapshots.py +++ b/archivebox/core/admin_snapshots.py @@ -346,19 +346,8 @@ class SnapshotAdminForm(forms.ModelForm): # Parse and save tags from tags_editor tags_str = self.cleaned_data.get("tags_editor", "") - if tags_str: - tag_names = [name.strip() for name in tags_str.split(",") if name.strip()] - tags = [] - for name in tag_names: - tag, _ = Tag.objects.get_or_create( - name__iexact=name, - defaults={"name": name}, - ) - tag = Tag.objects.filter(name__iexact=name).first() or tag - tags.append(tag) - instance.tags.set(tags) - else: - instance.tags.clear() + tag_names = [name.strip() for name in tags_str.split(",") if name.strip()] + instance.save_tags(tag_names) return instance @@ -367,6 +356,16 @@ class SnapshotAdmin(SearchResultsAdminMixin, ConfigEditorMixin, BaseModelAdmin): form = SnapshotAdminForm raw_id_fields = ("crawl", "parent_snapshot") list_select_related = () + + def save_related(self, request, form, formsets, change): + super().save_related(request, form, formsets, change) + tags_str = form.cleaned_data.get("tags_editor", "") + tag_names = [name.strip() for name in tags_str.split(",") if name.strip()] + form.instance.save_tags( + tag_names, + created_by=request.user if request.user.is_authenticated else None, + ) + list_display = ( "permissions_badge", "created_at", diff --git a/archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py b/archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py index 11775c69..fad6b4de 100644 --- a/archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py +++ b/archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py @@ -1,13 +1,36 @@ +import hashlib + from django.db import migrations, models -from django.db.models import Sum + + +STASH_KEY = "__archivebox_0052_original_row" +TEMP_PLUGIN_PREFIX = "__abx52_" + + +def _temporary_plugin_name(row_id, reserved_plugins): + """Return a deterministic 32-character plugin name unused by this snapshot.""" + salt = 0 + while True: + digest = hashlib.sha256(f"{row_id}:{salt}".encode()).hexdigest()[:24] + candidate = f"{TEMP_PLUGIN_PREFIX}{digest}" + if candidate not in reserved_plugins: + return candidate + salt += 1 def consolidate_archiveresults_per_plugin(apps, schema_editor): + """Temporarily make plugin names unique without deleting hook history. + + This migration shipped with a short-lived one-row-per-plugin model, while + 0054 restores the durable (snapshot, plugin, hook_name) identity. Fresh + upgrades still traverse both published migrations, so deleting duplicate + rows here would lose history before 0054 gets a chance to restore the + correct constraint. Rename only the non-canonical rows and stash their + original plugin/output_json values for 0054 to restore verbatim. + """ ArchiveResult = apps.get_model("core", "ArchiveResult") - Snapshot = apps.get_model("core", "Snapshot") duplicate_groups = ArchiveResult.objects.values("snapshot_id", "plugin").annotate(count=models.Count("id")).filter(count__gt=1) - affected_snapshot_ids = set() for group in duplicate_groups.iterator(chunk_size=200): rows = list( ArchiveResult.objects.filter( @@ -24,27 +47,22 @@ def consolidate_archiveresults_per_plugin(apps, schema_editor): str(row.id), ), ) - output_files = {} - mimetypes = set() - for row in rows: - output_files.update(row.output_files or {}) - mimetypes.update(part.strip() for part in (row.output_mimetypes or "").split(",") if part.strip()) - - winner.output_files = output_files - winner.output_size = max( - sum(int(metadata.get("size") or 0) for metadata in output_files.values() if isinstance(metadata, dict)), - *(int(row.output_size or 0) for row in rows), + reserved_plugins = set( + ArchiveResult.objects.filter(snapshot_id=group["snapshot_id"]).values_list("plugin", flat=True), ) - winner.output_mimetypes = ",".join(sorted(mimetypes)) - winner.start_ts = min((row.start_ts for row in rows if row.start_ts), default=None) - winner.end_ts = max((row.end_ts for row in rows if row.end_ts), default=None) - winner.save(update_fields=["output_files", "output_size", "output_mimetypes", "start_ts", "end_ts"]) - ArchiveResult.objects.filter(id__in=[row.id for row in rows if row.id != winner.id]).delete() - affected_snapshot_ids.add(group["snapshot_id"]) - - for snapshot_id in affected_snapshot_ids: - total = ArchiveResult.objects.filter(snapshot_id=snapshot_id).aggregate(total=Sum("output_size"))["total"] or 0 - Snapshot.objects.filter(id=snapshot_id).update(output_size=total) + for row in rows: + if row.id == winner.id: + continue + temporary_plugin = _temporary_plugin_name(row.id, reserved_plugins) + reserved_plugins.add(temporary_plugin) + row.output_json = { + STASH_KEY: { + "plugin": row.plugin, + "output_json": row.output_json, + }, + } + row.plugin = temporary_plugin + row.save(update_fields=["plugin", "output_json"]) class Migration(migrations.Migration): diff --git a/archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py b/archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py new file mode 100644 index 00000000..d19299ef --- /dev/null +++ b/archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py @@ -0,0 +1,90 @@ +import hashlib + +from django.db import migrations, models + + +STASH_KEY = "__archivebox_0052_original_row" +TEMP_PLUGIN_PREFIX = "__abx52_" + + +def _temporary_plugin_name(row_id, reserved_plugins): + salt = 0 + while True: + digest = hashlib.sha256(f"{row_id}:{salt}".encode()).hexdigest()[:24] + candidate = f"{TEMP_PLUGIN_PREFIX}{digest}" + if candidate not in reserved_plugins: + return candidate + salt += 1 + + +def stash_archiveresult_hook_rows(apps, schema_editor): + """Make plugin names temporarily unique when reversing to 0053.""" + ArchiveResult = apps.get_model("core", "ArchiveResult") + duplicate_groups = ArchiveResult.objects.values("snapshot_id", "plugin").annotate(count=models.Count("id")).filter(count__gt=1) + for group in duplicate_groups.iterator(chunk_size=200): + rows = list( + ArchiveResult.objects.filter(snapshot_id=group["snapshot_id"], plugin=group["plugin"]).order_by("created_at", "id"), + ) + winner = max( + rows, + key=lambda row: ( + bool(row.output_files), + int(row.output_size or 0), + row.modified_at, + str(row.id), + ), + ) + reserved_plugins = set( + ArchiveResult.objects.filter(snapshot_id=group["snapshot_id"]).values_list("plugin", flat=True), + ) + for row in rows: + if row.id == winner.id: + continue + temporary_plugin = _temporary_plugin_name(row.id, reserved_plugins) + reserved_plugins.add(temporary_plugin) + row.output_json = { + STASH_KEY: { + "plugin": row.plugin, + "output_json": row.output_json, + }, + } + row.plugin = temporary_plugin + row.save(update_fields=["plugin", "output_json"]) + + +def restore_archiveresult_hook_rows(apps, schema_editor): + """Undo 0052's temporary plugin renames after its constraint is removed.""" + ArchiveResult = apps.get_model("core", "ArchiveResult") + rows = ArchiveResult.objects.filter(plugin__startswith=TEMP_PLUGIN_PREFIX) + for row in rows.iterator(chunk_size=200): + stash = row.output_json + original = stash.get(STASH_KEY) if isinstance(stash, dict) else None + if not isinstance(original, dict) or "plugin" not in original: + continue + row.plugin = original["plugin"] + row.output_json = original.get("output_json") + row.save(update_fields=["plugin", "output_json"]) + + +class Migration(migrations.Migration): + dependencies = [ + ("core", "0053_alter_archiveresult_options"), + ] + + operations = [ + migrations.RemoveConstraint( + model_name="archiveresult", + name="unique_archiveresult_per_snapshot_plugin", + ), + migrations.RunPython( + restore_archiveresult_hook_rows, + reverse_code=stash_archiveresult_hook_rows, + ), + migrations.AddConstraint( + model_name="archiveresult", + constraint=models.UniqueConstraint( + fields=("snapshot", "plugin", "hook_name"), + name="unique_archiveresult_per_snapshot_hook", + ), + ), + ] diff --git a/archivebox/core/models.py b/archivebox/core/models.py index d99f2843..7f175359 100644 --- a/archivebox/core/models.py +++ b/archivebox/core/models.py @@ -12,7 +12,7 @@ from urllib.parse import urlparse from django.conf import settings from django.contrib import admin from django.core.exceptions import FieldDoesNotExist, ObjectDoesNotExist, ValidationError -from django.db import models, transaction +from django.db import IntegrityError, models, transaction from django.db.models import Case, F, Q, QuerySet, Sum, Value, When from django.db.models.fields.json import KT from django.db.models.functions import Coalesce, Concat @@ -85,6 +85,19 @@ class Tag(ModelWithUUID): modified_at = models.DateTimeField(auto_now=True) name = models.CharField(unique=True, blank=False, max_length=100) + @classmethod + def get_or_create_by_name(cls, name: str, *, defaults: Mapping[str, Any] | None = None) -> tuple["Tag", bool]: + tag = cls.objects.filter(name__iexact=name).first() + if tag: + return tag, False + try: + return cls.objects.create(name=name, **(defaults or {})), True + except IntegrityError: + tag = cls.objects.filter(name__iexact=name).first() + if tag is None: + raise + return tag, False + snapshot_set: models.Manager["Snapshot"] class Meta(ModelWithUUID.Meta): @@ -139,7 +152,7 @@ class Tag(ModelWithUUID): if not name: return None - tag, _ = Tag.objects.get_or_create(name=name) + tag, _ = Tag.get_or_create_by_name(name) # Auto-attach to snapshot if in overrides if overrides and "snapshot" in overrides and tag: @@ -606,13 +619,25 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW archiveresult_set: models.Manager["ArchiveResult"] def add_tag_ids(self, tag_ids: Iterable[int | str]) -> None: + tag_ids = [tag_id for tag_id in dict.fromkeys(tag_ids) if tag_id] + for tag_id in tag_ids: + try: + SnapshotTag(snapshot_id=self.pk, tag_id=tag_id).save(force_insert=True) + except IntegrityError: + # Only the unique (snapshot, tag) conflict is idempotent. Do + # not hide foreign-key or other integrity failures. + if SnapshotTag.objects.filter(snapshot_id=self.pk, tag_id=tag_id).exists(): + continue + raise + + def remove_tag_ids(self, tag_ids: Iterable[int | str]) -> int: tag_ids = [tag_id for tag_id in dict.fromkeys(tag_ids) if tag_id] if not tag_ids: - return - SnapshotTag.objects.bulk_create( - [SnapshotTag(snapshot_id=self.pk, tag_id=tag_id) for tag_id in tag_ids], - ignore_conflicts=True, - ) + return 0 + # QuerySet.delete() wraps even a fast through-table DELETE in atomic(). + # SnapshotTag has no delete hooks or child rows, so issue the same + # idempotent DELETE as one autocommit statement. + return SnapshotTag.objects.filter(snapshot_id=self.pk, tag_id__in=tag_ids)._raw_delete(SnapshotTag.objects.db) class Meta( ModelWithDeleteAfter.Meta, @@ -1785,8 +1810,6 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW def _merge_tags_from_index(self, index_data: dict): """Merge tags - union of both sources.""" - from django.db import transaction - index_tags = set(index_data.get("tags", "").split(",")) if index_data.get("tags") else set() index_tags = {t.strip() for t in index_tags if t.strip()} @@ -1794,14 +1817,13 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW new_tags = index_tags - db_tags if new_tags: - with transaction.atomic(): - for tag_name in new_tags: - tag, _ = Tag.objects.get_or_create(name=tag_name) - self.add_tag_ids([tag.pk]) + for tag_name in new_tags: + tag, _ = Tag.get_or_create_by_name(tag_name) + self.add_tag_ids([tag.pk]) def _merge_archive_results_from_index(self, index_data: dict, update_existing: bool = True): - """Merge ArchiveResults one row per plugin; retries update the existing row.""" - existing = {ar.plugin: ar for ar in ArchiveResult.objects.filter(snapshot=self)} + """Merge ArchiveResults one row per hook; retries update the existing row.""" + existing = {(ar.plugin, ar.hook_name): ar for ar in ArchiveResult.objects.filter(snapshot=self)} if update_existing: for archiveresult in existing.values(): normalized_status = ArchiveResult.normalize_status(archiveresult.status) @@ -1864,7 +1886,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW output_mimetypes = result_data.get("output_mimetypes", "") hook_name = result_data.get("hook_name", "") - existing_result = existing.get(plugin) + existing_result = existing.get((plugin, hook_name)) if existing_result: if not update_existing: return @@ -1873,9 +1895,6 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW if existing_result.status != status: existing_result.status = status update_fields.append("status") - if hook_name and existing_result.hook_name != hook_name: - existing_result.hook_name = hook_name - update_fields.append("hook_name") if output_str and existing_result.output_str != output_str: existing_result.output_str = output_str update_fields.append("output_str") @@ -1932,7 +1951,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW end_ts=end_ts, process=process, ) - existing[plugin] = archiveresult + existing[(plugin, hook_name)] = archiveresult def write_index_json(self): """Write index.json in 0.9.x format (deprecated, use write_index_jsonl).""" @@ -2262,15 +2281,18 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW for tag in dup.tags.all(): keeper.add_tag_ids([tag.pk]) - # Move ArchiveResults while preserving the one-row-per-plugin invariant. + # Move each hook result, merging only an exact identity collision. for result in ArchiveResult.objects.filter(snapshot=dup): - existing = ArchiveResult.objects.filter(snapshot=keeper, plugin=result.plugin).first() + existing = ArchiveResult.objects.filter( + snapshot=keeper, + plugin=result.plugin, + hook_name=result.hook_name, + ).first() if existing is None: result.snapshot = keeper result.save(update_fields=["snapshot", "modified_at"]) continue - prefer_result = result.output_size > existing.output_size - prior_output_sizes = (existing.output_size, result.output_size) + output_files = {**(existing.output_files or {}), **(result.output_files or {})} existing.output_files = output_files existing.output_size = max( @@ -2279,13 +2301,15 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW for metadata in output_files.values() if isinstance(metadata, dict) ), - *prior_output_sizes, + existing.output_size, + result.output_size, ) - if prefer_result: + if result.modified_at >= existing.modified_at: existing.status = result.status - existing.hook_name = result.hook_name existing.output_str = result.output_str existing.output_json = result.output_json + existing.start_ts = result.start_ts + existing.end_ts = result.end_ts existing.output_mimetypes = ",".join( sorted( { @@ -2764,10 +2788,13 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW def archive_size(self): return int(self.output_size or 0) - def save_tags(self, tags: Iterable[str] = ()) -> None: - tags_id = [Tag.objects.get_or_create(name=tag)[0].pk for tag in tags if tag.strip()] - self.tags.clear() - self.add_tag_ids(tags_id) + def save_tags(self, tags: Iterable[str] = (), *, created_by: Any = None) -> None: + from archivebox.core.tag_util import get_or_create_tag + + tag_ids = {get_or_create_tag(tag, created_by=created_by)[0].pk for tag in tags if tag.strip()} + existing_tag_ids = set(SnapshotTag.objects.filter(snapshot_id=self.pk).values_list("tag_id", flat=True)) + self.remove_tag_ids(existing_tag_ids - tag_ids) + self.add_tag_ids(tag_ids - existing_tag_ids) def pending_archiveresults(self) -> QuerySet["ArchiveResult"]: return self.archiveresult_set.exclude(status__in=ArchiveResult.FINAL_OR_ACTIVE_STATES) @@ -3035,8 +3062,8 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW - abx_plugins/plugins/*/on_Snapshot__*.{py,sh,js} - data/custom_plugins/*/on_Snapshot__*.{py,sh,js} - Creates one ArchiveResult per plugin. A queued plugin row runs all of that - plugin's Snapshot hooks; hook_name records the most recently completed hook. + Creates one ArchiveResult per hook (not per plugin), with hook_name set. + This enables step-based execution where all hooks in a step can run in parallel. """ try: self.validate_url_for_archiving() @@ -3054,21 +3081,17 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW hooks = ((hook_path.parent.name, hook_path.stem) for hook_path in discover_hooks("Snapshot", config=config)) archiveresults = [] - for plugin in dict.fromkeys(plugin for plugin, _hook_name in hooks): - # ArchiveResult output is one filesystem directory per plugin, so all - # hooks and retries update this row in place instead of creating siblings. - archiveresult, _created = ArchiveResult.objects.get_or_create( - snapshot=self, - plugin=plugin, + for plugin, hook_name in hooks: + # Hooks in one plugin share a filesystem directory, but each hook has + # its own durable result row and retries update that exact row. + archiveresult, _created = ArchiveResult.get_or_create_by_hook( + self, + plugin, + hook_name, defaults={ - "hook_name": "", "status": ArchiveResult.INITIAL_STATE, }, ) - if archiveresult.hook_name == self.BROWSER_EXTENSION_UPLOAD_HOOK_NAME: - archiveresult.hook_name = "" - archiveresult.status = ArchiveResult.INITIAL_STATE - archiveresult.save(update_fields=["hook_name", "status", "modified_at"]) if archiveresult.status == ArchiveResult.INITIAL_STATE: archiveresults.append(archiveresult) @@ -3157,10 +3180,10 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW snapshot=self, status=ArchiveResult.StatusChoices.FAILED, ) + legacy_result_count = retryable_results.filter(hook_name="").count() now = timezone.now() - count = retryable_results.update( + count = retryable_results.exclude(hook_name="").update( status=ArchiveResult.StatusChoices.QUEUED, - hook_name="", output_str="", output_json=None, output_files={}, @@ -3171,11 +3194,11 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW modified_at=now, ) - if count > 0: + if count + legacy_result_count > 0: self.refresh_from_db(fields=["modified_at", "retry_at", "status"]) self.queue_for_extraction(when=now) - return count + return count + legacy_result_count # ========================================================================= # URL Helper Properties (migrated from Link schema) @@ -3849,6 +3872,27 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): "backoff": cls.StatusChoices.BACKOFF, }.get(str(status or "").strip().lower(), cls.StatusChoices.FAILED) + @classmethod + def get_or_create_by_hook( + cls, + snapshot: Snapshot, + plugin: str, + hook_name: str, + *, + defaults: Mapping[str, Any] | None = None, + ) -> tuple["ArchiveResult", bool]: + lookup = {"snapshot": snapshot, "plugin": plugin, "hook_name": hook_name} + result = cls.objects.filter(**lookup).first() + if result: + return result, False + try: + return cls.objects.create(**lookup, **(defaults or {})), True + except IntegrityError: + result = cls.objects.filter(**lookup).first() + if result is None: + raise + return result, False + @staticmethod def output_files_upload_complete(output_files: dict[str, dict[str, Any]]) -> bool: if not output_files: @@ -3983,7 +4027,7 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): models.Index(fields=["-start_ts", "-id"], name="archiveresult_start_idx"), ] constraints: ClassVar[list[models.BaseConstraint]] = [ - models.UniqueConstraint(fields=["snapshot", "plugin"], name="unique_archiveresult_per_snapshot_plugin"), + models.UniqueConstraint(fields=["snapshot", "plugin", "hook_name"], name="unique_archiveresult_per_snapshot_hook"), ] def __str__(self): @@ -4094,16 +4138,16 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): except ArchiveResult.DoesNotExist: pass - # Get or create by snapshot_id + plugin. The filesystem has one output - # directory per plugin, so every hook and retry updates that same DB row. + # Get or create by the durable scheduler identity. Hooks in one plugin + # share an output directory, while retries update only their exact row. try: snapshot = Snapshot.objects.get(id=snapshot_id) - result, _ = ArchiveResult.objects.get_or_create( - snapshot=snapshot, - plugin=plugin, + result, _ = ArchiveResult.get_or_create_by_hook( + snapshot, + plugin, + record.get("hook_name", ""), defaults={ - "hook_name": record.get("hook_name", ""), "status": record.get("status", "queued"), "output_str": record.get("output_str", ""), }, @@ -4167,13 +4211,40 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): ), ) + def safe_update(self, update_fields: Mapping[str, Any], *, refresh: bool = True) -> bool: + """Compare-and-swap one loaded ArchiveResult without opening a transaction.""" + expected_modified_at = self.modified_at + previous_output_size = int(self.output_size or 0) + values = dict(update_fields) + values.setdefault("modified_at", timezone.now()) + updated = type(self).objects.filter(pk=self.pk, modified_at=expected_modified_at).update(**values) + if updated == 1: + for field, value in values.items(): + setattr(self, field, value) + if "output_size" in values: + size_delta = int(values["output_size"] or 0) - previous_output_size + if size_delta: + Snapshot.objects.filter(pk=self.snapshot_id).update( + output_size=F("output_size") + size_delta, + modified_at=timezone.now(), + ) + if refresh: + try: + self.refresh_from_db() + except type(self).DoesNotExist: + pass + return updated == 1 + def schedule_delete_cleanup(self, *, using: str | None = None) -> None: """Remove shared plugin output and refresh persisted Snapshot metadata after commit.""" snapshot_id = self.snapshot_id + plugin = self.plugin paths = self.validate_output_paths_for_delete(self.output_paths_for_delete()) def cleanup() -> None: - type(self).delete_output_paths(paths) + results = type(self).objects.using(using) if using else type(self).objects + if not results.filter(snapshot_id=snapshot_id, plugin=plugin).exists(): + type(self).delete_output_paths(paths) type(self).refresh_snapshot_output_sizes({snapshot_id}) snapshot = Snapshot.objects.filter(pk=snapshot_id).first() if snapshot: @@ -4207,7 +4278,6 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): def reset_for_retry(self, *, save: bool = True) -> None: self.status = self.StatusChoices.QUEUED - self.hook_name = "" self.retry_at = None self.output_str = "" self.output_json = None @@ -4220,7 +4290,6 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): self.save( update_fields=[ "status", - "hook_name", "retry_at", "output_str", "output_json", diff --git a/archivebox/core/recovery_util.py b/archivebox/core/recovery_util.py index 81055b6a..5a3f1266 100644 --- a/archivebox/core/recovery_util.py +++ b/archivebox/core/recovery_util.py @@ -10,6 +10,11 @@ def _is_signal_interrupted_exit(exit_code: int | None) -> bool: return exit_code is not None and (exit_code < 0 or exit_code >= 128) +def _canonical_hook_name(hook_name: str) -> str: + hook_name = Path(hook_name).name + return Path(hook_name).stem if Path(hook_name).suffix in {".py", ".js", ".sh"} else hook_name + + def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | None = None) -> dict[str, int]: from archivebox.crawls.models import Crawl from archivebox.core.models import ArchiveResult, Snapshot @@ -168,14 +173,23 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | orphaned_hook_processes = orphaned_hook_processes.filter(snapshot_pwd_filter) else: orphaned_hook_processes = orphaned_hook_processes.none() - for process in orphaned_hook_processes.only("id", "pwd", "cmd", "process_type", "status", "started_at", "ended_at").order_by( - "-started_at", - "-id", - ): + for process in orphaned_hook_processes.only( + "id", + "pwd", + "cmd", + "process_type", + "status", + "exit_code", + "stdout", + "stderr", + "started_at", + "ended_at", + ).order_by("-started_at", "-id"): hook_script_name = process.hook_script_name if not hook_script_name or not process.pwd: continue plugin_dir = Path(process.pwd) + hook_name = _canonical_hook_name(hook_script_name) if crawl_snapshot_ids is not None and plugin_dir.parent.name not in crawl_snapshot_ids: continue try: @@ -187,19 +201,20 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | continue if snapshot is None: continue - result, created = ArchiveResult.objects.get_or_create( - snapshot=snapshot, - plugin=plugin_dir.name, + result, created = ArchiveResult.get_or_create_by_hook( + snapshot, + plugin_dir.name, + hook_name, defaults={ - "hook_name": Path(hook_script_name).stem, "status": ArchiveResult.StatusChoices.QUEUED, }, ) - if created or result.status == ArchiveResult.StatusChoices.QUEUED: + process_is_newer = bool(process.started_at and (result.start_ts is None or process.started_at >= result.start_ts)) + if result.status == ArchiveResult.StatusChoices.QUEUED or process_is_newer: requeue_snapshot = False # A runner can die after the hook Process exits but before the # ProcessCompletedEvent projector links/finalizes ArchiveResult. - # Reconstruct the plugin row from its newest durable Process row. + # Reconstruct only that exact hook row from the durable Process row. manifest = OutputManifest.scan(plugin_dir, containment_root=snapshot.output_dir) output_files = manifest.as_mapping() output_size = manifest.total_size @@ -209,10 +224,9 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | for record in Process.parse_records_from_text(process.stdout or "") if record.get("type") == "ArchiveResult" and (record.get("plugin") or plugin_dir.name) == plugin_dir.name - and (record.get("hook_name") or Path(hook_script_name).stem) == Path(hook_script_name).stem + and _canonical_hook_name(str(record.get("hook_name") or hook_name)) == hook_name ] emitted_result = emitted_records[-1] if emitted_records else {} - result.hook_name = Path(hook_script_name).stem result.process = process result.start_ts = process.started_at result.end_ts = process.ended_at @@ -225,6 +239,7 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | result.output_size = 0 result.output_mimetypes = "" result.output_str = "" + result.output_json = None result.status = ArchiveResult.StatusChoices.QUEUED requeue_snapshot = True else: @@ -249,7 +264,6 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | ) result.save( update_fields=[ - "hook_name", "process", "start_ts", "end_ts", diff --git a/archivebox/core/tag_util.py b/archivebox/core/tag_util.py index 9fbeea22..4bbbf107 100644 --- a/archivebox/core/tag_util.py +++ b/archivebox/core/tag_util.py @@ -152,15 +152,8 @@ def get_or_create_tag(name: str, created_by: User | None = None) -> tuple[Tag, b if not normalized_name: raise ValueError("Tag name is required") - existing = Tag.objects.filter(name__iexact=normalized_name).first() - if existing: - return existing, False - - tag = Tag.objects.create( - name=normalized_name, - created_by=created_by, - ) - return tag, True + defaults = {"created_by": created_by} if created_by is not None else None + return Tag.get_or_create_by_name(normalized_name, defaults=defaults) def rename_tag(tag: Tag, name: str) -> Tag: diff --git a/archivebox/core/views.py b/archivebox/core/views.py index c6849e91..e3b71e16 100644 --- a/archivebox/core/views.py +++ b/archivebox/core/views.py @@ -691,26 +691,27 @@ def _resolve_archiveresult_relpath(snapshot: Snapshot, rel_path: str) -> tuple[s plugin = parts[0] plugin_relpath = posixpath.join(*parts[1:]) - result = ( - ArchiveResult.objects.filter(snapshot=snapshot, plugin=plugin, status=ArchiveResult.StatusChoices.SUCCEEDED) - .only("plugin", "output_files") - .first() + results = list( + ArchiveResult.objects.filter( + snapshot=snapshot, + plugin=plugin, + status=ArchiveResult.StatusChoices.SUCCEEDED, + ).only("plugin", "output_files"), ) - if not result: + if not results: return rel_path, None - if not result.output_files: - return rel_path, result - output_files = result.output_files or {} - for candidate in (plugin_relpath, rel_path): - file_info = output_files.get(candidate) - if not isinstance(file_info, dict): - continue - if file_info.get("root_relative"): - return candidate, result - return rel_path, result + for result in results: + output_files = result.output_files or {} + for candidate in (plugin_relpath, rel_path): + file_info = output_files.get(candidate) + if not isinstance(file_info, dict): + continue + if file_info.get("root_relative"): + return candidate, result + return rel_path, result - return rel_path, result + return rel_path, results[0] def _plugin_full_preview_response( diff --git a/archivebox/crawls/models.py b/archivebox/crawls/models.py index f858ec3e..c8a46841 100755 --- a/archivebox/crawls/models.py +++ b/archivebox/crawls/models.py @@ -1191,7 +1191,7 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith ), } if tag_names: - tag_ids = [Tag.objects.get_or_create(name=tag_name)[0].pk for tag_name in tag_names] + tag_ids = [Tag.get_or_create_by_name(tag_name)[0].pk for tag_name in tag_names] snapshot.add_tag_ids(tag_ids) existing_scope = Snapshot.objects if bool(self._config_value(config, "ONLY_NEW", True)) else self.snapshot_set diff --git a/archivebox/services/archive_result_service.py b/archivebox/services/archive_result_service.py index 87d77ae1..afa157dc 100644 --- a/archivebox/services/archive_result_service.py +++ b/archivebox/services/archive_result_service.py @@ -4,7 +4,6 @@ import asyncio import inspect import json import os -import re import signal import sys import time @@ -14,7 +13,6 @@ from pathlib import Path from typing import Any from asgiref.sync import sync_to_async -from django.db import IntegrityError from django.utils import timezone from abx_dl.events import PROCESS_EXIT_SKIPPED, ArchiveResultEvent, ProcessCompletedEvent, ProcessStartedEvent, SnapshotEvent @@ -87,11 +85,6 @@ def _normalize_status(status: str) -> str: return status or "failed" -def _snapshot_hook_order(hook_name: str) -> int: - match = re.match(r"^on_Snapshot__(\d+)", hook_name or "") - return int(match.group(1)) if match else -1 - - def _normalize_snapshot_title(candidate: str, *, snapshot_url: str) -> str: title = " ".join(line.strip() for line in str(candidate or "").splitlines() if line.strip()).strip() if not title: @@ -209,7 +202,6 @@ def _save_archiveresult_event_to_db( start_ts = parse_event_datetime(event.start_ts) end_ts = parse_event_datetime(event.end_ts) or timezone.now() defaults = { - "hook_name": event.hook_name, "status": _normalize_status(event.status), "output_str": event.output_str, "output_json": event.output_json, @@ -224,48 +216,13 @@ def _save_archiveresult_event_to_db( if event.error: defaults["notes"] = event.error - with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_lookup"): - result = ArchiveResult.objects.filter( - snapshot=snapshot, - plugin=event.plugin, - ).first() - if result is None: - try: - with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_create"): - result = ArchiveResult.objects.create( - snapshot=snapshot, - plugin=event.plugin, - **defaults, - ) - except IntegrityError: - with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_get_after_integrity"): - result = ArchiveResult.objects.get( - snapshot=snapshot, - plugin=event.plugin, - ) - - if result.output_files: - merged_output_files = {**result.output_files, **defaults["output_files"]} - defaults["output_files"], defaults["output_size"], defaults["output_mimetypes"] = _manifest_metadata( - OutputManifest.from_value(merged_output_files), + with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_get_or_create"): + result, _created = ArchiveResult.get_or_create_by_hook( + snapshot, + event.plugin, + event.hook_name, + defaults=defaults, ) - defaults["output_size"] = max(defaults["output_size"], int(result.output_size or 0)) - defaults["output_mimetypes"] = ",".join( - dict.fromkeys( - mimetype.strip() - for value in (defaults["output_mimetypes"], result.output_mimetypes) - for mimetype in value.split(",") - if mimetype.strip() - ), - ) - if result.start_ts and defaults["start_ts"]: - defaults["start_ts"] = min(result.start_ts, defaults["start_ts"]) - if result.end_ts and defaults["end_ts"]: - defaults["end_ts"] = max(result.end_ts, defaults["end_ts"]) - if _snapshot_hook_order(result.hook_name) > _snapshot_hook_order(event.hook_name): - for field in ("hook_name", "status", "output_str", "output_json", "process_id", "notes"): - if field in result.__dict__: - defaults[field] = result.__dict__[field] with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.diff_fields"): update_fields = [] @@ -307,7 +264,7 @@ def _save_archiveresult_event_to_db( def mark_archiveresult_started(event: ProcessStartedEvent, *, snapshot_id: str, process_id: str) -> None: - """Advance an existing queued plugin row after its OS process is persisted.""" + """Advance an existing queued hook row after its OS process is persisted.""" from archivebox.core.models import ArchiveResult started_at = parse_event_datetime(event.start_ts) @@ -316,9 +273,9 @@ def mark_archiveresult_started(event: ProcessStartedEvent, *, snapshot_id: str, ArchiveResult.objects.filter( snapshot_id=snapshot_id, plugin=event.plugin_name, + hook_name=event.hook_name, status=ArchiveResult.StatusChoices.QUEUED, ).update( - hook_name=event.hook_name, status=ArchiveResult.StatusChoices.STARTED, start_ts=started_at, end_ts=None, @@ -333,7 +290,7 @@ class ArchiveResultService(BaseService): def __init__(self, bus): self._completed_process_event_ids: set[str] = set() - self._save_locks: dict[tuple[str, str], asyncio.Lock] = {} + self._save_locks: dict[tuple[str, str, str], asyncio.Lock] = {} super().__init__(bus) self.bus.on(ArchiveResultEvent, self.on_ArchiveResultEvent__save_to_db) self.bus.on(ProcessCompletedEvent, self.on_ProcessCompletedEvent__save_to_db) @@ -348,7 +305,7 @@ class ArchiveResultService(BaseService): where=lambda candidate: self.bus.event_is_child_of(event, candidate), ) - key = (str(event.snapshot_id), event.plugin) + key = (str(event.snapshot_id), event.plugin, event.hook_name) lock = self._save_locks.setdefault(key, asyncio.Lock()) async with lock: await sync_to_async(_save_archiveresult_event_to_db, thread_sensitive=True)(event, process_started) diff --git a/archivebox/services/runner.py b/archivebox/services/runner.py index ebdab9cc..89b4cbdc 100644 --- a/archivebox/services/runner.py +++ b/archivebox/services/runner.py @@ -1046,7 +1046,7 @@ class CrawlRunner: ) return [plugin for plugin in queued_plugins if plugin in expanded_selected_plugins] - queued_plugins = None + selected_hooks_by_plugin = None if snapshot["status"] == "started": _reset_count, running_count = await sync_to_async(snapshot["_snapshot"].reset_abandoned_results, thread_sensitive=True)() if running_count: @@ -1064,8 +1064,8 @@ class CrawlRunner: ) return if not self.selected_plugins_from_args: - queued_plugins = await sync_to_async( - queued_plugins_for_snapshot, + queued_plugins, selected_hooks_by_plugin = await sync_to_async( + queued_plugins_and_hooks_for_snapshot, thread_sensitive=True, )(snapshot["id"]) if queued_plugins: @@ -1078,10 +1078,13 @@ class CrawlRunner: snapshot["id"], disabled_queued_plugins, ) + selected_hooks_by_plugin = { + plugin: hooks for plugin, hooks in (selected_hooks_by_plugin or {}).items() if plugin in queued_plugins + } snapshot_selected_plugins = queued_plugins elif not self.selected_plugins_from_args: - queued_plugins = await sync_to_async( - queued_plugins_for_snapshot, + queued_plugins, selected_hooks_by_plugin = await sync_to_async( + queued_plugins_and_hooks_for_snapshot, thread_sensitive=True, )(snapshot["id"]) if queued_plugins: @@ -1094,6 +1097,9 @@ class CrawlRunner: snapshot["id"], disabled_queued_plugins, ) + selected_hooks_by_plugin = { + plugin: hooks for plugin, hooks in (selected_hooks_by_plugin or {}).items() if plugin in queued_plugins + } snapshot_selected_plugins = queued_plugins if snapshot["depth"] > 0 and CrawlLimitState.from_config(snapshot["config"]).get_stop_reason() in ( "crawl_max_size", @@ -1108,10 +1114,10 @@ class CrawlRunner: if snapshot_selected_plugins else self.plugins ) - if queued_plugins is not None: - await sync_to_async(fail_unavailable_queued_plugins, thread_sensitive=True)( + if selected_hooks_by_plugin is not None: + await sync_to_async(fail_unavailable_queued_hooks, thread_sensitive=True)( snapshot["id"], - queued_plugins, + selected_hooks_by_plugin, plugins, ) remaining_queued_plugins = await sync_to_async( @@ -1131,6 +1137,7 @@ class CrawlRunner: return snapshot_selected_plugins = remaining_queued_plugins plugins = filter_plugins(self.plugins, snapshot_selected_plugins, include_providers=True) + selected_hooks_by_plugin = include_background_prerequisite_hooks(selected_hooks_by_plugin, plugins) abx_snapshot = AbxSnapshot( id=snapshot["id"], url=snapshot["url"], @@ -1155,7 +1162,7 @@ class CrawlRunner: snapshot_service=HookSnapshotService, timeout_padding=120.0, abort_requested=self.crawl_is_cancelled, - selected_hooks_by_plugin=None, + selected_hooks_by_plugin=selected_hooks_by_plugin, emit_discovered_snapshot_events=False, ) try: @@ -1363,19 +1370,37 @@ def run_binary(binary_id: str) -> None: asyncio.run(_run_binary(binary_id)) -def queued_plugins_for_snapshot(snapshot_id: str) -> list[str] | None: +def queued_plugins_and_hooks_for_snapshot(snapshot_id: str) -> tuple[list[str] | None, dict[str, set[str] | None] | None]: from archivebox.core.models import ArchiveResult - queued_plugins = list( + queued_results = list( ArchiveResult.objects.filter( snapshot_id=snapshot_id, status=ArchiveResult.StatusChoices.QUEUED, ) .exclude(plugin="") - .order_by("plugin") - .values_list("plugin", flat=True), + .only("id", "plugin", "hook_name"), ) - return queued_plugins or None + + selected_hooks_by_plugin: dict[str, set[str] | None] = {} + queued_plugins = sorted({result.plugin for result in queued_results}) + for result in queued_results: + # hook_name is the modern scheduler identity. Empty hook_name rows are + # legacy plugin-level work and must keep running the whole plugin. + if not result.hook_name: + selected_hooks_by_plugin[result.plugin] = None + elif result.plugin not in selected_hooks_by_plugin: + selected_hooks_by_plugin[result.plugin] = {result.hook_name} + elif selected_hooks_by_plugin[result.plugin] is not None: + selected_hooks_by_plugin[result.plugin].add(result.hook_name) + if queued_plugins: + return queued_plugins, selected_hooks_by_plugin + return None, None + + +def queued_plugins_for_snapshot(snapshot_id: str) -> list[str] | None: + queued_plugins, _selected_hooks_by_plugin = queued_plugins_and_hooks_for_snapshot(snapshot_id) + return queued_plugins def config_overrides_for_queued_plugins(selected_plugins: list[str], **overrides: Any) -> dict[str, Any]: @@ -1392,27 +1417,40 @@ def config_overrides_for_queued_plugins(selected_plugins: list[str], **overrides return config_overrides -def fail_unavailable_queued_plugins( +def fail_unavailable_queued_hooks( snapshot_id: str, - queued_plugins: list[str], + selected_hooks_by_plugin: dict[str, set[str] | None], plugins: dict[str, Plugin], ) -> None: from archivebox.core.models import ArchiveResult - unavailable_plugins = set(queued_plugins) - set(plugins) - if not unavailable_plugins: - return now = timezone.now() - ArchiveResult.objects.filter( - snapshot_id=snapshot_id, - plugin__in=unavailable_plugins, - status=ArchiveResult.StatusChoices.QUEUED, - ).update( - status=ArchiveResult.StatusChoices.FAILED, - start_ts=now, - end_ts=now, - output_str="Queued plugin is not available in this ArchiveBox installation", - ) + for plugin_name, selected_hook_names in selected_hooks_by_plugin.items(): + if selected_hook_names is None: + continue + if plugin_name in plugins: + available_hook_names = { + name for hook in plugins[plugin_name].filter_hooks("Snapshot") for name in (hook.name, Path(hook.name).stem) + } + else: + available_hook_names = set() + missing_hook_names = [hook_name for hook_name in selected_hook_names if hook_name not in available_hook_names] + if not missing_hook_names: + continue + # Hook-level resume rows are durable scheduler state. If a plugin is + # installed but no longer exposes a queued hook, mark that row failed so + # the snapshot is not retried forever with no hook left to execute. + ArchiveResult.objects.filter( + snapshot_id=snapshot_id, + plugin=plugin_name, + hook_name__in=missing_hook_names, + status=ArchiveResult.StatusChoices.QUEUED, + ).update( + status=ArchiveResult.StatusChoices.FAILED, + start_ts=now, + end_ts=now, + output_str="Queued hook is no longer available in the installed plugin", + ) def skip_disabled_queued_plugins(snapshot_id: str, plugin_names: list[str]) -> None: @@ -1437,6 +1475,35 @@ def skip_disabled_queued_plugins(snapshot_id: str, plugin_names: list[str]) -> N ) +def include_background_prerequisite_hooks( + selected_hooks_by_plugin: dict[str, set[str] | None], + plugins: dict[str, Plugin], +) -> dict[str, set[str] | None]: + expanded: dict[str, set[str] | None] = {} + for plugin_name, selected_hook_names in selected_hooks_by_plugin.items(): + if selected_hook_names is None or plugin_name not in plugins: + expanded[plugin_name] = selected_hook_names + continue + plugin_hooks = sorted(plugins[plugin_name].filter_hooks("Snapshot"), key=lambda hook: hook.sort_key) + selected_sort_keys = [ + hook.sort_key for hook in plugin_hooks if hook.name in selected_hook_names or Path(hook.name).stem in selected_hook_names + ] + if not selected_sort_keys: + expanded[plugin_name] = set(selected_hook_names) + continue + first_selected_sort_key = min(selected_sort_keys) + expanded_hook_names = set(selected_hook_names) + # Earlier background hooks publish live resources (e.g. Chrome tabs) + # needed by later foreground hooks, but completed foreground hooks stay + # final and are not rerun during hook-level resume. + for hook in plugin_hooks: + if hook.is_background and hook.sort_key < first_selected_sort_key: + expanded_hook_names.add(hook.name) + expanded_hook_names.add(Path(hook.name).stem) + expanded[plugin_name] = expanded_hook_names + return expanded + + def snapshot_hooks_for_pending_archiveresults(snapshot) -> list[tuple[str, str]]: from archivebox.config.common import get_config from archivebox.core.models import Snapshot @@ -1697,10 +1764,10 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru return False snapshot.refresh_from_db() if snapshot.status == Snapshot.StatusChoices.QUEUED: - has_server_workset = snapshot.archiveresult_set.exclude( + has_server_archiveresults = snapshot.archiveresult_set.exclude( hook_name=Snapshot.BROWSER_EXTENSION_UPLOAD_HOOK_NAME, ).exists() - if has_server_workset and snapshot.is_finished_processing(): + if has_server_archiveresults and snapshot.is_finished_processing(): finalize_completed_snapshot(str(snapshot.id), output_dir=Path(snapshot.output_dir)) snapshot.refresh_from_db() if snapshot.status == Snapshot.StatusChoices.SEALED: @@ -1709,10 +1776,10 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru snapshot.refresh_from_db() _runner_console_line(crawl_id=snapshot.crawl_id, snapshot=snapshot, status="SEALED") return True - if not has_server_workset: - # Materialize one durable row per configured plugin. Existing - # browser-uploaded rows are reused and queued so the server adds its - # outputs to that plugin result instead of creating a sibling row. + # A Snapshot with no server hook rows is fresh lifecycle work; materialize + # its configured hook set. Browser-extension uploads are completed external + # outputs, not the durable server workset, and must coexist with these rows. + if not has_server_archiveresults: snapshot.create_pending_archiveresults(hooks=snapshot_hooks_for_pending_archiveresults(snapshot)) snapshot.advance_lifecycle() snapshot.refresh_from_db() @@ -1729,11 +1796,11 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru _runner_console_line(crawl_id=snapshot.crawl_id, snapshot=snapshot, status="SEALED") return True if snapshot.status == Snapshot.StatusChoices.STARTED: - queued_plugins = queued_plugins_for_snapshot(str(snapshot.id)) - if queued_plugins: - fail_unavailable_queued_plugins( + queued_plugins, selected_hooks_by_plugin = queued_plugins_and_hooks_for_snapshot(str(snapshot.id)) + if queued_plugins and selected_hooks_by_plugin: + fail_unavailable_queued_hooks( str(snapshot.id), - queued_plugins, + selected_hooks_by_plugin, _discover_archivebox_plugins(), ) if not queued_plugins_for_snapshot(str(snapshot.id)): @@ -1757,10 +1824,10 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru ) snapshot.refresh_from_db() if queued_plugins_for_snapshot(str(snapshot.id)): - # Plugin-level resume work is tracked by queued ArchiveResult rows, not by + # Hook-level resume work is tracked by queued ArchiveResult rows, not by # the Snapshot lease. If a partial pass returns with rows still queued, # wake the Snapshot immediately so takeover does not wait out a stale - # active-state lock before running the remaining plugins. + # active-state lock before running the remaining hooks. snapshot.update_and_requeue(retry_at=timezone.now()) return True diff --git a/archivebox/services/tag_service.py b/archivebox/services/tag_service.py index 22d6685d..2844a3f8 100644 --- a/archivebox/services/tag_service.py +++ b/archivebox/services/tag_service.py @@ -1,5 +1,7 @@ from __future__ import annotations +from asgiref.sync import sync_to_async + from abx_dl.events import TagEvent from abx_dl.services.base import BaseService @@ -13,10 +15,10 @@ class TagService(BaseService): self.bus.on(TagEvent, self.on_TagEvent__save_to_db) async def on_TagEvent__save_to_db(self, event: TagEvent) -> None: - from archivebox.core.models import Snapshot, SnapshotTag, Tag + from archivebox.core.models import Snapshot, Tag snapshot = await Snapshot.objects.filter(id=event.snapshot_id).afirst() if snapshot is None: return - tag, _ = await Tag.objects.aget_or_create(name=event.name) - await SnapshotTag.objects.aget_or_create(snapshot=snapshot, tag=tag) + tag, _ = await sync_to_async(Tag.get_or_create_by_name, thread_sensitive=True)(event.name) + await sync_to_async(snapshot.add_tag_ids, thread_sensitive=True)([tag.pk]) diff --git a/archivebox/tests/test_api_v1_core_archiveresults.py b/archivebox/tests/test_api_v1_core_archiveresults.py index 1718d8b0..7380dbe7 100644 --- a/archivebox/tests/test_api_v1_core_archiveresults.py +++ b/archivebox/tests/test_api_v1_core_archiveresults.py @@ -3,6 +3,7 @@ from datetime import timedelta import pytest from django.core.files.uploadedfile import SimpleUploadedFile from django.db import connection +from django.test.client import BOUNDARY, MULTIPART_CONTENT, encode_multipart from django.test.utils import CaptureQueriesContext from django.utils import timezone @@ -14,7 +15,7 @@ from archivebox.tests.conftest import api_client_request pytestmark = pytest.mark.django_db(transaction=True) -def test_archiveresult_upload_upserts_one_snapshot_plugin_result(client, api_admin_user, api_headers): +def test_archiveresult_upload_upserts_by_snapshot_plugin_and_hook(client, api_admin_user, api_headers): crawl = Crawl.objects.create(urls="https://example.com", created_by=api_admin_user) snapshot = Snapshot.objects.create(url="https://example.com/unified-result", crawl=crawl) @@ -39,13 +40,130 @@ def test_archiveresult_upload_upserts_one_snapshot_plugin_result(client, api_adm assert server_response.status_code == 200, server_response.content assert extension_update.status_code == 200, extension_update.content assert extension_update.json()["id"] == extension_response.json()["id"] - assert server_response.json()["id"] == extension_response.json()["id"] + assert server_response.json()["id"] != extension_response.json()["id"] - results = ArchiveResult.objects.filter(snapshot=snapshot, plugin="screenshot") - assert results.count() == 1 - result = results.get() - assert result.hook_name == "on_Snapshot__archivebox_browser_extension_upload" - assert set(result.output_files) == {"browser.png", "server.png", "browser-2.png"} + results = ArchiveResult.objects.filter(snapshot=snapshot, plugin="screenshot").order_by("hook_name") + assert results.count() == 2 + extension_result = results.get(hook_name="on_Snapshot__archivebox_browser_extension_upload") + server_result = results.get(hook_name="on_Snapshot__50_screenshot") + assert set(extension_result.output_files) == {"browser.png", "browser-2.png"} + assert set(server_result.output_files) == {"server.png"} + snapshot.refresh_from_db() + assert snapshot.output_size == len(b"browser") + len(b"browser-2") + len(b"server") + + +def test_archiveresult_create_does_not_open_a_database_transaction(client, api_admin_user, api_headers): + crawl = Crawl.objects.create(urls="https://example.com", created_by=api_admin_user) + snapshot = Snapshot.objects.create(url="https://example.com/autocommit-result", crawl=crawl) + + with CaptureQueriesContext(connection) as queries: + response = client.post( + "/api/v1/core/archiveresults", + { + "snapshot_id": str(snapshot.id), + "plugin": "chrome_extension_screenshot", + "files": SimpleUploadedFile("screenshot.png", b"screenshot", content_type="image/png"), + "output_paths": "screenshot.png", + }, + **api_headers, + ) + + assert response.status_code == 200, response.content + if connection.vendor == "sqlite": + transaction_queries = [query["sql"] for query in queries if query["sql"].strip().upper() in {"BEGIN", "COMMIT"}] + assert transaction_queries == [] + + +def test_intermediate_archiveresult_chunks_only_write_to_disk(client, api_admin_user, api_headers): + crawl = Crawl.objects.create(urls="https://example.com", created_by=api_admin_user) + snapshot = Snapshot.objects.create(url="https://example.com/chunked-result", crawl=crawl) + create_response = client.post( + "/api/v1/core/archiveresults", + { + "snapshot_id": str(snapshot.id), + "plugin": "chrome_extension_mhtml", + "status": ArchiveResult.StatusChoices.STARTED, + }, + **api_headers, + ) + assert create_response.status_code == 200, create_response.content + result = ArchiveResult.objects.get(pk=create_response.json()["id"]) + original_modified_at = result.modified_at + + with CaptureQueriesContext(connection) as queries: + chunk_response = client.patch( + f"/api/v1/core/archiveresult/{result.id}", + encode_multipart( + BOUNDARY, + { + "files": SimpleUploadedFile("snapshot.mhtml.part-000000", b"first", content_type="multipart/related"), + "chunk_output_path": "snapshot.mhtml", + "chunk_index": "0", + "chunk_count": "2", + "chunk_offset": "0", + "chunk_total_size": "11", + "mime_type": "multipart/related", + "status": ArchiveResult.StatusChoices.STARTED, + }, + ), + content_type=MULTIPART_CONTENT, + **api_headers, + ) + + assert chunk_response.status_code == 200, chunk_response.content + result.refresh_from_db() + assert result.output_files == {} + assert result.output_size == 0 + assert result.modified_at == original_modified_at + writes = [query["sql"] for query in queries if query["sql"].lstrip().upper().startswith(("INSERT", "UPDATE", "DELETE"))] + assert writes == [] + + final_response = client.patch( + f"/api/v1/core/archiveresult/{result.id}", + encode_multipart( + BOUNDARY, + { + "files": SimpleUploadedFile("snapshot.mhtml.part-000001", b"-final", content_type="multipart/related"), + "chunk_output_path": "snapshot.mhtml", + "chunk_index": "1", + "chunk_count": "2", + "chunk_offset": "5", + "chunk_total_size": "11", + "mime_type": "multipart/related", + "status": ArchiveResult.StatusChoices.SUCCEEDED, + }, + ), + content_type=MULTIPART_CONTENT, + **api_headers, + ) + assert final_response.status_code == 200, final_response.content + result.refresh_from_db() + assert result.status == ArchiveResult.StatusChoices.SUCCEEDED + assert result.output_size == 11 + assert result.output_files["snapshot.mhtml"]["upload"]["complete"] is True + snapshot.refresh_from_db() + assert snapshot.output_size == 11 + + +def test_archiveresult_safe_update_rejects_stale_writers(api_admin_user): + crawl = Crawl.objects.create(urls="https://example.com", created_by=api_admin_user) + snapshot = Snapshot.objects.create(url="https://example.com/cas-result", crawl=crawl) + result = ArchiveResult.objects.create( + snapshot=snapshot, + plugin="chrome_extension_dom", + hook_name=Snapshot.BROWSER_EXTENSION_UPLOAD_HOOK_NAME, + status=ArchiveResult.StatusChoices.STARTED, + output_size=1, + ) + stale_result = ArchiveResult.objects.get(pk=result.pk) + + assert result.safe_update({"output_size": 2}) is True + assert stale_result.safe_update({"output_size": 3}) is False + + stale_result.refresh_from_db() + snapshot.refresh_from_db() + assert stale_result.output_size == 2 + assert snapshot.output_size == 2 def test_archiveresult_upload_api_queues_snapshot_maintenance_without_finalizing(client, api_admin_user, api_headers): diff --git a/archivebox/tests/test_api_v1_core_snapshots.py b/archivebox/tests/test_api_v1_core_snapshots.py index 30e95a7e..771ca85e 100644 --- a/archivebox/tests/test_api_v1_core_snapshots.py +++ b/archivebox/tests/test_api_v1_core_snapshots.py @@ -2,6 +2,8 @@ import threading import time import pytest +from django.db import connection +from django.test.utils import CaptureQueriesContext from archivebox.core.models import Snapshot from archivebox.crawls.locks import crawl_lifecycle_lock @@ -85,3 +87,66 @@ def test_existing_snapshot_metadata_sync_does_not_wait_for_active_crawl(client, assert response.status_code == 200, response.content assert response.json()["id"] == str(snapshot.id) assert elapsed < 1 + + +def test_new_snapshot_creation_does_not_open_a_database_transaction(client, api_admin_user, api_headers): + url = "https://example.com/browser-extension-new-snapshot" + crawl = Crawl.objects.create(urls=url, created_by=api_admin_user) + + with CaptureQueriesContext(connection) as queries: + response = client.post( + "/api/v1/core/snapshots", + data={ + "url": url, + "crawl_id": str(crawl.id), + "depth": 0, + "status": Snapshot.StatusChoices.QUEUED, + "tags": ["browser-extension-upload"], + }, + content_type="application/json", + **api_headers, + ) + + assert response.status_code == 200, response.content + assert Snapshot.objects.filter(url=url, crawl=crawl).count() == 1 + assert Snapshot.objects.get(url=url, crawl=crawl).tags.filter(name="browser-extension-upload").exists() + if connection.vendor == "sqlite": + transaction_queries = [query["sql"] for query in queries if query["sql"].strip().upper() in {"BEGIN", "COMMIT"}] + assert transaction_queries == [] + + +def test_new_snapshot_creation_does_not_wait_for_active_crawl(client, api_admin_user, api_headers): + url = "https://example.com/browser-extension-new-snapshot-active-crawl" + crawl = Crawl.objects.create(urls=url, created_by=api_admin_user) + lock_acquired = threading.Event() + release_lock = threading.Event() + + def hold_active_crawl_lock(): + with crawl_lifecycle_lock(str(crawl.id)): + lock_acquired.set() + release_lock.wait(timeout=3) + + holder = threading.Thread(target=hold_active_crawl_lock) + holder.start() + assert lock_acquired.wait(timeout=1) + + started_at = time.monotonic() + response = client.post( + "/api/v1/core/snapshots", + data={ + "url": url, + "crawl_id": str(crawl.id), + "depth": 0, + "status": Snapshot.StatusChoices.QUEUED, + "tags": ["browser-extension-upload"], + }, + content_type="application/json", + **api_headers, + ) + elapsed = time.monotonic() - started_at + release_lock.set() + holder.join(timeout=3) + + assert response.status_code == 200, response.content + assert Snapshot.objects.filter(url=url, crawl=crawl).count() == 1 + assert elapsed < 1 diff --git a/archivebox/tests/test_api_v1_core_tags_add_to_snapshot.py b/archivebox/tests/test_api_v1_core_tags_add_to_snapshot.py index c732ca39..6f8a1ff0 100644 --- a/archivebox/tests/test_api_v1_core_tags_add_to_snapshot.py +++ b/archivebox/tests/test_api_v1_core_tags_add_to_snapshot.py @@ -1,4 +1,6 @@ import pytest +from django.db import IntegrityError, connection +from django.test.utils import CaptureQueriesContext from archivebox.core.models import Snapshot, Tag from archivebox.crawls.models import Crawl @@ -13,13 +15,39 @@ def test_basic_success_case_request(client, api_admin_user, api_headers): snapshot = Snapshot.objects.create(url="https://example.com/tag-add", crawl=crawl) tag = Tag.objects.create(name="api-basic-add-tag", created_by=api_admin_user) - response = api_client_request( - client, - "post", - "/api/v1/core/tags/add-to-snapshot/", - payload={"snapshot_id": str(snapshot.id), "tag_id": tag.id}, - headers=api_headers, - ) + with CaptureQueriesContext(connection) as queries: + response = api_client_request( + client, + "post", + "/api/v1/core/tags/add-to-snapshot/", + payload={"snapshot_id": str(snapshot.id), "tag_id": tag.id}, + headers=api_headers, + ) assert response.status_code == 200, response.content assert response.json()["success"] is True + assert snapshot.tags.filter(pk=tag.pk).exists() + if connection.vendor == "sqlite": + transaction_queries = [query["sql"] for query in queries if query["sql"].strip().upper() in {"BEGIN", "COMMIT"}] + assert transaction_queries == [] + + +def test_add_tag_ids_reraises_non_duplicate_integrity_errors(snapshot): + with pytest.raises(IntegrityError): + snapshot.add_tag_ids([2**31 - 1]) + + +def test_add_tag_ids_treats_existing_snapshot_tag_as_idempotent(snapshot, admin_user): + tag = Tag.objects.create(name="already-attached", created_by=admin_user) + + snapshot.add_tag_ids([tag.pk]) + snapshot.add_tag_ids([tag.pk]) + + assert snapshot.tags.filter(pk=tag.pk).count() == 1 + + +def test_save_tags_without_creator_uses_tag_model_default(snapshot): + snapshot.save_tags(["model-default-creator"]) + + tag = Tag.objects.get(name="model-default-creator") + assert tag.created_by_id is not None diff --git a/archivebox/tests/test_api_v1_core_tags_remove_from_snapshot.py b/archivebox/tests/test_api_v1_core_tags_remove_from_snapshot.py index 2b56770c..94614fe0 100644 --- a/archivebox/tests/test_api_v1_core_tags_remove_from_snapshot.py +++ b/archivebox/tests/test_api_v1_core_tags_remove_from_snapshot.py @@ -1,4 +1,6 @@ import pytest +from django.db import connection +from django.test.utils import CaptureQueriesContext from archivebox.core.models import Snapshot, Tag from archivebox.crawls.models import Crawl @@ -14,13 +16,18 @@ def test_basic_success_case_request(client, api_admin_user, api_headers): tag = Tag.objects.create(name="api-basic-remove-tag", created_by=api_admin_user) snapshot.tags.add(tag) - response = api_client_request( - client, - "post", - "/api/v1/core/tags/remove-from-snapshot/", - payload={"snapshot_id": str(snapshot.id), "tag_id": tag.id}, - headers=api_headers, - ) + with CaptureQueriesContext(connection) as queries: + response = api_client_request( + client, + "post", + "/api/v1/core/tags/remove-from-snapshot/", + payload={"snapshot_id": str(snapshot.id), "tag_id": tag.id}, + headers=api_headers, + ) assert response.status_code == 200, response.content assert response.json()["success"] is True + assert not snapshot.tags.filter(pk=tag.pk).exists() + if connection.vendor == "sqlite": + transaction_queries = [query["sql"] for query in queries if query["sql"].strip().upper() in {"BEGIN", "COMMIT"}] + assert transaction_queries == [] diff --git a/archivebox/tests/test_archive_result_service.py b/archivebox/tests/test_archive_result_service.py index 8991d749..7dc21ff6 100644 --- a/archivebox/tests/test_archive_result_service.py +++ b/archivebox/tests/test_archive_result_service.py @@ -230,7 +230,7 @@ def test_archiveresult_event_retry_updates_existing_hook_row(tmp_path, hermetic_ _cleanup_machine_process_rows() -def test_archiveresult_duplicate_plugin_rows_are_rejected(): +def test_archiveresult_duplicate_hook_rows_are_rejected(): from django.db import IntegrityError, transaction from archivebox.core.models import ArchiveResult @@ -246,45 +246,35 @@ def test_archiveresult_duplicate_plugin_rows_are_rejected(): ArchiveResult.objects.create( snapshot=snapshot, plugin="wget", - hook_name="on_Snapshot__99_other_wget_hook", + hook_name="on_Snapshot__06_wget.finite.bg", status=ArchiveResult.StatusChoices.SUCCEEDED, ) -def test_archivewebpage_lifecycle_hooks_project_one_plugin_output(): +def test_archiveresult_event_create_uses_one_result_lookup(): from abx_dl.events import ArchiveResultEvent - from archivebox.core.models import ArchiveResult + from django.db import connection + from django.test.utils import CaptureQueriesContext + from archivebox.services.archive_result_service import _save_archiveresult_event_to_db snapshot = _create_snapshot() - _save_archiveresult_event_to_db( - ArchiveResultEvent( - snapshot_id=str(snapshot.id), - plugin="archivewebpage", - hook_name="on_Snapshot__16_archivewebpage_start", - status="succeeded", - output_str="recording started", - output_files=[OutputFile(path="recording.json", extension="json", mimetype="application/json", size=175)], - ), - None, - ) - _save_archiveresult_event_to_db( - ArchiveResultEvent( - snapshot_id=str(snapshot.id), - plugin="archivewebpage", - hook_name="on_Snapshot__65_archivewebpage_stop", - status="succeeded", - output_str="archivewebpage.wacz", - output_files=[OutputFile(path="archivewebpage.wacz", extension="wacz", size=2048)], - ), - None, + event = ArchiveResultEvent( + snapshot_id=str(snapshot.id), + plugin="review-query-count", + hook_name="on_Snapshot__99_review.py", + status="failed", ) - result = ArchiveResult.objects.get(snapshot=snapshot, plugin="archivewebpage") - assert result.hook_name == "on_Snapshot__65_archivewebpage_stop" - assert result.output_str == "archivewebpage.wacz" - assert set(result.output_files) == {"recording.json", "archivewebpage.wacz"} - assert ArchiveResult.objects.filter(snapshot=snapshot, plugin="archivewebpage").count() == 1 + with CaptureQueriesContext(connection) as queries: + _save_archiveresult_event_to_db(event, None) + + result_lookups = [ + query["sql"] + for query in queries + if query["sql"].lstrip().upper().startswith("SELECT") and 'FROM "core_archiveresult"' in query["sql"] + ] + assert len(result_lookups) == 1 def test_process_completed_projects_failed_archiveresult_from_shipped_hook(tmp_path, hermetic_lib_dir): @@ -470,13 +460,12 @@ def test_retry_failed_archiveresults_requeues_snapshot_in_queued_state(): reset_count = snapshot.retry_failed_archiveresults() snapshot.refresh_from_db() - result = ArchiveResult.objects.get(snapshot=snapshot, plugin="chrome") + result = ArchiveResult.objects.get(snapshot=snapshot, plugin="chrome", hook_name="on_Snapshot__11_chrome_wait") assert reset_count == 1 assert snapshot.status == Snapshot.StatusChoices.QUEUED assert snapshot.retry_at is not None assert snapshot.current_step == 0 assert result.status == ArchiveResult.StatusChoices.QUEUED - assert result.hook_name == "" assert result.output_str == "" assert result.output_json is None assert result.output_files == {} diff --git a/archivebox/tests/test_cli_archiveresult.py b/archivebox/tests/test_cli_archiveresult.py index 13dcfa9f..e7efcf0e 100644 --- a/archivebox/tests/test_cli_archiveresult.py +++ b/archivebox/tests/test_cli_archiveresult.py @@ -61,11 +61,10 @@ class TestArchiveResultCreate: ar = next(r for r in records if r["type"] == "ArchiveResult") assert ar["plugin"] == "title" - # Queue projection is one row per plugin, while a plugin can contain - # several ordered hooks. The runner records the concrete hook only - # after execution; inventing one here would make the pending row claim - # work that has not run and reintroduce hook-level duplicate results. - assert ar["hook_name"] == "" + # Hook identity is known at scheduling time and is the durable retry + # key. Emitting it here prevents multiple hooks in one plugin from + # collapsing into the same ArchiveResult row. + assert ar["hook_name"] == "on_Snapshot__54_title" assert "id" not in ar def test_create_with_specific_plugin(self, initialized_archive): @@ -92,11 +91,9 @@ class TestArchiveResultCreate: assert code == 0 records = parse_jsonl_output(stdout2) ar_records = [r for r in records if r.get("type") == "ArchiveResult"] - assert len(ar_records) >= 1 + assert len(ar_records) == 1 assert all(record["plugin"] == "screenshot" for record in ar_records) - # A requested plugin is the schedulable unit; its concrete hook is an - # execution result, not input metadata on this pre-execution request. - assert all(record["hook_name"] == "" for record in ar_records) + assert [record["hook_name"] for record in ar_records] == ["on_Snapshot__51_screenshot"] def test_create_pass_through_crawl(self, initialized_archive): """Pass-through Crawl records unchanged.""" diff --git a/archivebox/tests/test_cli_extract_input.py b/archivebox/tests/test_cli_extract_input.py index 6d99fd88..3f2655b7 100644 --- a/archivebox/tests/test_cli_extract_input.py +++ b/archivebox/tests/test_cli_extract_input.py @@ -22,7 +22,7 @@ def create_extract_snapshot(initialized_archive, env, url="https://example.com") ) -def test_extract_archiveresult_record_queues_one_plugin_result(initialized_archive): +def test_extract_archiveresult_record_queues_only_exact_hook(initialized_archive): env = cli_env(PLUGINS="archivewebpage") create_extract_snapshot(initialized_archive, env) @@ -52,7 +52,7 @@ def test_extract_archiveresult_record_queues_one_plugin_result(initialized_archi rows = list( ArchiveResult.objects.filter(snapshot_id=snapshot_id, plugin="archivewebpage").values_list("hook_name", "status"), ) - assert rows == [("", ArchiveResult.StatusChoices.QUEUED)] + assert rows == [(hook_name, ArchiveResult.StatusChoices.QUEUED)] def test_extract_runs_on_snapshot_id(initialized_archive): diff --git a/archivebox/tests/test_cli_run.py b/archivebox/tests/test_cli_run.py index a3f30dd3..11db15be 100644 --- a/archivebox/tests/test_cli_run.py +++ b/archivebox/tests/test_cli_run.py @@ -240,7 +240,7 @@ class TestRunWithArchiveResult: """Tests for `archivebox run` with ArchiveResult input.""" @pytest.mark.django_db(transaction=True) - def test_run_treats_no_id_archiveresult_request_as_plugin_work(self, initialized_archive): + def test_run_creates_and_runs_exact_no_id_archiveresult_request(self, initialized_archive): import json from archivebox.core.models import ArchiveResult @@ -284,9 +284,13 @@ class TestRunWithArchiveResult: "output_str", ), ) - assert len(rows) == 1 - assert rows[0][0] != missing_hook - assert rows[0][1] in ArchiveResult.FINAL_STATES + assert rows == [ + ( + missing_hook, + ArchiveResult.StatusChoices.FAILED, + "Queued hook is no longer available in the installed plugin", + ), + ] def test_run_requeues_failed_archiveresult(self, initialized_archive): """Run re-queues a failed ArchiveResult.""" @@ -1072,7 +1076,7 @@ class TestRecoverOrchestratorState: assert snapshot.retry_at is None assert snapshot.downloaded_at is not None - def test_create_pending_archiveresults_creates_one_plugin_row_not_hook_rows(self): + def test_create_pending_archiveresults_uses_canonical_hook_names(self): from django.utils import timezone from archivebox.base_models.models import get_or_create_system_user_pk @@ -1092,17 +1096,11 @@ class TestRecoverOrchestratorState: retry_at=timezone.now(), ) - results = snapshot.create_pending_archiveresults( - hooks=[ - ("chrome", "on_Snapshot__00_chrome_launch.daemon.bg"), - ("chrome", "on_Snapshot__01_chrome_tab.daemon.bg"), - ("chrome", "on_Snapshot__30_chrome_navigate"), - ("title", "on_Snapshot__54_title"), - ], - ) + snapshot.create_pending_archiveresults() - assert [(result.plugin, result.hook_name) for result in results] == [("chrome", ""), ("title", "")] - assert ArchiveResult.objects.filter(snapshot=snapshot).count() == 2 + hook_names = list(ArchiveResult.objects.filter(snapshot=snapshot).values_list("hook_name", flat=True)) + assert hook_names + assert all(not hook_name.endswith((".py", ".js", ".sh")) for hook_name in hook_names) def test_snapshot_hooks_for_pending_archiveresults_respects_disabled_plugins_when_plugins_empty(self): from django.utils import timezone @@ -1718,7 +1716,7 @@ class TestRecoverOrchestratorState: assert result.end_ts is not None @pytest.mark.django_db(transaction=True) - def test_run_due_snapshot_runs_plugin_for_obsolete_queued_hook_name(self): + def test_run_due_snapshot_fails_obsolete_queued_hook_name(self): from django.utils import timezone from archivebox.base_models.models import get_or_create_system_user_pk @@ -1749,8 +1747,8 @@ class TestRecoverOrchestratorState: result.refresh_from_db() snapshot.refresh_from_db() - assert result.status in ArchiveResult.FINAL_STATES - assert result.hook_name != "on_Snapshot__50_singlefile.py" + assert result.status == ArchiveResult.StatusChoices.FAILED + assert result.output_str == "Queued hook is no longer available in the installed plugin" assert snapshot.retry_at is None @pytest.mark.django_db(transaction=True) @@ -1829,10 +1827,7 @@ class TestRecoverOrchestratorState: recursive_test_site, chrome_isolation, ): - from pathlib import Path - - from archivebox.core.models import ArchiveResult, Snapshot - from archivebox.machine.models import Process + from archivebox.core.models import ArchiveResult from archivebox.tests.test_orm_helpers import use_archivebox_db env = cli_env(disable_extractors=True) @@ -1870,19 +1865,13 @@ class TestRecoverOrchestratorState: snapshot_id = navigate_record["snapshot_id"] with use_archivebox_db(initialized_archive): - snapshot = Snapshot.objects.get(id=snapshot_id) - tab_processes = [ - process - for process in Process.objects.filter(pwd=str(Path(snapshot.output_dir) / "chrome")).order_by("-started_at") - if Path(process.hook_script_name or "").stem == "on_Snapshot__01_chrome_tab.daemon.bg" - ] - assert tab_processes - first_tab_process_id = tab_processes[0].id - plugin_result = ArchiveResult.objects.get( + tab_result = ArchiveResult.objects.get( snapshot_id=snapshot_id, plugin="chrome", + hook_name="on_Snapshot__01_chrome_tab.daemon.bg", ) - assert plugin_result.hook_name == "on_Snapshot__30_chrome_navigate" + first_tab_process_id = tab_result.process_id + assert first_tab_process_id is not None update_process = run_archivebox_cmd( ["archiveresult", "update", "--status=queued"], @@ -1916,18 +1905,18 @@ class TestRecoverOrchestratorState: navigate_result = ArchiveResult.objects.get( snapshot_id=snapshot_id, plugin="chrome", + hook_name="on_Snapshot__30_chrome_navigate", + ) + tab_result = ArchiveResult.objects.get( + snapshot_id=snapshot_id, + plugin="chrome", + hook_name="on_Snapshot__01_chrome_tab.daemon.bg", ) - snapshot = Snapshot.objects.get(id=snapshot_id) - tab_processes = [ - process - for process in Process.objects.filter(pwd=str(Path(snapshot.output_dir) / "chrome")).order_by("-started_at") - if Path(process.hook_script_name or "").stem == "on_Snapshot__01_chrome_tab.daemon.bg" - ] assert run_process.returncode == 0 assert navigate_result.status == ArchiveResult.StatusChoices.SUCCEEDED - assert navigate_result.hook_name == "on_Snapshot__30_chrome_navigate" - assert tab_processes[0].id != first_tab_process_id + assert tab_result.process_id is not None + assert tab_result.process_id != first_tab_process_id def test_recover_orchestrator_state_ignores_sealed_downloaded_snapshot_without_results(self): from django.utils import timezone diff --git a/archivebox/tests/test_migration_archiveresult_hook_identity.py b/archivebox/tests/test_migration_archiveresult_hook_identity.py new file mode 100644 index 00000000..fa290489 --- /dev/null +++ b/archivebox/tests/test_migration_archiveresult_hook_identity.py @@ -0,0 +1,95 @@ +import pytest +from django.db import IntegrityError, connection, transaction +from django.db.migrations.executor import MigrationExecutor + + +pytestmark = pytest.mark.django_db(transaction=True) + + +def test_published_plugin_constraint_preserves_existing_hook_rows(): + try: + executor = MigrationExecutor(connection) + executor.migrate([("core", "0051_postgres_url_pattern_ops_index")]) + old_apps = executor.loader.project_state([("core", "0051_postgres_url_pattern_ops_index")]).apps + Crawl = old_apps.get_model("crawls", "Crawl") + Snapshot = old_apps.get_model("core", "Snapshot") + ArchiveResult = old_apps.get_model("core", "ArchiveResult") + crawl = Crawl.objects.create(urls="https://example.com") + snapshot = Snapshot.objects.create(url="https://example.com/history", crawl=crawl) + first = ArchiveResult.objects.create( + snapshot=snapshot, + plugin="responses", + hook_name="browser-upload", + output_json={"source": "browser"}, + ) + second = ArchiveResult.objects.create( + snapshot=snapshot, + plugin="responses", + hook_name="server-capture", + output_json={"source": "server"}, + ) + + executor = MigrationExecutor(connection) + executor.migrate([("core", "0054_restore_archiveresult_hook_identity")]) + new_apps = executor.loader.project_state([("core", "0054_restore_archiveresult_hook_identity")]).apps + ArchiveResult = new_apps.get_model("core", "ArchiveResult") + rows = list( + ArchiveResult.objects.filter(snapshot_id=snapshot.id, plugin="responses") + .order_by("hook_name") + .values_list("id", "hook_name", "output_json"), + ) + + assert rows == [ + (first.id, "browser-upload", {"source": "browser"}), + (second.id, "server-capture", {"source": "server"}), + ] + + executor = MigrationExecutor(connection) + executor.migrate([("core", "0053_alter_archiveresult_options")]) + temporary_apps = executor.loader.project_state([("core", "0053_alter_archiveresult_options")]).apps + ArchiveResult = temporary_apps.get_model("core", "ArchiveResult") + temporary_rows = list( + ArchiveResult.objects.filter(snapshot_id=snapshot.id).values_list("plugin", "output_json"), + ) + assert len(temporary_rows) == 2 + assert len({plugin for plugin, _output_json in temporary_rows}) == 2 + + executor = MigrationExecutor(connection) + executor.migrate([("core", "0054_restore_archiveresult_hook_identity")]) + restored_apps = executor.loader.project_state([("core", "0054_restore_archiveresult_hook_identity")]).apps + ArchiveResult = restored_apps.get_model("core", "ArchiveResult") + assert ( + list( + ArchiveResult.objects.filter(snapshot_id=snapshot.id, plugin="responses") + .order_by("hook_name") + .values_list("id", "hook_name", "output_json"), + ) + == rows + ) + finally: + MigrationExecutor(connection).migrate([("core", "0054_restore_archiveresult_hook_identity")]) + + +def test_migration_restores_distinct_hook_rows_after_published_plugin_constraint(): + try: + executor = MigrationExecutor(connection) + executor.migrate([("core", "0052_unique_archiveresult_per_snapshot_plugin")]) + old_apps = executor.loader.project_state([("core", "0052_unique_archiveresult_per_snapshot_plugin")]).apps + Crawl = old_apps.get_model("crawls", "Crawl") + Snapshot = old_apps.get_model("core", "Snapshot") + ArchiveResult = old_apps.get_model("core", "ArchiveResult") + crawl = Crawl.objects.create(urls="https://example.com") + snapshot = Snapshot.objects.create(url="https://example.com/migration", crawl=crawl) + ArchiveResult.objects.create(snapshot=snapshot, plugin="responses", hook_name="browser-upload") + + executor = MigrationExecutor(connection) + executor.migrate([("core", "0054_restore_archiveresult_hook_identity")]) + new_apps = executor.loader.project_state([("core", "0054_restore_archiveresult_hook_identity")]).apps + ArchiveResult = new_apps.get_model("core", "ArchiveResult") + + ArchiveResult.objects.create(snapshot_id=snapshot.id, plugin="responses", hook_name="server-capture") + assert ArchiveResult.objects.filter(snapshot_id=snapshot.id, plugin="responses").count() == 2 + with pytest.raises(IntegrityError), transaction.atomic(): + ArchiveResult.objects.create(snapshot_id=snapshot.id, plugin="responses", hook_name="browser-upload") + finally: + MigrationExecutor(connection).migrate([("core", "0054_restore_archiveresult_hook_identity")]) diff --git a/archivebox/tests/test_migration_archiveresult_unique_plugin.py b/archivebox/tests/test_migration_archiveresult_unique_plugin.py deleted file mode 100644 index c3297056..00000000 --- a/archivebox/tests/test_migration_archiveresult_unique_plugin.py +++ /dev/null @@ -1,53 +0,0 @@ -import pytest -from django.db import connection -from django.db.migrations.executor import MigrationExecutor - - -pytestmark = pytest.mark.django_db(transaction=True) - - -def test_migration_consolidates_plugin_rows_and_enforces_uniqueness(): - try: - executor = MigrationExecutor(connection) - executor.migrate([("core", "0051_postgres_url_pattern_ops_index")]) - old_apps = executor.loader.project_state([("core", "0051_postgres_url_pattern_ops_index")]).apps - Crawl = old_apps.get_model("crawls", "Crawl") - Snapshot = old_apps.get_model("core", "Snapshot") - ArchiveResult = old_apps.get_model("core", "ArchiveResult") - crawl = Crawl.objects.create(urls="https://example.com") - snapshot = Snapshot.objects.create(url="https://example.com/migration", crawl=crawl, output_size=10) - ArchiveResult.objects.create( - snapshot=snapshot, - plugin="responses", - hook_name="browser-upload", - status="succeeded", - output_files={"browser.png": {"size": 7}}, - output_size=7, - output_mimetypes="image/png", - ) - ArchiveResult.objects.create( - snapshot=snapshot, - plugin="responses", - hook_name="server-capture", - status="noresults", - output_files={"server.json": {"size": 3}}, - output_size=3, - output_mimetypes="application/json", - ) - - executor = MigrationExecutor(connection) - executor.migrate([("core", "0052_unique_archiveresult_per_snapshot_plugin")]) - new_apps = executor.loader.project_state([("core", "0052_unique_archiveresult_per_snapshot_plugin")]).apps - Snapshot = new_apps.get_model("core", "Snapshot") - ArchiveResult = new_apps.get_model("core", "ArchiveResult") - snapshot = Snapshot.objects.get(id=snapshot.id) - result = ArchiveResult.objects.get(snapshot=snapshot, plugin="responses") - - assert ArchiveResult.objects.filter(snapshot=snapshot, plugin="responses").count() == 1 - assert result.status == "succeeded" - assert set(result.output_files) == {"browser.png", "server.json"} - assert result.output_size == 10 - assert snapshot.output_size == 10 - assert set(result.output_mimetypes.split(",")) == {"application/json", "image/png"} - finally: - MigrationExecutor(connection).migrate([("core", "0052_unique_archiveresult_per_snapshot_plugin")]) diff --git a/archivebox/tests/test_migrations_08_to_09.py b/archivebox/tests/test_migrations_08_to_09.py index 744db687..f9265962 100644 --- a/archivebox/tests/test_migrations_08_to_09.py +++ b/archivebox/tests/test_migrations_08_to_09.py @@ -898,7 +898,10 @@ def test_update_preserves_legacy_plugin_directory_without_output_files(migration (snapshot["id"],), ).fetchall() conn.close() - assert rows == [("", "{}", "")] + assert len(rows) == 2 + assert all(output_str == "" for output_str, _output_files, _hook_name in rows) + assert len({hook_name for _output_str, _output_files, hook_name in rows}) == 2 + assert sum(not hook_name for _output_str, _output_files, hook_name in rows) == 1 def test_07_filesystem_hop_preserves_complete_output_tree(tmp_path): diff --git a/archivebox/tests/test_snapshot_service.py b/archivebox/tests/test_snapshot_service.py index cdc1036a..2f8f139d 100644 --- a/archivebox/tests/test_snapshot_service.py +++ b/archivebox/tests/test_snapshot_service.py @@ -40,7 +40,7 @@ def _snapshot_state(cwd: Path, url: str) -> dict[str, object]: } -def test_snapshot_merge_consolidates_plugin_output_identity(admin_user): +def test_snapshot_merge_consolidates_only_exact_hook_identity(admin_user): from archivebox.crawls.models import Crawl keeper_crawl = Crawl.objects.create(urls="https://example.com", created_by=admin_user) @@ -64,15 +64,24 @@ def test_snapshot_merge_consolidates_plugin_output_identity(admin_user): output_files={"archivewebpage.wacz": {"size": 7}}, output_size=7, ) + ArchiveResult.objects.create( + snapshot=duplicate, + plugin="archivewebpage", + hook_name="on_Snapshot__16_archivewebpage_start", + status=ArchiveResult.StatusChoices.SUCCEEDED, + output_str="recording started", + ) + Snapshot._merge_snapshots([keeper, duplicate]) assert not Snapshot.objects.filter(pk=duplicate.pk).exists() results = ArchiveResult.objects.filter(snapshot=keeper, plugin="archivewebpage") - assert results.count() == 1 - stop_result = results.get() + assert results.count() == 2 + stop_result = results.get(hook_name="on_Snapshot__65_archivewebpage_stop") assert stop_result.status == ArchiveResult.StatusChoices.SUCCEEDED assert stop_result.output_str == "archivewebpage.wacz" assert set(stop_result.output_files) == {"older.wacz", "archivewebpage.wacz"} + assert results.filter(hook_name="on_Snapshot__16_archivewebpage_start").exists() @pytest.mark.timeout(180) diff --git a/archivebox/tests/test_ui_admin_archiveresult.py b/archivebox/tests/test_ui_admin_archiveresult.py index f2f6f44f..0d85e4f5 100644 --- a/archivebox/tests/test_ui_admin_archiveresult.py +++ b/archivebox/tests/test_ui_admin_archiveresult.py @@ -142,8 +142,7 @@ class TestArchiveResultAdminListView: snapshot.refresh_from_db() assert snapshot.output_size == 0 - def test_duplicate_plugin_result_is_rejected_without_touching_output(self, snapshot): - from django.db import IntegrityError, transaction + def test_deleting_sibling_hook_preserves_shared_plugin_output(self, snapshot): from archivebox.core.models import ArchiveResult output_dir = Path(snapshot.output_dir) / "responses" @@ -157,13 +156,14 @@ class TestArchiveResultAdminListView: status=ArchiveResult.StatusChoices.SUCCEEDED, output_size=18, ) - with pytest.raises(IntegrityError), transaction.atomic(): - ArchiveResult.objects.create( - snapshot=snapshot, - plugin="responses", - hook_name="on_Snapshot__24_responses.daemon.bg.replayed", - status=ArchiveResult.StatusChoices.NORESULTS, - ) + duplicate = ArchiveResult.objects.create( + snapshot=snapshot, + plugin="responses", + hook_name="on_Snapshot__24_responses.daemon.bg.replayed", + status=ArchiveResult.StatusChoices.NORESULTS, + ) + + duplicate.delete() assert ArchiveResult.objects.filter(pk=primary.pk).exists() assert output_file.read_text() == "captured response\n" diff --git a/archivebox/tests/test_ui_admin_snapshot.py b/archivebox/tests/test_ui_admin_snapshot.py index 79df449f..fefa7bab 100644 --- a/archivebox/tests/test_ui_admin_snapshot.py +++ b/archivebox/tests/test_ui_admin_snapshot.py @@ -324,6 +324,35 @@ def test_snapshot_admin_preview_uses_extension_screenshot_when_standard_screensh assert "chrome_extension_screenshot/screenshot-2.png" not in preview["fallback_list"] +def test_snapshot_admin_attributes_new_tags_to_authenticated_user(client, snapshot, admin_user): + from archivebox.core.models import Tag + + client.force_login(admin_user) + response = client.post( + reverse("admin:core_snapshot_change", args=[snapshot.pk]), + { + "url": snapshot.url, + "title": snapshot.title or "", + "tags_editor": "admin-created-tag", + "permissions_config": "private", + "status": snapshot.status, + "retry_at": "", + "bookmarked_at_0": snapshot.bookmarked_at.date().isoformat(), + "bookmarked_at_1": snapshot.bookmarked_at.time().isoformat(), + "crawl": str(snapshot.crawl_id), + "config": '{"SAVE_ARCHIVE_DOT_ORG": "false"}', + "notes": "", + "_save": "Save", + }, + HTTP_HOST=ADMIN_TEST_HOST, + ) + + assert response.status_code == 302, response.context and response.context["adminform"].form.errors + tag = Tag.objects.get(name="admin-created-tag") + assert tag.created_by == admin_user + assert snapshot.tags.filter(pk=tag.pk).exists() + + class TestSnapshotProgressStats: """Tests for Snapshot.get_progress_stats() method.""" @@ -928,10 +957,9 @@ class TestSnapshotOutputDeletion: from archivebox.core.models import ArchiveResult first = self._create_output(snapshot, size=11) - second = self._create_output(snapshot, plugin="singlefile", hook_name="on_Snapshot__50_singlefile.py", size=13) + second = self._create_output(snapshot, hook_name="on_Snapshot__51_screenshot_retry.py", size=13) kept = self._create_output(snapshot, plugin="pdf", hook_name="on_Snapshot__60_pdf.py", size=7) deleted_dir = Path(first.output_dir) - second_deleted_dir = Path(second.output_dir) kept_dir = Path(kept.output_dir) hashes_dir = Path(snapshot.output_dir) / "hashes" hashes_dir.mkdir(parents=True, exist_ok=True) @@ -964,7 +992,6 @@ class TestSnapshotOutputDeletion: assert not ArchiveResult.objects.filter(pk__in=[first.pk, second.pk]).exists() assert ArchiveResult.objects.filter(pk=kept.pk).exists() assert not deleted_dir.exists() - assert not second_deleted_dir.exists() assert kept_dir.exists() snapshot.refresh_from_db() assert snapshot.output_size == 7 diff --git a/archivebox/tests/test_urls.py b/archivebox/tests/test_urls.py index b18528ce..e5a51afb 100644 --- a/archivebox/tests/test_urls.py +++ b/archivebox/tests/test_urls.py @@ -15,22 +15,26 @@ REPO_ROOT = Path(__file__).resolve().parents[2] @pytest.mark.django_db -def test_archiveresult_relpath_uses_plugin_result_that_owns_output(admin_user): +def test_archiveresult_relpath_uses_sibling_hook_that_owns_output(admin_user): from archivebox.core.models import ArchiveResult, Snapshot from archivebox.core.views import _resolve_archiveresult_relpath from archivebox.crawls.models import Crawl crawl = Crawl.objects.create(urls="https://example.com", created_by=admin_user) snapshot = Snapshot.objects.create(url="https://example.com", crawl=crawl) + ArchiveResult.objects.create( + snapshot=snapshot, + plugin="screenshot", + hook_name="on_Snapshot__archivebox_browser_extension_upload", + status=ArchiveResult.StatusChoices.SUCCEEDED, + output_files={"browser.png": {"size": 7}}, + ) server_result = ArchiveResult.objects.create( snapshot=snapshot, plugin="screenshot", hook_name="on_Snapshot__50_screenshot", status=ArchiveResult.StatusChoices.SUCCEEDED, - output_files={ - "browser.png": {"size": 7}, - "screenshot.png": {"size": 6, "root_relative": True}, - }, + output_files={"screenshot.png": {"size": 6, "root_relative": True}}, ) resolved_path, result = _resolve_archiveresult_relpath(snapshot, "screenshot/screenshot.png")