what2do/scripts/cache_source_documents.py
2026-09-07 12:53:26 +10:00

75 lines
3.7 KiB
Python

#!/usr/bin/env python3
import argparse
import hashlib
import json
import time
from datetime import datetime, timezone
from pathlib import Path
from app.crawler import _http_client
ROOT = Path(__file__).resolve().parents[1]
MANIFEST_PATH = ROOT / "tests" / "fixtures" / "illawarra_sources.json"
CACHE_DIR = ROOT / "tests" / "fixtures" / "source_cache"
INDEX_PATH = CACHE_DIR / "index.json"
def cache_key(source):
digest = hashlib.sha256(source["url"].encode()).hexdigest()[:12]
return f"{digest}.html"
def main():
parser = argparse.ArgumentParser(description="Manually cache each Illawarra source once for offline extraction tests.")
parser.add_argument("--delay", type=float, default=2.0, help="Seconds to wait between source requests (default: 2).")
parser.add_argument("--timeout", type=float, default=30.0, help="Per-request timeout in seconds (default: 30).")
parser.add_argument("--limit", type=int, help="Only fetch the first N sources, useful for a trial run.")
parser.add_argument("--force", action="store_true", help="Replace documents already present in the cache.")
args = parser.parse_args()
sources = json.loads(MANIFEST_PATH.read_text())
if args.limit is not None:
sources = sources[:args.limit]
CACHE_DIR.mkdir(parents=True, exist_ok=True)
previous = json.loads(INDEX_PATH.read_text()) if INDEX_PATH.exists() else {"sources": {}}
records = previous.get("sources", {})
with _http_client(timeout=args.timeout) as client:
for position, source in enumerate(sources, start=1):
filename = cache_key(source)
destination = CACHE_DIR / filename
existing = records.get(source["url"])
if destination.exists() and existing and not args.force:
print(f"[{position}/{len(sources)}] SKIP {source['name']} (already cached)")
continue
print(f"[{position}/{len(sources)}] GET {source['name']} — {source['url']}")
record = {"name": source["name"], "url": source["url"], "mode": source["mode"], "filename": filename, "fetched_at": datetime.now(timezone.utc).isoformat()}
try:
response = client.get(source["url"])
record.update({"status_code": response.status_code, "content_type": response.headers.get("content-type", ""), "final_url": str(response.url), "bytes": len(response.content)})
if response.status_code < 400 and "html" in record["content_type"].lower():
destination.write_text(response.text, encoding="utf-8")
record["cached"] = True
print(f" SAVED {filename} ({record['bytes']} bytes)")
else:
record["cached"] = False
record["error"] = f"Expected successful HTML response, got {response.status_code} {record['content_type'] or 'unknown content type'}"
destination.unlink(missing_ok=True)
print(f" ERROR {record['error']}")
except Exception as exc:
record.update({"cached": False, "error": f"{type(exc).__name__}: {exc}"})
destination.unlink(missing_ok=True)
print(f" ERROR {record['error']}")
records[source["url"]] = record
INDEX_PATH.write_text(json.dumps({"generated_at": datetime.now(timezone.utc).isoformat(), "sources": records}, indent=2) + "\n")
if position < len(sources) and args.delay > 0:
time.sleep(args.delay)
cached = sum(bool(record.get("cached")) for record in records.values())
print(f"Cache complete: {cached}/{len(records)} documents available in {CACHE_DIR}")
if __name__ == "__main__":
main()