__init__
This commit is contained in:
kstyagi23
+302
@@ -0,0 +1,302 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import zipfile
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.config import Settings
|
||||
from app.crawler import (
|
||||
DEFAULT_ASSET_TYPES,
|
||||
CaptureCancelled,
|
||||
CaptureError,
|
||||
CaptureOptions,
|
||||
SiteCrawler,
|
||||
)
|
||||
from app.db import Database
|
||||
from app.url_policy import NetworkRules, UrlPolicyError, resolve_and_validate
|
||||
|
||||
|
||||
def _remove_tree(path: Path, allowed_root: Path) -> None:
|
||||
"""Delete only a child path beneath the configured work root."""
|
||||
root = allowed_root.resolve()
|
||||
resolved = path.resolve()
|
||||
if resolved == root or not resolved.is_relative_to(root):
|
||||
raise RuntimeError("Refusing to remove a path outside SiteHarbor's work directory.")
|
||||
shutil.rmtree(resolved, ignore_errors=True)
|
||||
|
||||
|
||||
def _create_archive(
|
||||
source_dir: Path,
|
||||
destination: Path,
|
||||
archive_root: str,
|
||||
cancelled: Callable[[], bool],
|
||||
deadline_exceeded: Callable[[], bool],
|
||||
heartbeat: Callable[[], None],
|
||||
) -> int:
|
||||
partial = destination.with_suffix(".zip.partial")
|
||||
partial.unlink(missing_ok=True)
|
||||
last_heartbeat = time.monotonic()
|
||||
try:
|
||||
with zipfile.ZipFile(
|
||||
partial,
|
||||
mode="w",
|
||||
compression=zipfile.ZIP_DEFLATED,
|
||||
compresslevel=6,
|
||||
strict_timestamps=False,
|
||||
) as archive:
|
||||
for source_path in sorted(source_dir.rglob("*")):
|
||||
if cancelled():
|
||||
raise CaptureCancelled("Capture cancelled while packaging the archive")
|
||||
if deadline_exceeded():
|
||||
raise CaptureError(
|
||||
"The configured capture time limit was reached while packaging"
|
||||
)
|
||||
if time.monotonic() - last_heartbeat >= 5:
|
||||
heartbeat()
|
||||
last_heartbeat = time.monotonic()
|
||||
if not source_path.is_file() or source_path.is_symlink():
|
||||
continue
|
||||
relative = source_path.relative_to(source_dir).as_posix()
|
||||
with source_path.open("rb") as source_file:
|
||||
with archive.open(f"{archive_root}/{relative}", mode="w") as archive_file:
|
||||
while chunk := source_file.read(1_024 * 1_024):
|
||||
if cancelled():
|
||||
raise CaptureCancelled(
|
||||
"Capture cancelled while packaging the archive"
|
||||
)
|
||||
if deadline_exceeded():
|
||||
raise CaptureError(
|
||||
"The configured capture time limit was reached while packaging"
|
||||
)
|
||||
if time.monotonic() - last_heartbeat >= 5:
|
||||
heartbeat()
|
||||
last_heartbeat = time.monotonic()
|
||||
archive_file.write(chunk)
|
||||
os.replace(partial, destination)
|
||||
return destination.stat().st_size
|
||||
except Exception:
|
||||
partial.unlink(missing_ok=True)
|
||||
destination.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
class CaptureService:
|
||||
"""Runs one leased capture and persists a terminal SQLite job state."""
|
||||
|
||||
def __init__(self, database: Database, settings: Settings, worker_id: str) -> None:
|
||||
self.database = database
|
||||
self.settings = settings
|
||||
self.worker_id = worker_id
|
||||
|
||||
def run(self, job: dict[str, Any]) -> None:
|
||||
job_id = str(job["id"])
|
||||
config = dict(job["config"])
|
||||
attempt_id = uuid.uuid4().hex[:12]
|
||||
work_dir = self.settings.work_dir / f"{job_id}-{attempt_id}"
|
||||
artifact_path = self.settings.artifacts_dir / f"{job_id}-{attempt_id}.zip"
|
||||
report_path = self.settings.reports_dir / f"{job_id}-{attempt_id}.json"
|
||||
heartbeat_stop = threading.Event()
|
||||
lease_lost = threading.Event()
|
||||
heartbeat_thread: threading.Thread | None = None
|
||||
capture_deadline = time.monotonic() + int(config["max_duration_seconds"])
|
||||
|
||||
def capture_cancelled() -> bool:
|
||||
return (
|
||||
lease_lost.is_set()
|
||||
or self.database.is_cancel_requested(job_id)
|
||||
or not self.database.owns_lease(job_id, self.worker_id)
|
||||
)
|
||||
|
||||
def renew_lease() -> None:
|
||||
interval = min(20, max(3, self.settings.lease_seconds // 3))
|
||||
while not heartbeat_stop.wait(interval):
|
||||
try:
|
||||
if not self.database.renew_lease(
|
||||
job_id, self.worker_id, self.settings.lease_seconds
|
||||
):
|
||||
lease_lost.set()
|
||||
return
|
||||
except Exception:
|
||||
# Ownership checks stop this attempt if the lease eventually expires.
|
||||
continue
|
||||
|
||||
try:
|
||||
_remove_tree_if_present(work_dir, self.settings.work_dir)
|
||||
work_dir.mkdir(parents=True, exist_ok=True)
|
||||
if capture_cancelled():
|
||||
raise CaptureCancelled("Capture cancelled before it started")
|
||||
|
||||
heartbeat_thread = threading.Thread(
|
||||
target=renew_lease,
|
||||
name=f"siteharbor-lease-{job_id[-8:]}",
|
||||
daemon=True,
|
||||
)
|
||||
heartbeat_thread.start()
|
||||
|
||||
rules = NetworkRules(
|
||||
allow_private_networks=self.settings.allow_private_networks,
|
||||
allow_nonstandard_ports=self.settings.allow_nonstandard_ports,
|
||||
)
|
||||
resolve_and_validate(str(job["normalized_url"]), rules)
|
||||
remaining_seconds = int(capture_deadline - time.monotonic())
|
||||
if remaining_seconds <= 0:
|
||||
raise CaptureError(
|
||||
"The configured capture time limit was reached during validation"
|
||||
)
|
||||
options = CaptureOptions(
|
||||
source_url=str(job["normalized_url"]),
|
||||
include_external_assets=bool(config["include_external_assets"]),
|
||||
max_pages=int(config["max_pages"]),
|
||||
max_depth=int(config["max_depth"]),
|
||||
max_bytes=int(config["max_bytes"]),
|
||||
max_duration_seconds=remaining_seconds,
|
||||
parallel_connections=int(
|
||||
config.get("parallel_connections", self.settings.fetch_concurrency)
|
||||
),
|
||||
download_asset_types=tuple(
|
||||
config.get("download_asset_types", DEFAULT_ASSET_TYPES)
|
||||
),
|
||||
skip_file_types=tuple(config.get("skip_file_types", ())),
|
||||
use_proxy=bool(config.get("use_proxy", False)),
|
||||
)
|
||||
|
||||
def update(
|
||||
phase: str, message: str, stats: dict[str, int], warnings: list[str]
|
||||
) -> None:
|
||||
renewed = self.database.update_progress(
|
||||
job_id,
|
||||
self.worker_id,
|
||||
self.settings.lease_seconds,
|
||||
phase=phase,
|
||||
message=message,
|
||||
stats=stats,
|
||||
warnings=warnings,
|
||||
)
|
||||
if not renewed:
|
||||
raise CaptureCancelled("Capture lease was reassigned")
|
||||
|
||||
if not self.database.update_progress(
|
||||
job_id,
|
||||
self.worker_id,
|
||||
self.settings.lease_seconds,
|
||||
phase="Validating",
|
||||
message="Checking capture policy and target reachability",
|
||||
stats=_job_stats(job),
|
||||
warnings=[],
|
||||
):
|
||||
raise CaptureCancelled("Capture lease was reassigned")
|
||||
|
||||
crawler = SiteCrawler(
|
||||
settings=self.settings,
|
||||
options=options,
|
||||
cancelled=capture_cancelled,
|
||||
on_progress=update,
|
||||
)
|
||||
result = crawler.run(work_dir)
|
||||
|
||||
if capture_cancelled():
|
||||
raise CaptureCancelled("Capture cancelled before packaging")
|
||||
if time.monotonic() >= capture_deadline:
|
||||
raise CaptureError("The configured capture time limit was reached before packaging")
|
||||
|
||||
update("Packaging", "Creating a portable ZIP archive", result.stats, result.warnings)
|
||||
archive_size = _create_archive(
|
||||
result.site_dir,
|
||||
artifact_path,
|
||||
archive_root=f"siteharbor-{job_id}",
|
||||
cancelled=capture_cancelled,
|
||||
deadline_exceeded=lambda: time.monotonic() >= capture_deadline,
|
||||
heartbeat=lambda: update(
|
||||
"Packaging",
|
||||
"Creating a portable ZIP archive",
|
||||
result.stats,
|
||||
result.warnings,
|
||||
),
|
||||
)
|
||||
|
||||
report = {
|
||||
**result.report,
|
||||
"job_id": job_id,
|
||||
"archive": {
|
||||
"filename": artifact_path.name,
|
||||
"size_bytes": archive_size,
|
||||
"archive_root": f"siteharbor-{job_id}",
|
||||
},
|
||||
}
|
||||
if time.monotonic() >= capture_deadline:
|
||||
raise CaptureError(
|
||||
"The configured capture time limit was reached before completion"
|
||||
)
|
||||
report_path.write_text(
|
||||
json.dumps(report, indent=2, ensure_ascii=True), encoding="utf-8"
|
||||
)
|
||||
marked_ready = self.database.mark_ready(
|
||||
job_id,
|
||||
self.worker_id,
|
||||
artifact_path=str(artifact_path),
|
||||
report_path=str(report_path),
|
||||
archive_size=archive_size,
|
||||
entry_point=result.entry_point,
|
||||
stats=result.stats,
|
||||
warnings=result.warnings,
|
||||
retention_days=int(config["retention_days"]),
|
||||
)
|
||||
if not marked_ready:
|
||||
artifact_path.unlink(missing_ok=True)
|
||||
report_path.unlink(missing_ok=True)
|
||||
if self.database.is_cancel_requested(job_id):
|
||||
self.database.mark_cancelled(
|
||||
job_id,
|
||||
self.worker_id,
|
||||
"Capture cancelled before the archive became available",
|
||||
)
|
||||
except CaptureCancelled as error:
|
||||
artifact_path.unlink(missing_ok=True)
|
||||
report_path.unlink(missing_ok=True)
|
||||
if self.database.is_cancel_requested(job_id):
|
||||
self.database.mark_cancelled(job_id, self.worker_id, str(error))
|
||||
except UrlPolicyError as error:
|
||||
artifact_path.unlink(missing_ok=True)
|
||||
self.database.mark_failed(job_id, self.worker_id, "target_blocked", str(error))
|
||||
except CaptureError as error:
|
||||
artifact_path.unlink(missing_ok=True)
|
||||
self.database.mark_failed(job_id, self.worker_id, "capture_failed", str(error))
|
||||
except Exception as error:
|
||||
artifact_path.unlink(missing_ok=True)
|
||||
self.database.mark_failed(
|
||||
job_id,
|
||||
self.worker_id,
|
||||
"internal_error",
|
||||
f"Capture failed: {error}",
|
||||
)
|
||||
finally:
|
||||
heartbeat_stop.set()
|
||||
if heartbeat_thread:
|
||||
heartbeat_thread.join(timeout=1)
|
||||
try:
|
||||
_remove_tree_if_present(work_dir, self.settings.work_dir)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _job_stats(job: dict[str, Any]) -> dict[str, int]:
|
||||
return {
|
||||
"pages_found": int(job.get("pages_found", 0)),
|
||||
"pages_fetched": int(job.get("pages_fetched", 0)),
|
||||
"assets_found": int(job.get("assets_found", 0)),
|
||||
"assets_fetched": int(job.get("assets_fetched", 0)),
|
||||
"bytes_downloaded": int(job.get("bytes_downloaded", 0)),
|
||||
"files_written": int(job.get("files_written", 0)),
|
||||
}
|
||||
|
||||
|
||||
def _remove_tree_if_present(path: Path, root: Path) -> None:
|
||||
if path.exists():
|
||||
_remove_tree(path, root)
|
||||
Reference in New Issue
Block a user