From 2bda2abbf0a30b7975e4615c7a04f45ed277cf82 Mon Sep 17 00:00:00 2001 From: Nick Sweeting Date: Wed, 10 Jun 2026 21:04:48 -0700 Subject: [PATCH] Fix tagged internal add crawl sealing --- archivebox/core/models.py | 19 ++++++++--- archivebox/crawls/models.py | 4 +-- archivebox/services/runner.py | 3 ++ archivebox/tests/test_cli_add.py | 58 +++++++++++++++++++++++++++++++- 4 files changed, 76 insertions(+), 8 deletions(-) diff --git a/archivebox/core/models.py b/archivebox/core/models.py index 42d97603..6fc7cf37 100755 --- a/archivebox/core/models.py +++ b/archivebox/core/models.py @@ -136,7 +136,7 @@ class Tag(ModelWithUUID): # Auto-attach to snapshot if in overrides if overrides and "snapshot" in overrides and tag: - overrides["snapshot"].tags.add(tag) + overrides["snapshot"].add_tag_ids([tag.pk]) return tag @@ -571,6 +571,15 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW @property def sm(self) -> "SnapshotMachine": ... + 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] + 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, + ) + class Meta( ModelWithDeleteAfter.Meta, ModelWithOutputDir.Meta, @@ -974,7 +983,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW if missing_tags: Tag.objects.bulk_create(missing_tags, ignore_conflicts=True) tags_by_name = {tag.name: tag for tag in Tag.objects.filter(name__in=crawl_tag_names)} - self.tags.add(*[tag.pk for name in crawl_tag_names if (tag := tags_by_name.get(name))]) + self.add_tag_ids([tag.pk for name in crawl_tag_names if (tag := tags_by_name.get(name))]) # Snapshot.save() normally appends newly created URLs to Crawl.urls # so legacy/direct crawls can keep their queue text in sync. For # internal-input crawls that would corrupt the original submitted @@ -1690,7 +1699,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW with transaction.atomic(): for tag_name in new_tags: tag, _ = Tag.objects.get_or_create(name=tag_name) - self.tags.add(tag) + 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.""" @@ -2136,7 +2145,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW # Merge tags for tag in dup.tags.all(): - keeper.tags.add(tag) + keeper.add_tag_ids([tag.pk]) # Move ArchiveResults ArchiveResult.objects.filter(snapshot=dup).update( @@ -2584,7 +2593,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW 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.tags.add(*tags_id) + self.add_tag_ids(tags_id) def pending_archiveresults(self) -> QuerySet["ArchiveResult"]: return self.archiveresult_set.exclude(status__in=ArchiveResult.FINAL_OR_ACTIVE_STATES) diff --git a/archivebox/crawls/models.py b/archivebox/crawls/models.py index 482b6efa..d018fa7d 100755 --- a/archivebox/crawls/models.py +++ b/archivebox/crawls/models.py @@ -1082,7 +1082,7 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith # relation without clearing any non-crawl snapshot tags. Tag.objects.bulk_create(missing_tags, ignore_conflicts=True) tags_by_name.update({tag.name: tag for tag in Tag.objects.filter(name__in=missing_names)}) - snapshot.tags.add(*[tag.pk for tag_name in tag_names if (tag := tags_by_name.get(tag_name))]) + snapshot.add_tag_ids([tag.pk for tag_name in tag_names if (tag := tags_by_name.get(tag_name))]) # Symlink creation touches the filesystem and can be slow on remote disks. # Defer it until after any active DB transaction commits so SQLite does @@ -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] - snapshot.tags.add(*tag_ids) + snapshot.add_tag_ids(tag_ids) existing_scope = Snapshot.objects if bool(self._config_value(config, "ONLY_NEW", True)) else self.snapshot_set existing_urls = set(existing_scope.filter(url__in=deduped_records.keys()).values_list("url", flat=True)) diff --git a/archivebox/services/runner.py b/archivebox/services/runner.py index 1010582f..72b1200c 100644 --- a/archivebox/services/runner.py +++ b/archivebox/services/runner.py @@ -1409,6 +1409,7 @@ def include_background_prerequisite_hooks( def snapshot_hooks_for_pending_archiveresults(snapshot) -> list[tuple[str, str]]: from archivebox.config.common import get_config + from archivebox.core.models import Snapshot config = get_config(crawl=snapshot.crawl, snapshot=snapshot) snapshot_plugin_names = [name.strip() for name in str((snapshot.config or {}).get("PLUGINS") or "").split(",") if name.strip()] @@ -1420,6 +1421,8 @@ def snapshot_hooks_for_pending_archiveresults(snapshot) -> list[tuple[str, str]] if plugin_names else _discover_archivebox_plugins() ) + if snapshot.url == Snapshot.INTERNAL_INPUT_URL: + plugins = {name: plugin for name, plugin in plugins.items() if getattr(plugin.config, "x_accepts_internal_input", False)} return sorted((plugin.name, hook.name) for plugin in plugins.values() for hook in plugin.filter_hooks("Snapshot")) diff --git a/archivebox/tests/test_cli_add.py b/archivebox/tests/test_cli_add.py index 6de356f4..8201564f 100644 --- a/archivebox/tests/test_cli_add.py +++ b/archivebox/tests/test_cli_add.py @@ -12,7 +12,7 @@ import pytest from django.db import connection from django.utils import timezone -from archivebox.core.models import ArchiveResult, Snapshot +from archivebox.core.models import ArchiveResult, Snapshot, SnapshotTag from archivebox.crawls.models import Crawl from archivebox.machine.models import Process from archivebox.tests.conftest import ( @@ -486,6 +486,62 @@ def test_add_bg_queues_internal_input_root_snapshot(initialized_archive): assert root_input == "https://example.com" +@pytest.mark.timeout(180) +def test_add_tagged_single_url_seals_without_duplicate_snapshot_tags(initialized_archive): + env = os.environ.copy() + env.update( + { + "USE_COLOR": "False", + "SHOW_PROGRESS": "False", + "PLUGINS": "parse_txt_urls,title", + "TIMEOUT": "60", + "CRAWL_MAX_CONCURRENT_SNAPSHOTS": "1", + }, + ) + result = run_archivebox_cmd( + [ + "add", + "--bg", + "--depth=0", + "--tag=tagged-single-url", + "https://example.com/?archivebox-tagged-single-url=1", + ], + cwd=initialized_archive, + env=env, + timeout=60, + ) + assert result.returncode == 0, result.stderr or result.stdout + + run_queued_crawls(initialized_archive, env, timeout=180) + + with use_archivebox_db(initialized_archive): + crawl = Crawl.objects.get() + snapshots = list(Snapshot.objects.order_by("depth", "url")) + tag_counts = { + snapshot.url: SnapshotTag.objects.filter(snapshot=snapshot, tag__name="tagged-single-url").count() for snapshot in snapshots + } + results = list( + ArchiveResult.objects.select_related("snapshot") + .order_by("snapshot__depth", "snapshot__url", "plugin") + .values_list("snapshot__url", "plugin", "status", "output_str"), + ) + + assert crawl.status == Crawl.StatusChoices.SEALED + assert [(snapshot.url, snapshot.depth, snapshot.status) for snapshot in snapshots] == [ + (Snapshot.INTERNAL_INPUT_URL, 0, Snapshot.StatusChoices.SEALED), + ("https://example.com/?archivebox-tagged-single-url=1", 1, Snapshot.StatusChoices.SEALED), + ] + assert tag_counts == { + Snapshot.INTERNAL_INPUT_URL: 1, + "https://example.com/?archivebox-tagged-single-url=1": 1, + } + by_url_plugin = {(url, plugin): status for url, plugin, status, _output in results} + assert by_url_plugin[(Snapshot.INTERNAL_INPUT_URL, "parse_txt_urls")] == "succeeded" + assert by_url_plugin[("https://example.com/?archivebox-tagged-single-url=1", "title")] == "succeeded" + unexpected_failures = [(url, plugin, status, output) for url, plugin, status, output in results if status == "failed"] + assert not unexpected_failures + + def test_add_index_only_rejected_urls_leave_empty_crawl_for_runner_to_seal(initialized_archive): """Index-only add only creates the crawl; rejected URLs are sealed by the runner.""" env = cli_env(disable_extractors=True)