"""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