n0mad1k
293 lines
11 KiB
Python
293 lines
11 KiB
Python
"""Abstract base collector for Mosaic document acquisition."""
|
|
import os
|
|
import tempfile
|
|
import hashlib
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
import datetime
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
|
|
|
|
from storage.db import MosaicDB
|
|
from storage.models import Document
|
|
from utils.http_limiter import HTTPRateLimiter
|
|
from utils.cache import DownloadCache
|
|
from utils.sanitize import sanitize_filename, validate_content_type
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BaseCollector(ABC):
|
|
"""Abstract base for all document collectors."""
|
|
|
|
def __init__(self, db: MosaicDB, cache: DownloadCache,
|
|
rate_limiter: Optional[HTTPRateLimiter] = None,
|
|
config: Optional[dict] = None):
|
|
self.db = db
|
|
self.cache = cache
|
|
self.rate_limiter = rate_limiter or HTTPRateLimiter()
|
|
self.config = config or {}
|
|
self.user_agent = self.config.get('user_agent', 'Mozilla/5.0 (compatible; research-crawler/1.0)')
|
|
self.timeout = int(self.config.get('timeout', 30))
|
|
self.max_retries = int(self.config.get('max_retries', 3))
|
|
|
|
@abstractmethod
|
|
def discover(self, since: Optional[str] = None,
|
|
until: Optional[str] = None,
|
|
limit: Optional[int] = None) -> list[dict]:
|
|
"""Discover document URLs from the source.
|
|
|
|
Returns list of dicts: {'url': str, 'title': str, 'date': str|None, 'doc_type': str}
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_name(self) -> str:
|
|
"""Return the source name."""
|
|
pass
|
|
|
|
def expected_count(self) -> int:
|
|
"""Minimum expected document count for validation."""
|
|
return 0
|
|
|
|
def fetch(self, url: str, collection: str) -> Optional[Document]:
|
|
"""Download a URL with atomic writes and hash verification.
|
|
|
|
Returns Document if successful, None if failed.
|
|
"""
|
|
# Check if already cached
|
|
existing = self.db.get_document_by_url(url)
|
|
if existing and existing.status == 'CACHED':
|
|
cached_path = self.cache.get_cached_path(
|
|
existing.content_hash, collection,
|
|
self._url_extension(url)
|
|
)
|
|
if cached_path:
|
|
logger.debug("Already cached: %s", url)
|
|
return existing
|
|
|
|
# Rate limit
|
|
self.rate_limiter.wait()
|
|
|
|
# Download to temp file
|
|
try:
|
|
with httpx.Client(timeout=self.timeout, follow_redirects=True) as client:
|
|
response = client.get(url, headers={'User-Agent': self.user_agent})
|
|
response.raise_for_status()
|
|
|
|
content = response.content
|
|
content_hash = hashlib.sha256(content).hexdigest()
|
|
content_length = len(content)
|
|
etag = response.headers.get('etag')
|
|
last_modified = response.headers.get('last-modified')
|
|
content_type = response.headers.get('content-type', '').split(';')[0].strip()
|
|
|
|
# Write to temp file
|
|
ext = self._url_extension(url) or self._mime_to_ext(content_type)
|
|
tmp_fd, tmp_path = tempfile.mkstemp(suffix=f'.{ext}.tmp', dir=str(self.cache.base_dir))
|
|
try:
|
|
os.write(tmp_fd, content)
|
|
os.close(tmp_fd)
|
|
|
|
# Validate content type by magic bytes
|
|
try:
|
|
detected_type = validate_content_type(tmp_path)
|
|
except ValueError:
|
|
detected_type = content_type
|
|
|
|
# Atomic store
|
|
final_path = self.cache.store_atomic(
|
|
tmp_path, content_hash, collection, ext, content_length
|
|
)
|
|
except Exception:
|
|
# Clean up temp file on failure
|
|
os.close(tmp_fd) if not os.path.exists(tmp_path) else None
|
|
if os.path.exists(tmp_path):
|
|
os.unlink(tmp_path)
|
|
raise
|
|
|
|
# Build document
|
|
doc_type = self._detect_doc_type(detected_type)
|
|
doc = Document(
|
|
id=content_hash,
|
|
source=collection,
|
|
source_url=url,
|
|
doc_type=doc_type,
|
|
title=self._extract_title(url),
|
|
raw_path=final_path,
|
|
content_hash=content_hash,
|
|
fetch_date=datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
etag=etag,
|
|
last_modified=last_modified,
|
|
char_count=len(content.decode('utf-8', errors='replace')),
|
|
status='CACHED',
|
|
parse_status='PENDING',
|
|
)
|
|
|
|
self.db.insert_document(doc)
|
|
return doc
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error("HTTP %d for %s: %s", e.response.status_code, url, e)
|
|
except httpx.RequestError as e:
|
|
logger.error("Request error for %s: %s", url, e)
|
|
except Exception as e:
|
|
logger.error("Download failed for %s: %s", url, e)
|
|
|
|
# Record failure
|
|
if existing:
|
|
self.db.update_document_status(existing.id, status='DOWNLOAD_FAILED')
|
|
return None
|
|
|
|
def check_update(self, url: str) -> bool:
|
|
"""Check if a cached URL has been updated (HTTP HEAD with ETag/Last-Modified).
|
|
|
|
Returns True if the document has changed and should be re-downloaded.
|
|
"""
|
|
existing = self.db.get_document_by_url(url)
|
|
if not existing:
|
|
return True
|
|
|
|
try:
|
|
self.rate_limiter.wait()
|
|
with httpx.Client(timeout=self.timeout, follow_redirects=True) as client:
|
|
response = client.head(url, headers={'User-Agent': self.user_agent})
|
|
response.raise_for_status()
|
|
|
|
new_etag = response.headers.get('etag')
|
|
new_modified = response.headers.get('last-modified')
|
|
|
|
if existing.etag and new_etag:
|
|
return existing.etag != new_etag
|
|
if existing.last_modified and new_modified:
|
|
return existing.last_modified != new_modified
|
|
|
|
# Can't determine — assume changed
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.warning("HEAD check failed for %s: %s", url, e)
|
|
return False
|
|
|
|
def collect(self, since: Optional[str] = None, until: Optional[str] = None,
|
|
limit: Optional[int] = None, update: bool = False,
|
|
retry_failed: bool = False) -> dict:
|
|
"""Main collection entry point with progress tracking.
|
|
|
|
Returns dict with counts: {'discovered', 'downloaded', 'cached', 'failed', 'updated'}
|
|
"""
|
|
name = self.get_name()
|
|
stats = {'discovered': 0, 'downloaded': 0, 'cached': 0, 'failed': 0, 'updated': 0}
|
|
|
|
# Check disk space
|
|
if not self.cache.check_disk_space(required_mb=100):
|
|
logger.error("Insufficient disk space for collection")
|
|
return stats
|
|
|
|
# Discover URLs
|
|
urls = self.discover(since=since, until=until, limit=limit)
|
|
stats['discovered'] = len(urls)
|
|
|
|
# Validate against expected count
|
|
expected = self.expected_count()
|
|
if expected > 0 and len(urls) < expected * 0.1:
|
|
logger.warning(
|
|
"Discovered only %d URLs (expected >%d). "
|
|
"Site structure may have changed.",
|
|
len(urls), int(expected * 0.1)
|
|
)
|
|
|
|
# Filter for update mode
|
|
if update:
|
|
urls_to_process = []
|
|
for item in urls:
|
|
if self.check_update(item['url']):
|
|
urls_to_process.append(item)
|
|
stats['updated'] += 1
|
|
urls = urls_to_process
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
transient=True,
|
|
) as progress:
|
|
task = progress.add_task(f"Collecting {name}", total=len(urls))
|
|
|
|
for item in urls:
|
|
url = item['url']
|
|
|
|
# Skip already cached (unless update mode)
|
|
existing = self.db.get_document_by_url(url)
|
|
if existing and existing.status == 'CACHED' and not update:
|
|
stats['cached'] += 1
|
|
progress.advance(task)
|
|
continue
|
|
|
|
# Skip failed unless retry
|
|
if existing and existing.status == 'DOWNLOAD_FAILED' and not retry_failed:
|
|
stats['failed'] += 1
|
|
progress.advance(task)
|
|
continue
|
|
|
|
doc = self.fetch(url, name)
|
|
if doc:
|
|
# Update with discovered metadata
|
|
if item.get('title'):
|
|
doc.title = item['title']
|
|
if item.get('date'):
|
|
doc.date = item['date']
|
|
self.db.insert_document(doc)
|
|
stats['downloaded'] += 1
|
|
else:
|
|
stats['failed'] += 1
|
|
|
|
progress.advance(task)
|
|
|
|
return stats
|
|
|
|
def _url_extension(self, url: str) -> str:
|
|
"""Extract file extension from URL."""
|
|
from urllib.parse import urlparse
|
|
path = urlparse(url).path
|
|
ext = Path(path).suffix.lstrip('.')
|
|
if ext and len(ext) <= 10:
|
|
return sanitize_filename(ext).lower()
|
|
return 'html'
|
|
|
|
def _mime_to_ext(self, mime: str) -> str:
|
|
"""Convert MIME type to file extension."""
|
|
mapping = {
|
|
'text/html': 'html',
|
|
'application/pdf': 'pdf',
|
|
'text/plain': 'txt',
|
|
'message/rfc822': 'eml',
|
|
'application/json': 'json',
|
|
}
|
|
return mapping.get(mime, 'bin')
|
|
|
|
def _detect_doc_type(self, content_type: str) -> str:
|
|
"""Map content type to our doc_type enum."""
|
|
if 'html' in content_type:
|
|
return 'html'
|
|
elif 'pdf' in content_type:
|
|
return 'pdf'
|
|
elif 'email' in content_type or 'rfc822' in content_type:
|
|
return 'email'
|
|
elif 'text' in content_type:
|
|
return 'text'
|
|
return 'html' # default for web pages
|
|
|
|
def _extract_title(self, url: str) -> str:
|
|
"""Extract a title from URL path."""
|
|
from urllib.parse import urlparse, unquote
|
|
path = urlparse(url).path
|
|
parts = [p for p in path.split('/') if p]
|
|
if parts:
|
|
return unquote(parts[-1]).replace('-', ' ').replace('_', ' ')
|
|
return url
|