"""Base collector class for all data sources""" import asyncio from abc import ABC, abstractmethod from typing import Dict, List, Any, Optional from datetime import UTC, datetime from time import perf_counter from urllib.parse import urlparse import httpx from sqlalchemy import select, text from sqlalchemy.ext.asyncio import AsyncSession from app.core.collected_data_fields import build_dynamic_metadata, get_record_field from app.core.countries import normalize_country from app.core.enums import JobStatus, SnapshotStatus from app.core.logging import get_logger from app.core.time import to_iso8601_utc from app.core.websocket.broadcaster import broadcaster from app.services.business_logs import emit_business_log, exception_context logger = get_logger(__name__, service="collector") class BaseCollector(ABC): """Abstract base class for data collectors""" name: str = "base_collector" priority: str = "P1" module: str = "L1" frequency_hours: int = 4 data_type: str = "generic" fail_on_empty: bool = False def __init__(self): self._current_task = None self._db_session = None self._datasource_id = 1 self._resolved_url: Optional[str] = None self._last_broadcast_progress: Optional[int] = None self._last_save_summary: dict[str, int] = {} async def resolve_url(self, db: AsyncSession) -> None: from app.core.data_sources import get_data_sources_config config = get_data_sources_config() self._resolved_url = await config.get_url(self.name, db) async def _publish_task_update(self, force: bool = False): if not self._current_task: return progress = float(self._current_task.progress or 0.0) rounded_progress = int(round(progress)) if not force and self._last_broadcast_progress == rounded_progress: return await broadcaster.broadcast_datasource_task_update( { "datasource_id": getattr(self, "_datasource_id", None), "collector_name": self.name, "task_id": self._current_task.id, "status": self._current_task.status, "phase": self._current_task.phase, "phase_progress": self._current_task.phase_progress, "phase_message": self._current_task.phase_message, "phase_current": self._current_task.phase_current, "phase_total": self._current_task.phase_total, "phase_unit": self._current_task.phase_unit, "progress": progress, "records_processed": self._current_task.records_processed, "total_records": self._current_task.total_records, "started_at": to_iso8601_utc(self._current_task.started_at), "completed_at": to_iso8601_utc(self._current_task.completed_at), "error_message": self._current_task.error_message, } ) self._last_broadcast_progress = rounded_progress async def update_progress(self, records_processed: int, *, commit: bool = False, force: bool = False): """Update task progress - call this during data processing""" if self._current_task and self._db_session: self._current_task.records_processed = records_processed if self._current_task.total_records and self._current_task.total_records > 0: self._current_task.progress = ( records_processed / self._current_task.total_records ) * 100 else: self._current_task.progress = 0.0 if commit: await self._db_session.commit() await self._publish_task_update(force=force) async def set_phase(self, phase: str, *, message: str | None = None, reset_progress: bool = True): if self._current_task and self._db_session: self._current_task.phase = phase self._current_task.phase_message = message if reset_progress: self._current_task.phase_progress = None self._current_task.phase_current = None self._current_task.phase_total = None self._current_task.phase_unit = None await self._db_session.commit() await self._publish_task_update(force=True) async def update_phase_progress( self, *, current: int | None = None, total: int | None = None, unit: str | None = None, message: str | None = None, progress: float | None = None, commit: bool = False, force: bool = False, ): """Update progress for the current phase without changing task totals.""" if not self._current_task or not self._db_session: return if progress is None and current is not None and total and total > 0: progress = (current / total) * 100 if progress is not None: self._current_task.phase_progress = max(0.0, min(float(progress), 100.0)) if current is not None: self._current_task.phase_current = max(0, int(current)) if total is not None: self._current_task.phase_total = max(0, int(total)) if unit is not None: self._current_task.phase_unit = unit if message is not None: self._current_task.phase_message = message if commit: await self._db_session.commit() await self._publish_task_update(force=force) @abstractmethod async def fetch(self) -> List[Dict[str, Any]]: """Fetch raw data from source""" pass def transform(self, raw_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Transform raw data to internal format (default: pass through)""" return raw_data def _parse_reference_date(self, value: Any) -> Optional[datetime]: if not value: return None if isinstance(value, datetime): return value if isinstance(value, str): return datetime.fromisoformat(value.replace("Z", "+00:00")) return None def _build_comparable_payload(self, record: Any) -> Dict[str, Any]: return { "name": getattr(record, "name", None), "title": getattr(record, "title", None), "description": getattr(record, "description", None), "country": get_record_field(record, "country"), "city": get_record_field(record, "city"), "latitude": get_record_field(record, "latitude"), "longitude": get_record_field(record, "longitude"), "value": get_record_field(record, "value"), "unit": get_record_field(record, "unit"), "metadata": getattr(record, "extra_data", None) or {}, "reference_date": ( getattr(record, "reference_date", None).isoformat() if getattr(record, "reference_date", None) else None ), } async def _create_snapshot( self, db: AsyncSession, task_id: int, data: List[Dict[str, Any]], started_at: datetime, ) -> int: from app.models.data_snapshot import DataSnapshot reference_dates = [ parsed for parsed in (self._parse_reference_date(item.get("reference_date")) for item in data) if parsed is not None ] reference_date = max(reference_dates) if reference_dates else None result = await db.execute( select(DataSnapshot) .where(DataSnapshot.source == self.name, DataSnapshot.is_current.is_(True)) .order_by(DataSnapshot.completed_at.desc().nullslast(), DataSnapshot.id.desc()) .limit(1) ) previous_snapshot = result.scalar_one_or_none() snapshot = DataSnapshot( datasource_id=getattr(self, "_datasource_id", 1), task_id=task_id, source=self.name, snapshot_key=f"{self.name}:{task_id}", reference_date=reference_date, started_at=started_at, status="running", is_current=True, parent_snapshot_id=previous_snapshot.id if previous_snapshot else None, summary={}, ) db.add(snapshot) if previous_snapshot: previous_snapshot.is_current = False await db.commit() return snapshot.id async def _rollback_incomplete_run( self, db: AsyncSession, *, task_id: int, snapshot_id: Optional[int], reason: str, ) -> None: from app.models.collected_data import CollectedData from app.models.data_snapshot import DataSnapshot await db.execute(CollectedData.__table__.delete().where(CollectedData.task_id == task_id)) parent_snapshot_id: Optional[int] = None if snapshot_id is not None: snapshot = await db.get(DataSnapshot, snapshot_id) if snapshot: parent_snapshot_id = snapshot.parent_snapshot_id snapshot.status = SnapshotStatus.CANCELLED.value snapshot.is_current = False snapshot.completed_at = datetime.now(UTC) summary = dict(snapshot.summary or {}) summary["rollback"] = True summary["rollback_reason"] = reason snapshot.summary = summary await db.execute( text( """ UPDATE collected_data SET is_current = FALSE WHERE source = :source """ ), {"source": self.name}, ) if parent_snapshot_id is not None: parent_snapshot = await db.get(DataSnapshot, parent_snapshot_id) if parent_snapshot: parent_snapshot.is_current = True await db.execute( text( """ UPDATE collected_data SET is_current = TRUE WHERE snapshot_id = :snapshot_id """ ), {"snapshot_id": parent_snapshot_id}, ) async def run(self, db: AsyncSession) -> Dict[str, Any]: """Full pipeline: fetch -> transform -> save""" from app.services.collectors.registry import collector_registry from app.models.task import CollectionTask from app.models.data_snapshot import DataSnapshot start_time = datetime.now(UTC) started_at = perf_counter() datasource_id = getattr(self, "_datasource_id", 1) snapshot_id: Optional[int] = None if not collector_registry.is_active(self.name): await self._log_collection_event( "collector.run.skipped_disabled", "Collector skipped because it is disabled", level="info", context={"status": "skipped", "reason": "disabled"}, ) return {"status": "skipped", "reason": "Collector is disabled"} task = self._current_task if isinstance(self._current_task, CollectionTask) else None if task is None: task = CollectionTask( datasource_id=datasource_id, source=self.name, task_type="collect", status="running", phase="queued", started_at=start_time, ) db.add(task) else: task.datasource_id = datasource_id task.source = task.source or self.name task.task_type = task.task_type or "collect" task.status = JobStatus.RUNNING.value task.phase = "queued" task.started_at = task.started_at or start_time task.completed_at = None task.error_message = None await db.commit() task_id = task.id self._current_task = task self._db_session = db self._last_broadcast_progress = None await self.resolve_url(db) await self._publish_task_update(force=True) await self._log_collection_event( "collector.run.started", "Collector run started", context={"status": "running", "task_id": task_id}, ) try: phase_started_at = perf_counter() await self.set_phase("fetching", message="正在拉取原始数据") await self._log_collection_event( "collector.phase.fetching.start", "Collector fetch phase started", context={"task_id": task_id, "snapshot_id": snapshot_id}, ) raw_data = await self.fetch() task.total_records = len(raw_data) await db.commit() await self._publish_task_update(force=True) await self._log_collection_event( "collector.phase.fetching.success", "Collector fetch phase completed", context={ "task_id": task_id, "raw_count": len(raw_data), "duration_ms": self._duration_ms(phase_started_at), }, ) if self.fail_on_empty and not raw_data: raise RuntimeError(f"Collector {self.name} returned no data") phase_started_at = perf_counter() await self.set_phase("transforming", message="正在转换采集数据") await self._log_collection_event( "collector.phase.transforming.start", "Collector transform phase started", context={"task_id": task_id, "raw_count": len(raw_data)}, ) data = self.transform(raw_data) await self._log_collection_event( "collector.phase.transforming.success", "Collector transform phase completed", context={ "task_id": task_id, "raw_count": len(raw_data), "transformed_count": len(data), "duration_ms": self._duration_ms(phase_started_at), }, ) snapshot_id = await self._create_snapshot(db, task_id, data, start_time) phase_started_at = perf_counter() await self.set_phase("saving", message="正在保存采集数据") await self._log_collection_event( "collector.phase.saving.start", "Collector save phase started", context={"task_id": task_id, "snapshot_id": snapshot_id, "transformed_count": len(data)}, ) records_count = await self._save_data(db, data, task_id=task_id, snapshot_id=snapshot_id) await self._log_collection_event( "collector.phase.saving.success", "Collector save phase completed", context={ "task_id": task_id, "snapshot_id": snapshot_id, "saved_count": records_count, **self._last_save_summary, "duration_ms": self._duration_ms(phase_started_at), }, ) task.status = JobStatus.SUCCESS.value task.phase = "completed" task.phase_progress = 100.0 task.phase_message = "采集完成" task.phase_current = records_count task.phase_total = records_count task.phase_unit = "records" task.records_processed = records_count task.progress = 100.0 task.completed_at = datetime.now(UTC) await db.commit() await self._publish_task_update(force=True) await self._log_collection_event( "collector.run.completed", "Collector run completed", context={ "status": "success", "task_id": task_id, "snapshot_id": snapshot_id, "raw_count": len(raw_data), "transformed_count": len(data), "saved_count": records_count, **self._last_save_summary, "duration_ms": self._duration_ms(started_at), }, ) return { "status": "success", "task_id": task_id, "records_processed": records_count, "execution_time_seconds": (datetime.now(UTC) - start_time).total_seconds(), } except asyncio.CancelledError: await db.rollback() task.status = JobStatus.CANCELLED.value task.phase = "cancelled" task.phase_message = "采集已取消" task.error_message = "Collection cancelled by operator and rolled back" task.completed_at = datetime.now(UTC) if snapshot_id is not None: await self._rollback_incomplete_run( db, task_id=task_id, snapshot_id=snapshot_id, reason="cancelled_by_operator", ) await db.commit() await self._publish_task_update(force=True) await self._log_collection_event( "collector.run.cancelled", "Collector run cancelled", level="warning", context={ "status": "cancelled", "task_id": task_id, "snapshot_id": snapshot_id, "duration_ms": self._duration_ms(started_at), }, ) raise except Exception as e: await db.rollback() task.status = JobStatus.FAILED.value task.phase = "failed" task.phase_message = str(e) task.error_message = str(e) task.completed_at = datetime.now(UTC) if snapshot_id is not None: snapshot = await db.get(DataSnapshot, snapshot_id) if snapshot: snapshot.status = SnapshotStatus.FAILED.value snapshot.completed_at = datetime.now(UTC) snapshot.summary = {"error": str(e)} await db.commit() await self._publish_task_update(force=True) await self._log_collection_event( "collector.run.failed", "Collector run failed", level="error", context=exception_context( e, { "status": "failed", "task_id": task_id, "snapshot_id": snapshot_id, "duration_ms": self._duration_ms(started_at), }, ), ) return { "status": "failed", "task_id": task_id, "error": str(e), "execution_time_seconds": (datetime.now(UTC) - start_time).total_seconds(), } async def _save_data( self, db: AsyncSession, data: List[Dict[str, Any]], task_id: Optional[int] = None, snapshot_id: Optional[int] = None, ) -> int: """Save transformed data to database""" from app.models.collected_data import CollectedData from app.models.data_snapshot import DataSnapshot if not data: self._last_save_summary = {"created": 0, "updated": 0, "unchanged": 0, "deleted": 0} if snapshot_id is not None: snapshot = await db.get(DataSnapshot, snapshot_id) if snapshot: snapshot.record_count = 0 snapshot.summary = {"created": 0, "updated": 0, "unchanged": 0} snapshot.status = SnapshotStatus.SUCCESS.value snapshot.completed_at = datetime.now(UTC) await db.commit() return 0 collected_at = datetime.now(UTC) records_added = 0 created_count = 0 updated_count = 0 unchanged_count = 0 seen_entity_keys: set[str] = set() progress_commit_interval = 1000 previous_current_result = await db.execute( select(CollectedData) .where( CollectedData.source == self.name, CollectedData.is_current.is_(True), ) .order_by(CollectedData.entity_key.asc(), CollectedData.collected_at.desc().nullslast(), CollectedData.id.desc()) ) previous_current_records = previous_current_result.scalars().all() previous_current_keys = {record.entity_key for record in previous_current_records if record.entity_key} previous_current_map: dict[str, CollectedData] = {} stale_previous_records: list[CollectedData] = [] for existing_record in previous_current_records: entity_key = existing_record.entity_key if not entity_key: continue if entity_key not in previous_current_map: previous_current_map[entity_key] = existing_record continue stale_previous_records.append(existing_record) for stale_record in stale_previous_records: stale_record.is_current = False for i, item in enumerate(data): raw_metadata = item.get("metadata", {}) extra_data = build_dynamic_metadata( raw_metadata, country=item.get("country"), city=item.get("city"), latitude=item.get("latitude"), longitude=item.get("longitude"), value=item.get("value"), unit=item.get("unit"), ) normalized_country = normalize_country(item.get("country")) if normalized_country is not None: extra_data["country"] = normalized_country if item.get("country") and normalized_country != item.get("country"): extra_data["raw_country"] = item.get("country") if normalized_country is None: extra_data["country_validation"] = "invalid" source_id = item.get("source_id") or item.get("id") reference_date = ( self._parse_reference_date(item.get("reference_date")) ) source_id_str = str(source_id) if source_id is not None else None entity_key = f"{self.name}:{source_id_str}" if source_id_str else f"{self.name}:{i}" previous_record = None if entity_key and entity_key not in seen_entity_keys: previous_record = previous_current_map.get(entity_key) if previous_record is not None: previous_record.is_current = False record = CollectedData( snapshot_id=snapshot_id, task_id=task_id, source=self.name, source_id=source_id_str, entity_key=entity_key, data_type=self.data_type, name=item.get("name"), title=item.get("title"), description=item.get("description"), extra_data=extra_data, collected_at=collected_at, reference_date=reference_date, is_valid=1, is_current=True, previous_record_id=previous_record.id if previous_record else None, deleted_at=None, ) if previous_record is None: record.change_type = "created" record.change_summary = {} created_count += 1 else: previous_payload = self._build_comparable_payload(previous_record) current_payload = self._build_comparable_payload(record) if current_payload == previous_payload: record.change_type = "unchanged" record.change_summary = {} unchanged_count += 1 else: changed_fields = [ key for key in current_payload.keys() if current_payload[key] != previous_payload.get(key) ] record.change_type = "updated" record.change_summary = {"changed_fields": changed_fields} updated_count += 1 db.add(record) seen_entity_keys.add(entity_key) records_added += 1 if (i + 1) % progress_commit_interval == 0: await self.update_progress(i + 1, commit=True) if snapshot_id is not None: deleted_keys = previous_current_keys - seen_entity_keys await db.execute( text( """ UPDATE collected_data SET is_current = FALSE WHERE source = :source AND snapshot_id IS DISTINCT FROM :snapshot_id AND COALESCE(is_current, TRUE) = TRUE """ ), {"source": self.name, "snapshot_id": snapshot_id}, ) snapshot = await db.get(DataSnapshot, snapshot_id) if snapshot: snapshot.record_count = records_added snapshot.status = SnapshotStatus.SUCCESS.value snapshot.completed_at = datetime.now(UTC) snapshot.summary = { "created": created_count, "updated": updated_count, "unchanged": unchanged_count, "deleted": len(deleted_keys), } self._last_save_summary = { "created": created_count, "updated": updated_count, "unchanged": unchanged_count, "deleted": len(deleted_keys), } else: self._last_save_summary = { "created": created_count, "updated": updated_count, "unchanged": unchanged_count, "deleted": 0, } await db.commit() await self.update_progress(len(data), force=True) return records_added @staticmethod def _duration_ms(started_at: float) -> int: return int((perf_counter() - started_at) * 1000) async def _log_collection_event( self, event: str, message: str, *, level: str = "info", context: Dict[str, Any] | None = None, ) -> None: await emit_business_log( logger, event=event, message=message, category="collector", level=level, service="collector", module=__name__, context={ "collector_name": self.name, "datasource_id": getattr(self, "_datasource_id", None), **(context or {}), }, ) async def save(self, db: AsyncSession, data: List[Dict[str, Any]]) -> int: """Save data to database (legacy method, use _save_data instead)""" return await self._save_data(db, data) class HTTPCollector(BaseCollector): """Base class for HTTP API collectors""" base_url: str = "" headers: Dict[str, str] = {} async def fetch(self) -> List[Dict[str, Any]]: started_at = perf_counter() request_host = urlparse(self.base_url).netloc await emit_business_log( logger, event="collector.http.fetch.start", message="Collector HTTP request started", category="collector", service="collector", module=__name__, context={ "collector_name": self.name, "datasource_id": getattr(self, "_datasource_id", None), "url_host": request_host, }, ) async with httpx.AsyncClient(timeout=60.0) as client: try: response = await client.get(self.base_url, headers=self.headers) response.raise_for_status() payload = response.json() parsed = self.parse_response(payload) await emit_business_log( logger, event="collector.http.fetch.success", message="Collector HTTP request completed", category="collector", service="collector", module=__name__, context={ "collector_name": self.name, "datasource_id": getattr(self, "_datasource_id", None), "url_host": request_host, "status_code": response.status_code, "response_bytes": len(response.content or b""), "parsed_count": len(parsed), "duration_ms": BaseCollector._duration_ms(started_at), }, ) return parsed except Exception as exc: await emit_business_log( logger, event="collector.http.fetch.failed", message="Collector HTTP request failed", category="collector", level="error", service="collector", module=__name__, context=exception_context( exc, { "collector_name": self.name, "datasource_id": getattr(self, "_datasource_id", None), "url_host": request_host, "duration_ms": BaseCollector._duration_ms(started_at), }, ), ) raise @abstractmethod def parse_response(self, response: Dict[str, Any]) -> List[Dict[str, Any]]: pass class IntervalCollector(BaseCollector): """Base class for collectors that run on intervals""" async def run(self, db: AsyncSession) -> Dict[str, Any]: return await super().run(db) async def log_task( db: AsyncSession, datasource_id: int, status: str, records_processed: int = 0, error_message: Optional[str] = None, ): """Log collection task to database""" from app.models.task import CollectionTask task = CollectionTask( datasource_id=datasource_id, status=status, records_processed=records_processed, error_message=error_message, started_at=datetime.now(UTC), completed_at=datetime.now(UTC), ) db.add(task) await db.commit()