Files
website-downloader/app/db.py
T

700 lines
26 KiB
Python

from __future__ import annotations
import json
import sqlite3
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any
from app.config import MAX_ARCHIVE_RETENTION_DAYS
def utcnow() -> datetime:
return datetime.now(UTC)
def timestamp(value: datetime | None = None) -> str:
return (value or utcnow()).isoformat()
class Database:
"""Small SQLite persistence layer with lease-based job claiming."""
def __init__(self, path: Path) -> None:
self.path = path
@contextmanager
def connection(self) -> Iterator[sqlite3.Connection]:
connection = sqlite3.connect(self.path, timeout=15, isolation_level=None)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA foreign_keys = ON")
connection.execute("PRAGMA busy_timeout = 15000")
try:
yield connection
finally:
connection.close()
def initialize(self) -> None:
with self.connection() as connection:
connection.execute("PRAGMA journal_mode = WAL")
connection.execute("PRAGMA synchronous = NORMAL")
connection.executescript(
"""
CREATE TABLE IF NOT EXISTS jobs (
id TEXT PRIMARY KEY,
source_url TEXT NOT NULL,
normalized_url TEXT NOT NULL,
owner_token TEXT NOT NULL DEFAULT '',
config_json TEXT NOT NULL,
state TEXT NOT NULL,
phase TEXT NOT NULL,
message TEXT NOT NULL,
error_code TEXT,
error_message TEXT,
warnings_json TEXT NOT NULL DEFAULT '[]',
pages_found INTEGER NOT NULL DEFAULT 0,
pages_fetched INTEGER NOT NULL DEFAULT 0,
assets_found INTEGER NOT NULL DEFAULT 0,
assets_fetched INTEGER NOT NULL DEFAULT 0,
bytes_downloaded INTEGER NOT NULL DEFAULT 0,
files_written INTEGER NOT NULL DEFAULT 0,
artifact_path TEXT,
report_path TEXT,
archive_size INTEGER,
entry_point TEXT,
cancel_requested INTEGER NOT NULL DEFAULT 0,
worker_id TEXT,
lease_expires_at TEXT,
attempt_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
started_at TEXT,
completed_at TEXT,
expires_at TEXT
);
CREATE INDEX IF NOT EXISTS jobs_state_created_idx
ON jobs(state, created_at);
CREATE INDEX IF NOT EXISTS jobs_lease_idx
ON jobs(lease_expires_at);
CREATE TABLE IF NOT EXISTS job_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
job_id TEXT NOT NULL REFERENCES jobs(id) ON DELETE CASCADE,
kind TEXT NOT NULL,
payload_json TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS job_events_job_id_idx
ON job_events(job_id, id);
"""
)
columns = {
str(row["name"])
for row in connection.execute("PRAGMA table_info(jobs)").fetchall()
}
if "owner_token" not in columns:
connection.execute(
"ALTER TABLE jobs ADD COLUMN owner_token TEXT NOT NULL DEFAULT ''"
)
connection.execute(
"CREATE INDEX IF NOT EXISTS jobs_owner_created_idx "
"ON jobs(owner_token, created_at DESC)"
)
def create_job(
self,
source_url: str,
normalized_url: str,
config: dict[str, Any],
owner_token: str = "",
) -> dict[str, Any]:
job_id = f"cap_{uuid.uuid4().hex}"
now = timestamp()
payload = json.dumps(config, separators=(",", ":"))
with self.connection() as connection:
connection.execute(
"""
INSERT INTO jobs (
id, source_url, normalized_url, owner_token, config_json, state, phase, message,
created_at, updated_at, expires_at
) VALUES (
?, ?, ?, ?, ?, 'queued', 'Queued', 'Waiting for a local worker', ?, ?, NULL
)
""",
(
job_id,
source_url,
normalized_url,
owner_token,
payload,
now,
now,
),
)
self._append_event(
connection,
job_id,
"state",
{"state": "queued", "phase": "Queued", "message": "Waiting for a local worker"},
)
return self.get_job(job_id) or {}
def get_job(self, job_id: str) -> dict[str, Any] | None:
with self.connection() as connection:
row = connection.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
return self._row_to_job(row) if row else None
def get_owned_job(self, job_id: str, owner_token: str) -> dict[str, Any] | None:
with self.connection() as connection:
row = connection.execute(
"SELECT * FROM jobs WHERE id = ? AND owner_token = ?", (job_id, owner_token)
).fetchone()
return self._row_to_job(row) if row else None
def list_jobs(self, limit: int = 30) -> list[dict[str, Any]]:
with self.connection() as connection:
rows = connection.execute(
"SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?", (limit,)
).fetchall()
return [self._row_to_job(row) for row in rows]
def list_owned_jobs(self, owner_token: str, limit: int = 30) -> list[dict[str, Any]]:
with self.connection() as connection:
rows = connection.execute(
"""
SELECT * FROM jobs
WHERE owner_token = ?
ORDER BY created_at DESC
LIMIT ?
""",
(owner_token, limit),
).fetchall()
return [self._row_to_job(row) for row in rows]
def public_stats(self) -> dict[str, int]:
with self.connection() as connection:
row = connection.execute(
"""
SELECT
COUNT(CASE WHEN state = 'ready' THEN 1 END) AS websites_cloned,
COALESCE(SUM(bytes_downloaded), 0) AS bytes_scraped,
COUNT(CASE WHEN state IN ('queued', 'running') THEN 1 END) AS active_crawls,
COALESCE(SUM(files_written), 0) AS files_archived
FROM jobs
"""
).fetchone()
return {
"websites_cloned": int(row["websites_cloned"] or 0),
"bytes_scraped": int(row["bytes_scraped"] or 0),
"active_crawls": int(row["active_crawls"] or 0),
"files_archived": int(row["files_archived"] or 0),
}
def get_events(self, job_id: str, after_id: int = 0) -> list[dict[str, Any]]:
with self.connection() as connection:
rows = connection.execute(
"""
SELECT id, kind, payload_json, created_at
FROM job_events
WHERE job_id = ? AND id > ?
ORDER BY id ASC
LIMIT 100
""",
(job_id, after_id),
).fetchall()
return [
{
"id": row["id"],
"kind": row["kind"],
"payload": json.loads(row["payload_json"]),
"created_at": row["created_at"],
}
for row in rows
]
def claim_next(self, worker_id: str, lease_seconds: int) -> dict[str, Any] | None:
now = utcnow()
now_text = timestamp(now)
lease = timestamp(now + timedelta(seconds=lease_seconds))
with self.connection() as connection:
try:
connection.execute("BEGIN IMMEDIATE")
cancelled_rows = connection.execute(
"""
SELECT id FROM jobs
WHERE cancel_requested = 1
AND (
state = 'queued'
OR (
state = 'running'
AND lease_expires_at IS NOT NULL
AND lease_expires_at < ?
)
)
""",
(now_text,),
).fetchall()
connection.execute(
"""
UPDATE jobs
SET state = 'cancelled', phase = 'Cancelled',
message = 'Cancelled before starting',
completed_at = ?, updated_at = ?
WHERE cancel_requested = 1
AND (
state = 'queued'
OR (
state = 'running'
AND lease_expires_at IS NOT NULL
AND lease_expires_at < ?
)
)
""",
(now_text, now_text, now_text),
)
for cancelled_row in cancelled_rows:
self._append_event(
connection,
cancelled_row["id"],
"state",
{
"state": "cancelled",
"phase": "Cancelled",
"message": "Cancelled before starting",
},
)
row = connection.execute(
"""
SELECT * FROM jobs
WHERE cancel_requested = 0
AND (
state = 'queued'
OR (
state = 'running'
AND lease_expires_at IS NOT NULL
AND lease_expires_at < ?
)
)
ORDER BY created_at ASC
LIMIT 1
""",
(now_text,),
).fetchone()
if not row:
connection.execute("COMMIT")
return None
job_id = row["id"]
message = (
"Recovering a leased capture"
if row["state"] == "running"
else "Preparing capture"
)
connection.execute(
"""
UPDATE jobs
SET state = 'running', phase = 'Preparing', message = ?, worker_id = ?,
lease_expires_at = ?, attempt_count = attempt_count + 1,
started_at = COALESCE(started_at, ?), updated_at = ?
WHERE id = ?
""",
(message, worker_id, lease, now_text, now_text, job_id),
)
self._append_event(
connection,
job_id,
"state",
{"state": "running", "phase": "Preparing", "message": message},
)
claimed = connection.execute(
"SELECT * FROM jobs WHERE id = ?", (job_id,)
).fetchone()
connection.execute("COMMIT")
return self._row_to_job(claimed)
except Exception:
connection.execute("ROLLBACK")
raise
def update_progress(
self,
job_id: str,
worker_id: str,
lease_seconds: int,
*,
phase: str,
message: str,
stats: dict[str, int],
warnings: list[str] | None = None,
) -> bool:
now = utcnow()
values = {
"phase": phase,
"message": message,
"pages_found": stats.get("pages_found", 0),
"pages_fetched": stats.get("pages_fetched", 0),
"assets_found": stats.get("assets_found", 0),
"assets_fetched": stats.get("assets_fetched", 0),
"bytes_downloaded": stats.get("bytes_downloaded", 0),
"files_written": stats.get("files_written", 0),
"warnings_json": json.dumps(warnings or []),
"worker_id": worker_id,
"lease_expires_at": timestamp(now + timedelta(seconds=lease_seconds)),
"updated_at": timestamp(now),
"current_time": timestamp(now),
}
with self.connection() as connection:
cursor = connection.execute(
"""
UPDATE jobs SET
phase = :phase, message = :message,
pages_found = :pages_found, pages_fetched = :pages_fetched,
assets_found = :assets_found, assets_fetched = :assets_fetched,
bytes_downloaded = :bytes_downloaded, files_written = :files_written,
warnings_json = :warnings_json, worker_id = :worker_id,
lease_expires_at = :lease_expires_at, updated_at = :updated_at
WHERE id = :job_id
AND state = 'running'
AND worker_id = :worker_id
AND lease_expires_at >= :current_time
""",
{**values, "job_id": job_id},
)
if cursor.rowcount:
self._append_event(
connection,
job_id,
"progress",
{
"phase": phase,
"message": message,
"stats": stats,
"warnings": warnings or [],
},
)
return True
return False
def renew_lease(self, job_id: str, worker_id: str, lease_seconds: int) -> bool:
now = utcnow()
now_text = timestamp(now)
with self.connection() as connection:
cursor = connection.execute(
"""
UPDATE jobs
SET lease_expires_at = ?, updated_at = ?
WHERE id = ?
AND state = 'running'
AND worker_id = ?
AND lease_expires_at >= ?
""",
(
timestamp(now + timedelta(seconds=lease_seconds)),
now_text,
job_id,
worker_id,
now_text,
),
)
return bool(cursor.rowcount)
def is_cancel_requested(self, job_id: str) -> bool:
with self.connection() as connection:
row = connection.execute(
"SELECT cancel_requested FROM jobs WHERE id = ?", (job_id,)
).fetchone()
return bool(row and row["cancel_requested"])
def owns_lease(self, job_id: str, worker_id: str) -> bool:
now = timestamp()
with self.connection() as connection:
row = connection.execute(
"""
SELECT state, worker_id, lease_expires_at
FROM jobs
WHERE id = ?
""",
(job_id,),
).fetchone()
return bool(
row
and row["state"] == "running"
and row["worker_id"] == worker_id
and row["lease_expires_at"]
and row["lease_expires_at"] >= now
)
def request_cancel(self, job_id: str) -> dict[str, Any] | None:
now = timestamp()
with self.connection() as connection:
try:
connection.execute("BEGIN IMMEDIATE")
row = connection.execute(
"SELECT state FROM jobs WHERE id = ?", (job_id,)
).fetchone()
if not row:
connection.execute("COMMIT")
return None
if row["state"] in {"ready", "failed", "cancelled", "expired"}:
connection.execute("COMMIT")
return self.get_job(job_id)
if row["state"] == "queued":
cursor = connection.execute(
"""
UPDATE jobs
SET cancel_requested = 1, state = 'cancelled', phase = 'Cancelled',
message = 'Cancelled before starting', completed_at = ?, updated_at = ?
WHERE id = ? AND state = 'queued'
""",
(now, now, job_id),
)
payload = {
"state": "cancelled",
"phase": "Cancelled",
"message": "Cancelled before starting",
}
else:
cursor = connection.execute(
"""
UPDATE jobs
SET cancel_requested = 1, message = 'Cancellation requested', updated_at = ?
WHERE id = ? AND state = 'running' AND cancel_requested = 0
""",
(now, job_id),
)
payload = {
"state": "running",
"phase": "Cancelling",
"message": "Cancellation requested",
}
if cursor.rowcount:
self._append_event(connection, job_id, "state", payload)
connection.execute("COMMIT")
except Exception:
connection.execute("ROLLBACK")
raise
return self.get_job(job_id)
def mark_ready(
self,
job_id: str,
worker_id: str,
*,
artifact_path: str,
report_path: str,
archive_size: int,
entry_point: str,
stats: dict[str, int],
warnings: list[str],
retention_days: int,
) -> bool:
now_datetime = utcnow()
now = timestamp(now_datetime)
retention_days = min(retention_days, MAX_ARCHIVE_RETENTION_DAYS)
with self.connection() as connection:
cursor = connection.execute(
"""
UPDATE jobs
SET state = 'ready', phase = 'Ready', message = 'Archive is ready to download',
artifact_path = ?, report_path = ?, archive_size = ?, entry_point = ?,
pages_found = ?, pages_fetched = ?, assets_found = ?, assets_fetched = ?,
bytes_downloaded = ?, files_written = ?, warnings_json = ?,
worker_id = NULL, lease_expires_at = NULL, expires_at = ?,
completed_at = ?, updated_at = ?
WHERE id = ?
AND state = 'running'
AND worker_id = ?
AND cancel_requested = 0
AND lease_expires_at >= ?
""",
(
artifact_path,
report_path,
archive_size,
entry_point,
stats.get("pages_found", 0),
stats.get("pages_fetched", 0),
stats.get("assets_found", 0),
stats.get("assets_fetched", 0),
stats.get("bytes_downloaded", 0),
stats.get("files_written", 0),
json.dumps(warnings),
timestamp(now_datetime + timedelta(days=retention_days)),
now,
now,
job_id,
worker_id,
now,
),
)
if cursor.rowcount:
self._append_event(
connection,
job_id,
"state",
{"state": "ready", "phase": "Ready", "message": "Archive is ready to download"},
)
return True
return False
def mark_cancelled(
self, job_id: str, worker_id: str, message: str = "Capture cancelled"
) -> bool:
now = timestamp()
with self.connection() as connection:
cursor = connection.execute(
"""
UPDATE jobs
SET state = 'cancelled', phase = 'Cancelled', message = ?, worker_id = NULL,
lease_expires_at = NULL, completed_at = ?, updated_at = ?
WHERE id = ?
AND state = 'running'
AND worker_id = ?
AND cancel_requested = 1
AND lease_expires_at >= ?
""",
(message, now, now, job_id, worker_id, now),
)
if cursor.rowcount:
self._append_event(
connection,
job_id,
"state",
{"state": "cancelled", "phase": "Cancelled", "message": message},
)
return True
return False
def mark_failed(self, job_id: str, worker_id: str, code: str, message: str) -> bool:
now = timestamp()
with self.connection() as connection:
cursor = connection.execute(
"""
UPDATE jobs
SET state = 'failed', phase = 'Failed', message = ?, error_code = ?,
error_message = ?, worker_id = NULL, lease_expires_at = NULL,
completed_at = ?, updated_at = ?
WHERE id = ?
AND state = 'running'
AND worker_id = ?
AND cancel_requested = 0
AND lease_expires_at >= ?
""",
(message, code, message, now, now, job_id, worker_id, now),
)
if cursor.rowcount:
self._append_event(
connection,
job_id,
"state",
{"state": "failed", "phase": "Failed", "message": message, "code": code},
)
return True
return False
def mark_expired(self, job_id: str) -> None:
now = timestamp()
with self.connection() as connection:
connection.execute(
"""
UPDATE jobs
SET state = 'expired', phase = 'Expired',
message = 'Artifact retention period ended',
artifact_path = NULL, report_path = NULL, updated_at = ?
WHERE id = ?
""",
(now, job_id),
)
self._append_event(
connection,
job_id,
"state",
{
"state": "expired",
"phase": "Expired",
"message": "Artifact retention period ended",
},
)
def mark_deleted(self, job_id: str) -> None:
now = timestamp()
with self.connection() as connection:
cursor = connection.execute(
"""
UPDATE jobs
SET state = 'expired', phase = 'Deleted', message = 'Artifact deleted by user',
artifact_path = NULL, report_path = NULL, updated_at = ?
WHERE id = ? AND state = 'ready'
""",
(now, job_id),
)
if cursor.rowcount:
self._append_event(
connection,
job_id,
"state",
{
"state": "expired",
"phase": "Deleted",
"message": "Artifact deleted by user",
},
)
def expired_ready_jobs(self) -> list[dict[str, Any]]:
now_datetime = utcnow()
now = timestamp(now_datetime)
maximum_completion_age = timestamp(
now_datetime - timedelta(days=MAX_ARCHIVE_RETENTION_DAYS)
)
with self.connection() as connection:
rows = connection.execute(
"""
SELECT * FROM jobs
WHERE state = 'ready'
AND (
expires_at <= ?
OR COALESCE(completed_at, created_at) <= ?
)
""",
(now, maximum_completion_age),
).fetchall()
return [self._row_to_job(row) for row in rows]
def referenced_data_paths(self) -> set[str]:
with self.connection() as connection:
rows = connection.execute(
"""
SELECT artifact_path, report_path
FROM jobs
WHERE state = 'ready'
"""
).fetchall()
return {
path
for row in rows
for path in (row["artifact_path"], row["report_path"])
if isinstance(path, str) and path
}
def _append_event(
self, connection: sqlite3.Connection, job_id: str, kind: str, payload: dict[str, Any]
) -> None:
connection.execute(
"INSERT INTO job_events (job_id, kind, payload_json, created_at) VALUES (?, ?, ?, ?)",
(job_id, kind, json.dumps(payload, separators=(",", ":")), timestamp()),
)
@staticmethod
def _row_to_job(row: sqlite3.Row) -> dict[str, Any]:
data = dict(row)
data["config"] = json.loads(data.pop("config_json"))
data["warnings"] = json.loads(data.pop("warnings_json"))
data["cancel_requested"] = bool(data["cancel_requested"])
return data