diff --git a/archivebox/api/v1_core.py b/archivebox/api/v1_core.py index 36902152..c426d41d 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 @@ -533,51 +533,48 @@ def create_archiveresult( "plugin": plugin_name, "hook_name": hook, } - existing_result = ArchiveResult.objects.filter(**result_lookup).first() - existing_output_files = dict(existing_result.output_files or {}) if existing_result else {} - output_files = _write_archiveresult_files( + 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(**result_lookup).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.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) @@ -591,44 +588,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) @@ -945,52 +944,41 @@ 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 - - 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, - ) - - 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) - + snapshot = Snapshot.objects.filter(url=data.url, crawl=crawl).first() + if snapshot is None: try: - snapshot.ensure_crawl_symlink() - except Exception: - pass + 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 + + 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) + + try: + snapshot.ensure_crawl_symlink() + except Exception: + pass setattr(request, "with_archiveresults", False) return snapshot @@ -1425,8 +1413,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, @@ -1454,8 +1441,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/core/admin_snapshots.py b/archivebox/core/admin_snapshots.py index 45687c52..6e00575b 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 diff --git a/archivebox/core/models.py b/archivebox/core/models.py index 02719146..8855670e 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 @@ -88,6 +88,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): @@ -142,7 +155,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: @@ -591,13 +604,23 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW 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] + for tag_id in tag_ids: + try: + SnapshotTag(snapshot_id=self.pk, tag_id=tag_id).save(force_insert=True) + except IntegrityError: + # The unique (snapshot, tag) row already exists. In autocommit + # mode the failed INSERT is fully rolled back before continuing. + continue + + 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, @@ -1725,8 +1748,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()} @@ -1734,10 +1755,9 @@ 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 hook; retries update the existing row.""" @@ -2707,9 +2727,12 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW 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) + from archivebox.core.tag_util import get_or_create_tag + + tag_ids = {get_or_create_tag(tag)[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) @@ -3000,10 +3023,10 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW 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.objects.get_or_create( - snapshot=self, - plugin=plugin, - hook_name=hook_name, + archiveresult, _created = ArchiveResult.get_or_create_by_hook( + self, + plugin, + hook_name, defaults={ "status": ArchiveResult.INITIAL_STATE, }, @@ -3830,6 +3853,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: @@ -4080,10 +4124,10 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes): try: snapshot = Snapshot.objects.get(id=snapshot_id) - result, _ = ArchiveResult.objects.get_or_create( - snapshot=snapshot, - plugin=plugin, - hook_name=record.get("hook_name", ""), + result, _ = ArchiveResult.get_or_create_by_hook( + snapshot, + plugin, + record.get("hook_name", ""), defaults={ "status": record.get("status", "queued"), "output_str": record.get("output_str", ""), @@ -4148,6 +4192,30 @@ 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 diff --git a/archivebox/core/recovery_util.py b/archivebox/core/recovery_util.py index 22bda353..3403ead5 100644 --- a/archivebox/core/recovery_util.py +++ b/archivebox/core/recovery_util.py @@ -201,10 +201,10 @@ 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, - hook_name=hook_name, + result, created = ArchiveResult.get_or_create_by_hook( + snapshot, + plugin_dir.name, + hook_name, defaults={ "status": ArchiveResult.StatusChoices.QUEUED, }, diff --git a/archivebox/core/tag_util.py b/archivebox/core/tag_util.py index 9fbeea22..2336b1f7 100644 --- a/archivebox/core/tag_util.py +++ b/archivebox/core/tag_util.py @@ -152,15 +152,7 @@ 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 + return Tag.get_or_create_by_name(normalized_name, defaults={"created_by": created_by}) def rename_tag(tag: Tag, name: str) -> Tag: diff --git a/archivebox/crawls/models.py b/archivebox/crawls/models.py index b92fd456..a1af3793 100755 --- a/archivebox/crawls/models.py +++ b/archivebox/crawls/models.py @@ -1194,7 +1194,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 0ed9ff31..6689fa6e 100644 --- a/archivebox/services/archive_result_service.py +++ b/archivebox/services/archive_result_service.py @@ -15,7 +15,6 @@ from pathlib import Path from typing import Any, Protocol, runtime_checkable 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 @@ -340,21 +339,13 @@ def _save_archiveresult_event_to_db( hook_name=event.hook_name, ).first() if result is None: - try: - with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_create"): - result = ArchiveResult.objects.create( - snapshot=snapshot, - plugin=event.plugin, - hook_name=event.hook_name, - **defaults, - ) - except IntegrityError: - with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_get_after_integrity"): - result = ArchiveResult.objects.get( - snapshot=snapshot, - plugin=event.plugin, - hook_name=event.hook_name, - ) + with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.result_create"): + result, _created = ArchiveResult.get_or_create_by_hook( + snapshot, + event.plugin, + event.hook_name, + defaults=defaults, + ) with _perf_span("archivebox.ArchiveResultService.on_ArchiveResultEvent.diff_fields"): update_fields = [] 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 07ccc4a4..01a5e687 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 @@ -47,6 +48,121 @@ def test_archiveresult_upload_upserts_by_snapshot_plugin_and_hook(client, api_ad 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 + 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..4c01b820 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,28 @@ 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() + 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_api_v1_core_tags_add_to_snapshot.py b/archivebox/tests/test_api_v1_core_tags_add_to_snapshot.py index c732ca39..2f2f5221 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 connection +from django.test.utils import CaptureQueriesContext from archivebox.core.models import Snapshot, Tag from archivebox.crawls.models import Crawl @@ -13,13 +15,17 @@ 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() + 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_api_v1_core_tags_remove_from_snapshot.py b/archivebox/tests/test_api_v1_core_tags_remove_from_snapshot.py index 2b56770c..54ed29b4 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,17 @@ 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() + transaction_queries = [query["sql"] for query in queries if query["sql"].strip().upper() in {"BEGIN", "COMMIT"}] + assert transaction_queries == []