From 5b99586c1f744a12c9f66ff2bf4e07c957735f87 Mon Sep 17 00:00:00 2001 From: Nick Sweeting Date: Tue, 1 Sep 2026 11:44:54 -0700 Subject: [PATCH] Preserve published ArchiveResult migration history --- ...nique_archiveresult_per_snapshot_plugin.py | 71 +++++++++++++++++++ ...054_restore_archiveresult_hook_identity.py | 21 ++++++ ...t_migration_archiveresult_hook_identity.py | 31 ++++++++ 3 files changed, 123 insertions(+) create mode 100644 archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py create mode 100644 archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py create mode 100644 archivebox/tests/test_migration_archiveresult_hook_identity.py 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..11775c69 --- /dev/null +++ b/archivebox/core/migrations/0052_unique_archiveresult_per_snapshot_plugin.py @@ -0,0 +1,71 @@ +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"), + ) + 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()) + + 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), + ) + 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) + + +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/migrations/0054_restore_archiveresult_hook_identity.py b/archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py new file mode 100644 index 00000000..78eafcea --- /dev/null +++ b/archivebox/core/migrations/0054_restore_archiveresult_hook_identity.py @@ -0,0 +1,21 @@ +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("core", "0053_alter_archiveresult_options"), + ] + + operations = [ + migrations.RemoveConstraint( + model_name="archiveresult", + name="unique_archiveresult_per_snapshot_plugin", + ), + migrations.AddConstraint( + model_name="archiveresult", + constraint=models.UniqueConstraint( + fields=("snapshot", "plugin", "hook_name"), + name="unique_archiveresult_per_snapshot_hook", + ), + ), + ] 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..f1bb3b52 --- /dev/null +++ b/archivebox/tests/test_migration_archiveresult_hook_identity.py @@ -0,0 +1,31 @@ +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_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")])