Preserve ArchiveResult hook history through migrations

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):

archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py

@@ -1,6 +1,71 @@
+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"),
@@ -11,6 +76,10 @@ class Migration(migrations.Migration):
             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(

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."""

archivebox/tests/test_migration_archiveresult_hook_identity.py

@@ -6,6 +6,70 @@ 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)