mirror of
https://github.com/ArchiveBox/ArchiveBox.git
synced 2026-09-12 19:50:57 +05:00
543 lines
19 KiB
Python
543 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""
|
|
archivebox snapshot <action> [args...] [--filters]
|
|
|
|
Manage Snapshot records.
|
|
|
|
Actions:
|
|
create - Create Snapshots from URLs or Crawl JSONL
|
|
list - List Snapshots as JSONL (with optional filters)
|
|
update - Update Snapshots from stdin JSONL
|
|
delete - Delete Snapshots from stdin JSONL
|
|
|
|
Examples:
|
|
# Create
|
|
archivebox snapshot create https://example.com --tag=news
|
|
archivebox crawl create https://example.com | archivebox snapshot create
|
|
|
|
# List with filters
|
|
archivebox snapshot list --status=queued
|
|
archivebox snapshot list --url__icontains=example.com
|
|
|
|
# Update
|
|
archivebox snapshot list --tag=old | archivebox snapshot update --tag=new
|
|
|
|
# Delete
|
|
archivebox snapshot list --url__icontains=spam.com | archivebox snapshot delete --yes
|
|
"""
|
|
|
|
__package__ = "archivebox.cli"
|
|
__command__ = "archivebox snapshot"
|
|
|
|
import sys
|
|
from collections.abc import Iterable, Iterator
|
|
from itertools import islice
|
|
|
|
import rich_click as click
|
|
from rich import print as rprint
|
|
from django.db.models import QuerySet
|
|
|
|
SNAPSHOT_FILTER_TYPE_CHOICES = ("exact", "substring", "regex", "domain", "tag", "timestamp")
|
|
SNAPSHOT_LIST_CHUNK_SIZE = 5000
|
|
|
|
|
|
def iter_snapshot_json(queryset: QuerySet) -> Iterator[dict[str, object]]:
|
|
from archivebox.config import VERSION
|
|
from archivebox.core.models import SnapshotTag
|
|
|
|
fields = (
|
|
"id",
|
|
"crawl_id",
|
|
"url",
|
|
"title",
|
|
"bookmarked_at",
|
|
"created_at",
|
|
"timestamp",
|
|
"depth",
|
|
"status",
|
|
"fs_version",
|
|
"output_size",
|
|
)
|
|
rows = queryset.values(*fields).iterator(chunk_size=SNAPSHOT_LIST_CHUNK_SIZE)
|
|
while batch := list(islice(rows, SNAPSHOT_LIST_CHUNK_SIZE)):
|
|
tags_by_snapshot = {row["id"]: [] for row in batch}
|
|
tag_rows = (
|
|
SnapshotTag.objects.filter(snapshot_id__in=tags_by_snapshot).order_by("tag__name").values_list("snapshot_id", "tag__name")
|
|
)
|
|
for snapshot_id, tag_name in tag_rows:
|
|
tags_by_snapshot[snapshot_id].append(tag_name)
|
|
|
|
for row in batch:
|
|
archive_size = int(row["output_size"] or 0)
|
|
yield {
|
|
"type": "Snapshot",
|
|
"schema_version": VERSION,
|
|
"id": str(row["id"]),
|
|
"crawl_id": str(row["crawl_id"]),
|
|
"url": row["url"],
|
|
"title": row["title"],
|
|
"tags": ",".join(sorted(tags_by_snapshot[row["id"]])),
|
|
"bookmarked_at": row["bookmarked_at"].isoformat() if row["bookmarked_at"] else None,
|
|
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
|
|
"timestamp": row["timestamp"],
|
|
"depth": row["depth"],
|
|
"status": row["status"],
|
|
"fs_version": row["fs_version"],
|
|
"archive_size": archive_size,
|
|
"output_size": archive_size,
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# CREATE
|
|
# =============================================================================
|
|
|
|
|
|
def create_snapshots(
|
|
urls: Iterable[str],
|
|
tag: str = "",
|
|
status: str = "queued",
|
|
depth: int = 0,
|
|
created_by_id: int | None = None,
|
|
) -> int:
|
|
"""
|
|
Create Snapshots from URLs or stdin JSONL (Crawl or Snapshot records).
|
|
Pass-through: Records that are not Crawl/Snapshot/URL are output unchanged.
|
|
|
|
Exit codes:
|
|
0: Success
|
|
1: Failure
|
|
"""
|
|
from archivebox.misc.jsonl import (
|
|
read_args_or_stdin,
|
|
write_record,
|
|
TYPE_SNAPSHOT,
|
|
TYPE_CRAWL,
|
|
)
|
|
from archivebox.misc.util import validate_url
|
|
from archivebox.base_models.models import get_or_create_system_user_pk
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.crawls.models import Crawl
|
|
|
|
created_by_id = created_by_id or get_or_create_system_user_pk()
|
|
is_tty = sys.stdout.isatty()
|
|
|
|
# Collect all input records
|
|
records = list(read_args_or_stdin(urls))
|
|
|
|
if not records:
|
|
rprint("[yellow]No URLs or Crawls provided. Pass URLs as arguments or via stdin.[/yellow]", file=sys.stderr)
|
|
return 1
|
|
|
|
# Process each record - handle Crawls and plain URLs/Snapshots
|
|
created_snapshots = []
|
|
pass_through_count = 0
|
|
|
|
for record in records:
|
|
record_type = record.get("type", "")
|
|
|
|
try:
|
|
if record_type == TYPE_CRAWL:
|
|
# Pass through the Crawl record itself first
|
|
if not is_tty:
|
|
write_record(record)
|
|
|
|
# Input is a Crawl - get or create it, then create Snapshots for its URLs
|
|
crawl = None
|
|
crawl_id = record.get("id")
|
|
if crawl_id:
|
|
try:
|
|
crawl = Crawl.objects.get(id=crawl_id)
|
|
except Crawl.DoesNotExist:
|
|
crawl = Crawl.from_json(record, overrides={"created_by_id": created_by_id})
|
|
else:
|
|
crawl = Crawl.from_json(record, overrides={"created_by_id": created_by_id})
|
|
|
|
if not crawl:
|
|
continue
|
|
|
|
# Create snapshots for each URL in the crawl
|
|
for url in crawl.get_urls_list():
|
|
try:
|
|
validate_url(url)
|
|
except ValueError as err:
|
|
rprint(f"[red]Error creating snapshot: {err}[/red]", file=sys.stderr)
|
|
continue
|
|
merged_tags = crawl.tags_str
|
|
if tag:
|
|
merged_tags = f"{merged_tags},{tag}" if merged_tags else tag
|
|
snapshot_record = {
|
|
"url": url,
|
|
"tags": merged_tags,
|
|
"crawl_id": str(crawl.id),
|
|
"depth": depth,
|
|
"status": status,
|
|
}
|
|
snapshot = Snapshot.from_json(snapshot_record, overrides={"created_by_id": created_by_id})
|
|
if snapshot:
|
|
created_snapshots.append(snapshot)
|
|
if not is_tty:
|
|
write_record(snapshot.to_json())
|
|
|
|
elif record_type == TYPE_SNAPSHOT or record.get("url"):
|
|
# Input is a Snapshot or plain URL
|
|
if record.get("url"):
|
|
validate_url(str(record["url"]))
|
|
if tag and not record.get("tags"):
|
|
record["tags"] = tag
|
|
if status:
|
|
record["status"] = status
|
|
record["depth"] = record.get("depth", depth)
|
|
|
|
snapshot = Snapshot.from_json(record, overrides={"created_by_id": created_by_id})
|
|
if snapshot:
|
|
created_snapshots.append(snapshot)
|
|
if not is_tty:
|
|
write_record(snapshot.to_json())
|
|
|
|
else:
|
|
# Pass-through: output records we don't handle
|
|
if not is_tty:
|
|
write_record(record)
|
|
pass_through_count += 1
|
|
|
|
except Exception as e:
|
|
rprint(f"[red]Error creating snapshot: {e}[/red]", file=sys.stderr)
|
|
continue
|
|
|
|
if not created_snapshots:
|
|
if pass_through_count > 0:
|
|
rprint(f"[dim]Passed through {pass_through_count} records, no new snapshots[/dim]", file=sys.stderr)
|
|
return 0
|
|
rprint("[red]No snapshots created[/red]", file=sys.stderr)
|
|
return 1
|
|
|
|
rprint(f"[green]Created {len(created_snapshots)} snapshots[/green]", file=sys.stderr)
|
|
|
|
if is_tty:
|
|
for snapshot in created_snapshots:
|
|
rprint(f" [dim]{snapshot.id}[/dim] {snapshot.url[:60]}", file=sys.stderr)
|
|
|
|
return 0
|
|
|
|
|
|
# =============================================================================
|
|
# LIST
|
|
# =============================================================================
|
|
|
|
|
|
def snapshot_filter_options(*, default_filter_type: str):
|
|
def decorate(func):
|
|
for decorator in reversed(
|
|
(
|
|
click.option("--status", "-s", help="Filter by status (queued, started, sealed)"),
|
|
click.option("--url__icontains", help="Filter by URL contains"),
|
|
click.option("--url__istartswith", help="Filter by URL starts with"),
|
|
click.option("--tag", "-t", help="Filter by tag name"),
|
|
click.option("--crawl-id", help="Filter by crawl ID"),
|
|
click.option("--limit", "-n", type=int, help="Limit number of results"),
|
|
click.option("--sort", "-o", type=str, help="Field to sort by, e.g. url, created_at, bookmarked_at, downloaded_at"),
|
|
click.option("--search", help="Search mode to use for positional query"),
|
|
click.option("--before", type=float, help="Only snapshots bookmarked before timestamp"),
|
|
click.option("--after", type=float, help="Only snapshots bookmarked after timestamp"),
|
|
click.option(
|
|
"--filter-type",
|
|
"-f",
|
|
type=click.Choice(SNAPSHOT_FILTER_TYPE_CHOICES),
|
|
default=default_filter_type,
|
|
help="Type of pattern matching to use for positional filters",
|
|
),
|
|
click.argument("filter_patterns", nargs=-1),
|
|
),
|
|
):
|
|
func = decorator(func)
|
|
return func
|
|
|
|
return decorate
|
|
|
|
|
|
def snapshot_output_options(func):
|
|
for decorator in reversed(
|
|
(
|
|
click.option("--csv", "-C", type=str, help="Print output as CSV with the provided fields, e.g.: timestamp,url,title"),
|
|
click.option("--json", "as_json", is_flag=True, help="Print output as a JSON array"),
|
|
click.option("--html", "as_html", is_flag=True, help="Print output as HTML"),
|
|
click.option("--with-headers", is_flag=True, help="Include column headers in structured output"),
|
|
),
|
|
):
|
|
func = decorator(func)
|
|
return func
|
|
|
|
|
|
def build_snapshot_queryset(
|
|
**kwargs,
|
|
) -> QuerySet:
|
|
from archivebox.core.models import Snapshot
|
|
|
|
return Snapshot.objects.order_by("-created_at").search(**kwargs)
|
|
|
|
|
|
def list_snapshots(
|
|
csv: str | None = None,
|
|
as_json: bool = False,
|
|
as_html: bool = False,
|
|
with_headers: bool = False,
|
|
**kwargs,
|
|
) -> int:
|
|
"""
|
|
List Snapshots as JSONL with optional filters.
|
|
|
|
Exit codes:
|
|
0: Success (even if no results)
|
|
"""
|
|
from archivebox.misc.jsonl import write_record
|
|
|
|
output_formats = sum(bool(output_format) for output_format in (csv, as_json, as_html))
|
|
if output_formats > 1:
|
|
rprint("[red]Choose only one output format: --csv, --json, or --html[/red]", file=sys.stderr)
|
|
return 2
|
|
if with_headers and not output_formats:
|
|
rprint("[red]--with-headers requires --csv, --json, or --html[/red]", file=sys.stderr)
|
|
return 2
|
|
|
|
is_tty = sys.stdout.isatty() and not output_formats
|
|
|
|
try:
|
|
queryset = build_snapshot_queryset(**kwargs)
|
|
except ValueError as err:
|
|
rprint(f"[red]{err}[/red]", file=sys.stderr)
|
|
return 2
|
|
|
|
count = 0
|
|
if as_json:
|
|
queryset = queryset.prefetch_related("tags")
|
|
output = queryset.to_json(with_headers=with_headers)
|
|
sys.stdout.write(output)
|
|
if output and not output.endswith("\n"):
|
|
sys.stdout.write("\n")
|
|
rprint(f"[dim]Listed {queryset.count()} snapshots[/dim]", file=sys.stderr)
|
|
return 0
|
|
|
|
if as_html:
|
|
queryset = queryset.prefetch_related("tags")
|
|
output = queryset.to_html(with_headers=with_headers)
|
|
sys.stdout.write(output)
|
|
if output and not output.endswith("\n"):
|
|
sys.stdout.write("\n")
|
|
rprint(f"[dim]Listed {queryset.count()} snapshots[/dim]", file=sys.stderr)
|
|
return 0
|
|
|
|
if csv:
|
|
cols = [col.strip() for col in csv.split(",") if col.strip()]
|
|
if not cols:
|
|
rprint("[red]No CSV columns provided[/red]", file=sys.stderr)
|
|
return 2
|
|
if with_headers:
|
|
sys.stdout.write(",".join(cols))
|
|
sys.stdout.write("\n")
|
|
for snapshot in queryset.prefetch_related("tags").iterator(chunk_size=SNAPSHOT_LIST_CHUNK_SIZE):
|
|
sys.stdout.write(snapshot.to_csv(cols=cols, separator=","))
|
|
sys.stdout.write("\n")
|
|
count += 1
|
|
rprint(f"[dim]Listed {count} snapshots[/dim]", file=sys.stderr)
|
|
return 0
|
|
|
|
if not is_tty:
|
|
for snapshot_json in iter_snapshot_json(queryset):
|
|
write_record(snapshot_json)
|
|
count += 1
|
|
rprint(f"[dim]Listed {count} snapshots[/dim]", file=sys.stderr)
|
|
return 0
|
|
|
|
for snapshot in queryset.iterator(chunk_size=SNAPSHOT_LIST_CHUNK_SIZE):
|
|
status_color = {
|
|
"queued": "yellow",
|
|
"started": "blue",
|
|
"sealed": "green",
|
|
}.get(snapshot.status, "dim")
|
|
rprint(f"[{status_color}]{snapshot.status:8}[/{status_color}] [dim]{snapshot.id}[/dim] {snapshot.url[:60]}")
|
|
count += 1
|
|
|
|
rprint(f"[dim]Listed {count} snapshots[/dim]", file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
# =============================================================================
|
|
# UPDATE
|
|
# =============================================================================
|
|
|
|
|
|
def update_snapshots(
|
|
status: str | None = None,
|
|
tag: str | None = None,
|
|
) -> int:
|
|
"""
|
|
Update Snapshots from stdin JSONL.
|
|
|
|
Reads Snapshot records from stdin and applies updates.
|
|
Uses PATCH semantics - only specified fields are updated.
|
|
|
|
Exit codes:
|
|
0: Success
|
|
1: No input or error
|
|
"""
|
|
from django.utils import timezone
|
|
|
|
from archivebox.misc.jsonl import read_stdin, write_record
|
|
from archivebox.core.models import Snapshot
|
|
|
|
is_tty = sys.stdout.isatty()
|
|
|
|
records = list(read_stdin())
|
|
if not records:
|
|
rprint("[yellow]No records provided via stdin[/yellow]", file=sys.stderr)
|
|
return 1
|
|
|
|
updated_count = 0
|
|
for record in records:
|
|
snapshot_id = record.get("id")
|
|
if not snapshot_id:
|
|
continue
|
|
|
|
try:
|
|
snapshot = Snapshot.objects.get(id=snapshot_id)
|
|
|
|
if status:
|
|
if status not in Snapshot.StatusChoices.values:
|
|
rprint(f"[red]Invalid snapshot status: {status}[/red]", file=sys.stderr)
|
|
continue
|
|
if status == Snapshot.StatusChoices.SEALED:
|
|
snapshot.cancel()
|
|
elif status == Snapshot.StatusChoices.PAUSED:
|
|
snapshot.pause()
|
|
elif status == Snapshot.StatusChoices.QUEUED:
|
|
if snapshot.status == Snapshot.StatusChoices.PAUSED:
|
|
snapshot.resume()
|
|
else:
|
|
snapshot.update_and_requeue(status=Snapshot.StatusChoices.QUEUED, retry_at=timezone.now())
|
|
elif status == Snapshot.StatusChoices.STARTED:
|
|
snapshot.update_and_requeue(status=Snapshot.StatusChoices.STARTED, retry_at=timezone.now())
|
|
if tag:
|
|
from archivebox.core.models import Tag
|
|
|
|
tag_obj, _ = Tag.objects.get_or_create(name=tag)
|
|
snapshot.tags.add(tag_obj)
|
|
snapshot.safe_update({"modified_at": timezone.now()}, refresh=False)
|
|
|
|
if not status and not tag:
|
|
snapshot.safe_update({"modified_at": timezone.now()}, refresh=False)
|
|
updated_count += 1
|
|
|
|
if not is_tty:
|
|
snapshot.refresh_from_db()
|
|
write_record(snapshot.to_json())
|
|
|
|
except Snapshot.DoesNotExist:
|
|
rprint(f"[yellow]Snapshot not found: {snapshot_id}[/yellow]", file=sys.stderr)
|
|
continue
|
|
|
|
rprint(f"[green]Updated {updated_count} snapshots[/green]", file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
# =============================================================================
|
|
# DELETE
|
|
# =============================================================================
|
|
|
|
|
|
def delete_snapshots(yes: bool = False, dry_run: bool = False) -> int:
|
|
"""
|
|
Delete Snapshots from stdin JSONL.
|
|
|
|
Requires --yes flag to confirm deletion.
|
|
|
|
Exit codes:
|
|
0: Success
|
|
1: No input or missing --yes flag
|
|
"""
|
|
from archivebox.misc.jsonl import read_stdin
|
|
from archivebox.core.models import Snapshot
|
|
|
|
records = list(read_stdin())
|
|
if not records:
|
|
rprint("[yellow]No records provided via stdin[/yellow]", file=sys.stderr)
|
|
return 1
|
|
|
|
snapshot_ids = [r.get("id") for r in records if r.get("id")]
|
|
|
|
if not snapshot_ids:
|
|
rprint("[yellow]No valid snapshot IDs in input[/yellow]", file=sys.stderr)
|
|
return 1
|
|
|
|
snapshots = Snapshot.objects.filter(id__in=snapshot_ids)
|
|
count = snapshots.count()
|
|
|
|
if count == 0:
|
|
rprint("[yellow]No matching snapshots found[/yellow]", file=sys.stderr)
|
|
return 0
|
|
|
|
if dry_run:
|
|
rprint(f"[yellow]Would delete {count} snapshots (dry run)[/yellow]", file=sys.stderr)
|
|
for snapshot in snapshots:
|
|
rprint(f" [dim]{snapshot.id}[/dim] {snapshot.url[:60]}", file=sys.stderr)
|
|
return 0
|
|
|
|
if not yes:
|
|
rprint("[red]Use --yes to confirm deletion[/red]", file=sys.stderr)
|
|
return 1
|
|
|
|
# Perform deletion
|
|
deleted_count, _ = snapshots.delete()
|
|
rprint(f"[green]Deleted {deleted_count} snapshots[/green]", file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
# =============================================================================
|
|
# CLI Commands
|
|
# =============================================================================
|
|
|
|
|
|
@click.group()
|
|
def main():
|
|
"""Manage Snapshot records."""
|
|
pass
|
|
|
|
|
|
@main.command("create")
|
|
@click.argument("urls", nargs=-1)
|
|
@click.option("--tag", "-t", default="", help="Comma-separated tags to add")
|
|
@click.option("--status", "-s", default="queued", help="Initial status (default: queued)")
|
|
@click.option("--depth", "-d", type=int, default=0, help="Crawl depth (default: 0)")
|
|
def create_cmd(urls: tuple, tag: str, status: str, depth: int):
|
|
"""Create Snapshots from URLs or stdin JSONL."""
|
|
sys.exit(create_snapshots(urls, tag=tag, status=status, depth=depth))
|
|
|
|
|
|
@main.command("list")
|
|
@snapshot_output_options
|
|
@snapshot_filter_options(default_filter_type="substring")
|
|
def list_cmd(**kwargs):
|
|
"""List Snapshots as JSONL."""
|
|
sys.exit(list_snapshots(**kwargs))
|
|
|
|
|
|
@main.command("update")
|
|
@click.option("--status", "-s", help="Set status")
|
|
@click.option("--tag", "-t", help="Add tag")
|
|
def update_cmd(status: str | None, tag: str | None):
|
|
"""Update Snapshots from stdin JSONL."""
|
|
sys.exit(update_snapshots(status=status, tag=tag))
|
|
|
|
|
|
@main.command("delete")
|
|
@click.option("--yes", "-y", is_flag=True, help="Confirm deletion")
|
|
@click.option("--dry-run", is_flag=True, help="Show what would be deleted")
|
|
def delete_cmd(yes: bool, dry_run: bool):
|
|
"""Delete Snapshots from stdin JSONL."""
|
|
sys.exit(delete_snapshots(yes=yes, dry_run=dry_run))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|