diff --git a/archivebox/api/v1_core.py b/archivebox/api/v1_core.py index 3e823d82..b41b6f71 100644 --- a/archivebox/api/v1_core.py +++ b/archivebox/api/v1_core.py @@ -389,13 +389,6 @@ def _queue_archiveresult_snapshot_maintenance(snapshot: Snapshot) -> None: snapshot.safe_update(updates, refresh=False) -def _merge_archiveresult_output_file_maps(results: list[ArchiveResult]) -> dict[str, dict[str, Any]]: - output_files: dict[str, dict[str, Any]] = {} - for result in results: - output_files.update(result.output_file_map()) - return output_files - - def _write_archiveresult_files( request: HttpRequest, snapshot: Snapshot, @@ -535,15 +528,8 @@ 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 - matching_results = list( - ArchiveResult.objects.filter( - snapshot=snapshot, - plugin=plugin_name, - hook_name=hook, - ).order_by("created_at", "id"), - ) - existing_result = matching_results[0] if matching_results else None - existing_output_files = _merge_archiveresult_output_file_maps(matching_results) + 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( request, snapshot, @@ -554,22 +540,13 @@ def create_archiveresult( now = timezone.now() with transaction.atomic(): - matching_results = list( - ArchiveResult.objects.filter( - snapshot=snapshot, - plugin=plugin_name, - hook_name=hook, - ).order_by("created_at", "id"), - ) - if matching_results: - existing_result = matching_results[0] + 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 = { - **_merge_archiveresult_output_file_maps(matching_results), + **dict(existing_result.output_files or {}), **output_files, } - duplicate_ids = [result.id for result in matching_results[1:]] - if duplicate_ids: - ArchiveResult.objects.filter(id__in=duplicate_ids).delete() result = existing_result else: existing_result = None @@ -588,6 +565,7 @@ def create_archiveresult( 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 diff --git a/archivebox/cli/archivebox_archiveresult.py b/archivebox/cli/archivebox_archiveresult.py index 1b79f322..d83968cb 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 requests for all enabled snapshot hooks. + Otherwise, emits one request for each enabled Snapshot plugin. Exit codes: 0: Success @@ -151,11 +151,9 @@ def create_archiveresults( else: config = get_config(crawl=snapshot.crawl, snapshot=snapshot) hooks = discover_hooks("Snapshot", config=config) - for hook_path in hooks: - hook_name = hook_path.stem - plugin_name = hook_path.parent.name + 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, hook_name=hook_name, status=status)) + 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) @@ -256,9 +254,11 @@ def update_archiveresults( # Apply updates from CLI flags if status: - result.status = status - - result.save() + if status == ArchiveResult.StatusChoices.QUEUED: + result.reset_for_retry() + else: + result.status = status + result.save(update_fields=["status", "modified_at"]) updated_count += 1 if not is_tty: diff --git a/archivebox/cli/archivebox_extract.py b/archivebox/cli/archivebox_extract.py index f1dd0bb3..7e9b4675 100644 --- a/archivebox/cli/archivebox_extract.py +++ b/archivebox/cli/archivebox_extract.py @@ -122,7 +122,6 @@ 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() @@ -193,30 +192,21 @@ def run_plugins( if snapshot_id in existing_snapshot_ids for plugin_name in plugin_names ) - plugins_by_name = discover_plugins(runtime="archivebox") - requested_rows: set[tuple[str, str, str]] = set() - for snapshot_id, plugin_name in requested_pairs: - 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: set[tuple[str, str]] = set() + if preserve_queued and requested_pairs: queued_rows = { - (str(snapshot_id), plugin_name, hook_name) - for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( + (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, _hook_name in requested_rows}, + plugin__in={plugin_name for _snapshot_id, plugin_name in requested_pairs}, status=ArchiveResult.StatusChoices.QUEUED, - ).values_list("snapshot_id", "plugin", "hook_name") + ).values_list("snapshot_id", "plugin") } - rows_to_queue = requested_rows - queued_rows + rows_to_queue = requested_pairs - queued_rows reset_fields = { "status": ArchiveResult.StatusChoices.QUEUED, + "hook_name": "", "output_str": "", "output_json": None, "output_files": {}, @@ -227,28 +217,28 @@ def run_plugins( "modified_at": timezone.now(), } if rows_to_queue and plugins_list: - 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( + 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( **reset_fields, ) elif rows_to_queue and requested_plugins_by_snapshot: - 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( + 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( **reset_fields, ) existing_rows = ( { - (str(snapshot_id), plugin_name, hook_name) - for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( + (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, _hook_name in rows_to_queue}, - ).values_list("snapshot_id", "plugin", "hook_name") + plugin__in={plugin_name for _snapshot_id, plugin_name in rows_to_queue}, + ).values_list("snapshot_id", "plugin") } if rows_to_queue else set() @@ -260,10 +250,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, hook_name in sorted(missing_rows) + for snapshot_id, plugin_name in sorted(missing_rows) ], batch_size=500, ) @@ -271,13 +261,13 @@ def run_plugins( processed_count = len(existing_snapshot_ids) queue_at = timezone.now() if existing_snapshot_ids: - if requested_rows: - # Search indexing on a sealed Snapshot is the only targeted hook + if requested_pairs: + # Search indexing on a sealed Snapshot is the only targeted plugin # allowed to bypass the normal lifecycle. Every other requested # plugin requeues its Snapshot through the unified state machine. - affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in rows_to_queue} + affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name in rows_to_queue} if preserve_queued and queued_rows: - queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in queued_rows} + queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name in queued_rows} affected_snapshot_ids.update( str(snapshot_id) for snapshot_id in Snapshot.objects.filter(id__in=queued_snapshot_ids) @@ -292,7 +282,7 @@ def run_plugins( ) ) requested_plugins_by_id: dict[str, set[str]] = defaultdict(set) - for snapshot_id, plugin_name, _hook_name in requested_rows: + for snapshot_id, plugin_name in requested_pairs: 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 @@ -328,7 +318,7 @@ def run_plugins( refresh=False, extra_filter={"status": snapshot.status}, ) - if existing_crawl_ids and not requested_rows: + if existing_crawl_ids and not requested_pairs: from archivebox.crawls.models import Crawl for crawl in Crawl.objects.filter(id__in=existing_crawl_ids).only("id", "status", "retry_at", "modified_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 new file mode 100644 index 00000000..3a1989e5 --- /dev/null +++ b/archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py @@ -0,0 +1,88 @@ +from django.db import migrations, models +from django.db.models import Sum + + +def consolidate_archiveresults_per_plugin(apps, schema_editor): + 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( + snapshot_id=group["snapshot_id"], + plugin=group["plugin"], + ).order_by("created_at", "id"), + ) + canonical = rows[0] + winner = max( + rows, + key=lambda row: ( + bool(row.output_files), + int(row.output_size or 0), + row.modified_at, + 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()) + + canonical.hook_name = winner.hook_name + canonical.status = winner.status + canonical.output_str = winner.output_str + canonical.output_json = winner.output_json + canonical.output_files = output_files + canonical.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), + ) + canonical.output_mimetypes = ",".join(sorted(mimetypes)) + canonical.start_ts = min((row.start_ts for row in rows if row.start_ts), default=None) + canonical.end_ts = max((row.end_ts for row in rows if row.end_ts), default=None) + canonical.save( + update_fields=[ + "hook_name", + "status", + "output_str", + "output_json", + "output_files", + "output_size", + "output_mimetypes", + "start_ts", + "end_ts", + ], + ) + ArchiveResult.objects.filter(id__in=[row.id for row in rows[1:]]).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) + + +class Migration(migrations.Migration): + dependencies = [ + ("core", "0051_postgres_url_pattern_ops_index"), + ] + + operations = [ + migrations.RemoveConstraint( + model_name="archiveresult", + name="unique_archiveresult_per_snapshot_hook", + ), + migrations.RunPython( + consolidate_archiveresults_per_plugin, + reverse_code=migrations.RunPython.noop, + ), + migrations.AddConstraint( + model_name="archiveresult", + constraint=models.UniqueConstraint( + fields=("snapshot", "plugin"), + name="unique_archiveresult_per_snapshot_plugin", + ), + ), + ] diff --git a/archivebox/core/models.py b/archivebox/core/models.py index d6fcf2c3..c0fd9042 100644 --- a/archivebox/core/models.py +++ b/archivebox/core/models.py @@ -1740,8 +1740,8 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW 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 hook; retries update the existing row.""" - existing = {(ar.plugin, ar.hook_name): ar for ar in ArchiveResult.objects.filter(snapshot=self)} + """Merge ArchiveResults one row per plugin; retries update the existing row.""" + existing = {ar.plugin: ar for ar in ArchiveResult.objects.filter(snapshot=self)} if update_existing: for archiveresult in existing.values(): normalized_status = ArchiveResult.normalize_status(archiveresult.status) @@ -1804,7 +1804,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, hook_name)) + existing_result = existing.get(plugin) if existing_result: if not update_existing: return @@ -1813,6 +1813,9 @@ 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") @@ -1869,7 +1872,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW end_ts=end_ts, process=process, ) - existing[(plugin, hook_name)] = archiveresult + existing[plugin] = archiveresult def write_index_json(self): """Write index.json in 0.9.x format (deprecated, use write_index_jsonl).""" @@ -2199,11 +2202,42 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW for tag in dup.tags.all(): keeper.add_tag_ids([tag.pk]) - # Move ArchiveResults - ArchiveResult.objects.filter(snapshot=dup).update( - snapshot=keeper, - modified_at=timezone.now(), - ) + # Move ArchiveResults while preserving the one-row-per-plugin invariant. + for result in ArchiveResult.objects.filter(snapshot=dup): + existing = ArchiveResult.objects.filter(snapshot=keeper, plugin=result.plugin).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( + sum( + ArchiveResult._coerce_output_file_size(metadata.get("size")) + for metadata in output_files.values() + if isinstance(metadata, dict) + ), + *prior_output_sizes, + ) + if prefer_result: + existing.status = result.status + existing.hook_name = result.hook_name + existing.output_str = result.output_str + existing.output_json = result.output_json + existing.output_mimetypes = ",".join( + sorted( + { + mimetype.strip() + for value in (existing.output_mimetypes, result.output_mimetypes) + for mimetype in value.split(",") + if mimetype.strip() + }, + ), + ) + existing.save() + result.delete() # Delete dup.delete() @@ -2942,8 +2976,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 hook (not per plugin), with hook_name set. - This enables step-based execution where all hooks in a step can run in parallel. + 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. """ try: self.validate_url_for_archiving() @@ -2961,17 +2995,21 @@ 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, hook_name in hooks: - # ArchiveResult output is one filesystem directory per plugin hook, so - # retries must update this row in place instead of creating siblings. + 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, - hook_name=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) @@ -3060,10 +3098,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.exclude(hook_name="").update( + count = retryable_results.update( status=ArchiveResult.StatusChoices.QUEUED, + hook_name="", output_str="", output_json=None, output_files={}, @@ -3074,11 +3112,11 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW modified_at=now, ) - if count + legacy_result_count > 0: + if count > 0: self.refresh_from_db(fields=["modified_at", "retry_at", "status"]) self.queue_for_extraction(when=now) - return count + legacy_result_count + return count # ========================================================================= # URL Helper Properties (migrated from Link schema) @@ -3928,7 +3966,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", "hook_name"], name="unique_archiveresult_per_snapshot_hook"), + models.UniqueConstraint(fields=["snapshot", "plugin"], name="unique_archiveresult_per_snapshot_plugin"), ] def __str__(self): @@ -4039,16 +4077,16 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): except ArchiveResult.DoesNotExist: pass - # Get or create by snapshot_id + plugin + hook_name. The filesystem has a - # single output dir for each hook, so retries update that same DB row. + # 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. try: snapshot = Snapshot.objects.get(id=snapshot_id) result, _ = ArchiveResult.objects.get_or_create( snapshot=snapshot, plugin=plugin, - hook_name=record.get("hook_name", ""), defaults={ + "hook_name": record.get("hook_name", ""), "status": record.get("status", "queued"), "output_str": record.get("output_str", ""), }, @@ -4115,13 +4153,10 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): 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: - 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).delete_output_paths(paths) type(self).refresh_snapshot_output_sizes({snapshot_id}) snapshot = Snapshot.objects.filter(pk=snapshot_id).first() if snapshot: @@ -4155,6 +4190,7 @@ 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 @@ -4167,6 +4203,7 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): self.save( update_fields=[ "status", + "hook_name", "retry_at", "output_str", "output_json", diff --git a/archivebox/services/archive_result_service.py b/archivebox/services/archive_result_service.py index 0ed9ff31..eab86c9b 100644 --- a/archivebox/services/archive_result_service.py +++ b/archivebox/services/archive_result_service.py @@ -4,6 +4,7 @@ import asyncio import inspect import json import os +import re import signal import sys import time @@ -202,6 +203,11 @@ 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: @@ -319,6 +325,7 @@ 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, @@ -337,7 +344,6 @@ def _save_archiveresult_event_to_db( result = ArchiveResult.objects.filter( snapshot=snapshot, plugin=event.plugin, - hook_name=event.hook_name, ).first() if result is None: try: @@ -345,7 +351,6 @@ def _save_archiveresult_event_to_db( result = ArchiveResult.objects.create( snapshot=snapshot, plugin=event.plugin, - hook_name=event.hook_name, **defaults, ) except IntegrityError: @@ -353,9 +358,30 @@ def _save_archiveresult_event_to_db( result = ArchiveResult.objects.get( snapshot=snapshot, plugin=event.plugin, - hook_name=event.hook_name, ) + if result.output_files: + merged_output_files = {**result.output_files, **defaults["output_files"]} + defaults["output_files"] = merged_output_files + defaults["output_size"], defaults["output_mimetypes"] = _summarize_output_files(merged_output_files) + 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 = [] for field, value in defaults.items(): @@ -396,7 +422,7 @@ def _save_archiveresult_event_to_db( def mark_archiveresult_started(event: ProcessStartedEvent, *, snapshot_id: str, process_id: str) -> None: - """Advance an existing queued hook row after its OS process is persisted.""" + """Advance an existing queued plugin row after its OS process is persisted.""" from archivebox.core.models import ArchiveResult started_at = parse_event_datetime(event.start_ts) @@ -405,9 +431,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, @@ -422,7 +448,7 @@ class ArchiveResultService(BaseService): def __init__(self, bus): self._completed_process_event_ids: set[str] = set() - self._save_locks: dict[tuple[str, str, str], asyncio.Lock] = {} + self._save_locks: dict[tuple[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) @@ -437,7 +463,7 @@ class ArchiveResultService(BaseService): where=lambda candidate: self.bus.event_is_child_of(event, candidate), ) - key = (str(event.snapshot_id), event.plugin, event.hook_name) + key = (str(event.snapshot_id), event.plugin) 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 5035ef48..6ad9d29f 100644 --- a/archivebox/services/runner.py +++ b/archivebox/services/runner.py @@ -1038,7 +1038,7 @@ class CrawlRunner: ) return [plugin for plugin in queued_plugins if plugin in expanded_selected_plugins] - selected_hooks_by_plugin = None + queued_plugins = None if snapshot["status"] == "started": _reset_count, running_count = await sync_to_async(snapshot["_snapshot"].reset_abandoned_results, thread_sensitive=True)() if running_count: @@ -1056,8 +1056,8 @@ class CrawlRunner: ) return if not self.selected_plugins_from_args: - queued_plugins, selected_hooks_by_plugin = await sync_to_async( - queued_plugins_and_hooks_for_snapshot, + queued_plugins = await sync_to_async( + queued_plugins_for_snapshot, thread_sensitive=True, )(snapshot["id"]) if queued_plugins: @@ -1070,13 +1070,10 @@ 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, selected_hooks_by_plugin = await sync_to_async( - queued_plugins_and_hooks_for_snapshot, + queued_plugins = await sync_to_async( + queued_plugins_for_snapshot, thread_sensitive=True, )(snapshot["id"]) if queued_plugins: @@ -1089,9 +1086,6 @@ 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", @@ -1106,10 +1100,10 @@ class CrawlRunner: if snapshot_selected_plugins else self.plugins ) - if selected_hooks_by_plugin is not None: - await sync_to_async(fail_unavailable_queued_hooks, thread_sensitive=True)( + if queued_plugins is not None: + await sync_to_async(fail_unavailable_queued_plugins, thread_sensitive=True)( snapshot["id"], - selected_hooks_by_plugin, + queued_plugins, plugins, ) remaining_queued_plugins = await sync_to_async( @@ -1129,7 +1123,6 @@ 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"], @@ -1153,7 +1146,7 @@ class CrawlRunner: snapshot_cleanup_enabled=True, snapshot_cleanup_phase_timeout=snapshot_phase_timeout, abort_requested=self.crawl_is_cancelled, - selected_hooks_by_plugin=selected_hooks_by_plugin, + selected_hooks_by_plugin=None, ) try: snapshot_event = SnapshotEvent( @@ -1375,37 +1368,19 @@ def run_binary(binary_id: str) -> None: asyncio.run(_run_binary(binary_id)) -def queued_plugins_and_hooks_for_snapshot(snapshot_id: str) -> tuple[list[str] | None, dict[str, set[str] | None] | None]: +def queued_plugins_for_snapshot(snapshot_id: str) -> list[str] | None: from archivebox.core.models import ArchiveResult - queued_results = list( + queued_plugins = list( ArchiveResult.objects.filter( snapshot_id=snapshot_id, status=ArchiveResult.StatusChoices.QUEUED, ) .exclude(plugin="") - .only("id", "plugin", "hook_name"), + .order_by("plugin") + .values_list("plugin", flat=True), ) - - 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 + return queued_plugins or None def config_overrides_for_queued_plugins(selected_plugins: list[str], **overrides: Any) -> dict[str, Any]: @@ -1422,40 +1397,27 @@ def config_overrides_for_queued_plugins(selected_plugins: list[str], **overrides return config_overrides -def fail_unavailable_queued_hooks( +def fail_unavailable_queued_plugins( snapshot_id: str, - selected_hooks_by_plugin: dict[str, set[str] | None], + queued_plugins: list[str], 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() - 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", - ) + 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", + ) def skip_disabled_queued_plugins(snapshot_id: str, plugin_names: list[str]) -> None: @@ -1480,35 +1442,6 @@ 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 @@ -1772,10 +1705,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_archiveresults = snapshot.archiveresult_set.exclude( + has_server_workset = snapshot.archiveresult_set.exclude( hook_name=Snapshot.BROWSER_EXTENSION_UPLOAD_HOOK_NAME, ).exists() - if has_server_archiveresults and snapshot.is_finished_processing(): + if has_server_workset 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: @@ -1784,10 +1717,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 - # 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: + 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. snapshot.create_pending_archiveresults(hooks=snapshot_hooks_for_pending_archiveresults(snapshot)) snapshot.sm.tick() snapshot.refresh_from_db() @@ -1804,11 +1737,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, 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( + queued_plugins = queued_plugins_for_snapshot(str(snapshot.id)) + if queued_plugins: + fail_unavailable_queued_plugins( str(snapshot.id), - selected_hooks_by_plugin, + queued_plugins, _discover_archivebox_plugins(), ) if not queued_plugins_for_snapshot(str(snapshot.id)): @@ -1832,10 +1765,10 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru ) snapshot.refresh_from_db() if queued_plugins_for_snapshot(str(snapshot.id)): - # Hook-level resume work is tracked by queued ArchiveResult rows, not by + # Plugin-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 hooks. + # active-state lock before running the remaining plugins. snapshot.update_and_requeue(retry_at=timezone.now()) return True diff --git a/archivebox/tests/test_api_v1_core_archiveresults.py b/archivebox/tests/test_api_v1_core_archiveresults.py index 09c3af9e..ea804dd6 100644 --- a/archivebox/tests/test_api_v1_core_archiveresults.py +++ b/archivebox/tests/test_api_v1_core_archiveresults.py @@ -1,6 +1,7 @@ from datetime import timedelta import pytest +from django.core.files.uploadedfile import SimpleUploadedFile from django.db import connection from django.test.utils import CaptureQueriesContext from django.utils import timezone @@ -13,6 +14,43 @@ from archivebox.tests.conftest import api_client_request pytestmark = pytest.mark.django_db(transaction=True) +def test_archiveresult_upload_upserts_one_row_per_plugin(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) + + first_response = client.post( + "/api/v1/core/archiveresults", + { + "snapshot_id": str(snapshot.id), + "plugin": "screenshot", + "hook_name": "on_Snapshot__archivebox_browser_extension_upload", + "files": SimpleUploadedFile("browser.png", b"browser", content_type="image/png"), + "output_paths": "browser.png", + }, + **api_headers, + ) + second_response = client.post( + "/api/v1/core/archiveresults", + { + "snapshot_id": str(snapshot.id), + "plugin": "screenshot", + "hook_name": "on_Snapshot__50_screenshot.py", + "files": SimpleUploadedFile("server.png", b"server", content_type="image/png"), + "output_paths": "server.png", + }, + **api_headers, + ) + + assert first_response.status_code == 200, first_response.content + assert second_response.status_code == 200, second_response.content + assert second_response.json()["id"] == first_response.json()["id"] + result = ArchiveResult.objects.get(snapshot=snapshot, plugin="screenshot") + assert result.hook_name == "on_Snapshot__50_screenshot.py" + assert set(result.output_files) == {"browser.png", "server.png"} + assert (result.output_dir / "browser.png").read_bytes() == b"browser" + assert (result.output_dir / "server.png").read_bytes() == b"server" + + def test_archiveresult_upload_api_queues_snapshot_maintenance_without_finalizing(client, api_admin_user, api_headers): crawl = Crawl.objects.create( urls="https://example.com", diff --git a/archivebox/tests/test_archive_result_service.py b/archivebox/tests/test_archive_result_service.py index 774b2ca1..8c51927c 100644 --- a/archivebox/tests/test_archive_result_service.py +++ b/archivebox/tests/test_archive_result_service.py @@ -199,7 +199,7 @@ def test_process_completed_projects_inline_archiveresult(tmp_path, hermetic_lib_ _cleanup_machine_process_rows() -def test_archiveresult_event_retry_updates_existing_hook_row(tmp_path, hermetic_lib_dir): +def test_archiveresult_events_from_multiple_hooks_update_one_plugin_row(tmp_path, hermetic_lib_dir): from archivebox.core.models import ArchiveResult snapshot = _create_snapshot() @@ -221,16 +221,61 @@ def test_archiveresult_event_retry_updates_existing_hook_row(tmp_path, hermetic_ snapshot, plugin="hashes", hook_name="on_Snapshot__93_hashes.py", + event_hook_name="on_Snapshot__99_hashes_followup.py", lib_dir=hermetic_lib_dir, env={"HASHES_ENABLED": "True"}, ) assert retry_result.id == first_result_id assert retry_result.status == ArchiveResult.StatusChoices.SUCCEEDED - assert ArchiveResult.objects.filter(snapshot=snapshot, plugin="hashes", hook_name="on_Snapshot__93_hashes.py").count() == 1 + assert retry_result.hook_name == "on_Snapshot__99_hashes_followup.py" + assert ArchiveResult.objects.filter(snapshot=snapshot, plugin="hashes").count() == 1 _cleanup_machine_process_rows() -def test_archiveresult_duplicate_hook_rows_are_rejected(): +def test_late_background_event_cannot_overwrite_plugin_result(): + from django.utils import timezone + from abx_dl.events import ArchiveResultEvent + from archivebox.core.models import ArchiveResult + from archivebox.services.archive_result_service import _save_archiveresult_event_to_db + + snapshot = _create_snapshot() + now = timezone.now().isoformat() + _save_archiveresult_event_to_db( + ArchiveResultEvent( + snapshot_id=str(snapshot.id), + plugin="chrome", + hook_name="on_Snapshot__30_chrome_navigate.js", + status="succeeded", + output_str="navigation complete", + output_files=[OutputFile(path="navigate.json", extension="json", mimetype="application/json", size=7)], + start_ts=now, + end_ts=now, + ), + None, + ) + _save_archiveresult_event_to_db( + ArchiveResultEvent( + snapshot_id=str(snapshot.id), + plugin="chrome", + hook_name="on_Snapshot__01_chrome_tab.daemon.bg.js", + status="skipped", + output_str="background cleanup", + output_files=[OutputFile(path="tab.json", extension="json", mimetype="application/json", size=3)], + start_ts=now, + end_ts=now, + ), + None, + ) + + result = ArchiveResult.objects.get(snapshot=snapshot, plugin="chrome") + assert result.hook_name == "on_Snapshot__30_chrome_navigate.js" + assert result.status == ArchiveResult.StatusChoices.SUCCEEDED + assert result.output_str == "navigation complete" + assert set(result.output_files) == {"navigate.json", "tab.json"} + assert result.output_size == 10 + + +def test_archiveresult_duplicate_plugin_rows_are_rejected(): from django.db import IntegrityError, transaction from archivebox.core.models import ArchiveResult @@ -246,11 +291,61 @@ def test_archiveresult_duplicate_hook_rows_are_rejected(): ArchiveResult.objects.create( snapshot=snapshot, plugin="wget", - hook_name="on_Snapshot__06_wget.finite.bg", + hook_name="on_Snapshot__99_other_wget_hook", status=ArchiveResult.StatusChoices.SUCCEEDED, ) +def test_pending_archiveresults_create_one_row_per_plugin(): + from archivebox.core.models import ArchiveResult + + snapshot = _create_snapshot() + results = snapshot.create_pending_archiveresults( + hooks=[ + ("chrome", "on_Snapshot__00_chrome_launch.daemon.bg.js"), + ("chrome", "on_Snapshot__01_chrome_tab.daemon.bg.js"), + ("chrome", "on_Snapshot__30_chrome_navigate.js"), + ], + ) + + assert len(results) == 1 + assert results[0].plugin == "chrome" + assert results[0].hook_name == "" + assert ArchiveResult.objects.filter(snapshot=snapshot, plugin="chrome").count() == 1 + + +def test_snapshot_merge_preserves_one_combined_plugin_result(): + from archivebox.core.models import ArchiveResult, Snapshot + + keeper = _create_snapshot() + duplicate = Snapshot.objects.create( + url="https://example.com/duplicate", + crawl=keeper.crawl, + status=Snapshot.StatusChoices.STARTED, + ) + ArchiveResult.objects.create( + snapshot=keeper, + plugin="responses", + status=ArchiveResult.StatusChoices.SUCCEEDED, + output_files={"browser.png": {"size": 7}}, + output_size=7, + ) + ArchiveResult.objects.create( + snapshot=duplicate, + plugin="responses", + status=ArchiveResult.StatusChoices.NORESULTS, + output_files={"server.json": {"size": 3}}, + output_size=3, + ) + + Snapshot._merge_snapshots([keeper, duplicate]) + + result = ArchiveResult.objects.get(snapshot=keeper, plugin="responses") + assert set(result.output_files) == {"browser.png", "server.json"} + assert result.output_size == 10 + assert not Snapshot.objects.filter(pk=duplicate.pk).exists() + + def test_process_completed_projects_failed_archiveresult_from_shipped_hook(tmp_path, hermetic_lib_dir): from archivebox.core.models import ArchiveResult @@ -434,12 +529,13 @@ 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", hook_name="on_Snapshot__11_chrome_wait") + result = ArchiveResult.objects.get(snapshot=snapshot, plugin="chrome") 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_run.py b/archivebox/tests/test_cli_run.py index 8efd66b2..7a9edec1 100644 --- a/archivebox/tests/test_cli_run.py +++ b/archivebox/tests/test_cli_run.py @@ -970,7 +970,7 @@ class TestRecoverOrchestratorState: assert snapshot.retry_at is None assert snapshot.downloaded_at is not None - def test_create_pending_archiveresults_uses_canonical_hook_names(self): + def test_create_pending_archiveresults_creates_one_plugin_row(self): from django.utils import timezone from archivebox.base_models.models import get_or_create_system_user_pk @@ -992,9 +992,10 @@ class TestRecoverOrchestratorState: snapshot.create_pending_archiveresults() - 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) + results = list(ArchiveResult.objects.filter(snapshot=snapshot)) + assert results + assert all(result.hook_name == "" for result in results) + assert len(results) == len({result.plugin for result in results}) def test_snapshot_hooks_for_pending_archiveresults_respects_disabled_plugins_when_plugins_empty(self): from django.utils import timezone @@ -1610,7 +1611,7 @@ class TestRecoverOrchestratorState: assert result.end_ts is not None @pytest.mark.django_db(transaction=True) - def test_run_due_snapshot_fails_obsolete_queued_hook_name(self): + def test_run_due_snapshot_fails_unavailable_queued_plugin(self): from django.utils import timezone from archivebox.base_models.models import get_or_create_system_user_pk @@ -1632,8 +1633,8 @@ class TestRecoverOrchestratorState: ) result = ArchiveResult.objects.create( snapshot=snapshot, - plugin="singlefile", - hook_name="on_Snapshot__50_singlefile.py", + plugin="removed_plugin", + hook_name="", status=ArchiveResult.StatusChoices.QUEUED, ) @@ -1642,7 +1643,7 @@ class TestRecoverOrchestratorState: result.refresh_from_db() snapshot.refresh_from_db() assert result.status == ArchiveResult.StatusChoices.FAILED - assert result.output_str == "Queued hook is no longer available in the installed plugin" + assert result.output_str == "Queued plugin is not available in this ArchiveBox installation" assert snapshot.retry_at is None @pytest.mark.django_db(transaction=True) @@ -1715,7 +1716,7 @@ class TestRecoverOrchestratorState: @pytest.mark.django_db(transaction=True) @pytest.mark.timeout(300) @pytest.mark.parametrize("chrome_isolation", ["crawl", "snapshot"]) - def test_resume_queued_chrome_navigate_reruns_background_prerequisites( + def test_retry_queued_chrome_result_reruns_full_plugin( self, initialized_archive, recursive_test_site, @@ -1755,21 +1756,18 @@ class TestRecoverOrchestratorState: ) assert list_process.returncode == 0, list_process.stderr or list_process.stdout chrome_results = parse_jsonl_output(list_process.stdout) - navigate_record = next(record for record in chrome_results if record["hook_name"] == "on_Snapshot__30_chrome_navigate") - snapshot_id = navigate_record["snapshot_id"] + assert len(chrome_results) == 1 + chrome_record = chrome_results[0] + snapshot_id = chrome_record["snapshot_id"] with use_archivebox_db(initialized_archive): - tab_result = ArchiveResult.objects.get( - snapshot_id=snapshot_id, - plugin="chrome", - hook_name="on_Snapshot__01_chrome_tab.daemon.bg", - ) - first_tab_process_id = tab_result.process_id - assert first_tab_process_id is not None + chrome_result = ArchiveResult.objects.get(snapshot_id=snapshot_id, plugin="chrome") + first_process_id = chrome_result.process_id + assert first_process_id is not None update_process = run_archivebox_cmd( ["archiveresult", "update", "--status=queued"], - stdin=next(line for line in list_process.stdout.splitlines() if navigate_record["id"] in line) + "\n", + stdin=next(line for line in list_process.stdout.splitlines() if chrome_record["id"] in line) + "\n", cwd=initialized_archive, env=env, timeout=60, @@ -1796,21 +1794,12 @@ class TestRecoverOrchestratorState: cleanup_process_group(run_process.pid) with use_archivebox_db(initialized_archive): - 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", - ) + chrome_result = ArchiveResult.objects.get(snapshot_id=snapshot_id, plugin="chrome") assert run_process.returncode == 0 - assert navigate_result.status == ArchiveResult.StatusChoices.SUCCEEDED - assert tab_result.process_id is not None - assert tab_result.process_id != first_tab_process_id + assert chrome_result.status == ArchiveResult.StatusChoices.SUCCEEDED + assert chrome_result.process_id is not None + assert chrome_result.process_id != first_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_unique_plugin.py b/archivebox/tests/test_migration_archiveresult_unique_plugin.py new file mode 100644 index 00000000..c3297056 --- /dev/null +++ b/archivebox/tests/test_migration_archiveresult_unique_plugin.py @@ -0,0 +1,53 @@ +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_ui_admin_archiveresult.py b/archivebox/tests/test_ui_admin_archiveresult.py index 63d90464..c40c60f0 100644 --- a/archivebox/tests/test_ui_admin_archiveresult.py +++ b/archivebox/tests/test_ui_admin_archiveresult.py @@ -141,29 +141,3 @@ class TestArchiveResultAdminListView: assert not output_dir.exists() snapshot.refresh_from_db() assert snapshot.output_size == 0 - - def test_deleting_duplicate_preserves_shared_plugin_output(self, snapshot): - from archivebox.core.models import ArchiveResult - - output_dir = Path(snapshot.output_dir) / "responses" - output_dir.mkdir(parents=True, exist_ok=True) - output_file = output_dir / "index.jsonl" - output_file.write_text("captured response\n") - primary = ArchiveResult.objects.create( - snapshot=snapshot, - plugin="responses", - hook_name="on_Snapshot__24_responses.daemon.bg", - status=ArchiveResult.StatusChoices.SUCCEEDED, - output_size=18, - ) - 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"