The_Great_Archive/ArchiveBox
1
0
Fork 0
mirror of https://github.com/ArchiveBox/ArchiveBox.git synced 2026-09-12 19:50:57 +05:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Eliminate request-side SQLite write transactions

Browse Source
This commit is contained in:
Nick Sweeting 2026-08-28 15:26:50 -07:00
parent bdc236f076
commit c062a3c54f
No known key found for this signature in database
12 changed files with 385 additions and 202 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

222
archivebox/api/v1_core.py
View File

@ -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,

15
archivebox/core/admin_snapshots.py
View File

@ -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

116
archivebox/core/models.py
View File

@ -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

8
archivebox/core/recovery_util.py
View File

@ -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,
},

10
archivebox/core/tag_util.py
View File

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

2
archivebox/crawls/models.py
View File

@ -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

23
archivebox/services/archive_result_service.py
View File

@ -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 = []

8
archivebox/services/tag_service.py
View File

@ -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])

116
archivebox/tests/test_api_v1_core_archiveresults.py
View File

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

27
archivebox/tests/test_api_v1_core_snapshots.py
View File

@ -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 == []

20
archivebox/tests/test_api_v1_core_tags_add_to_snapshot.py
View File

@ -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 == []

20
archivebox/tests/test_api_v1_core_tags_remove_from_snapshot.py
View File

@ -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 == []