"""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 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.config import settings from app.core.countries import normalize_country from app.core.time import to_iso8601_utc from app.core.websocket.broadcaster import broadcaster 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 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, "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): if self._current_task and self._db_session: self._current_task.phase = phase await self._db_session.commit() await self._publish_task_update(force=True) @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 == 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 = "cancelled" 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) datasource_id = getattr(self, "_datasource_id", 1) snapshot_id: Optional[int] = None if not collector_registry.is_active(self.name): return {"status": "skipped", "reason": "Collector is disabled"} task = CollectionTask( datasource_id=datasource_id, status="running", phase="queued", started_at=start_time, ) db.add(task) 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) try: await self.set_phase("fetching") raw_data = await self.fetch() task.total_records = len(raw_data) await db.commit() await self._publish_task_update(force=True) if self.fail_on_empty and not raw_data: raise RuntimeError(f"Collector {self.name} returned no data") await self.set_phase("transforming") data = self.transform(raw_data) snapshot_id = await self._create_snapshot(db, task_id, data, start_time) await self.set_phase("saving") records_count = await self._save_data(db, data, task_id=task_id, snapshot_id=snapshot_id) task.status = "success" task.phase = "completed" 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) return { "status": "success", "task_id": task_id, "records_processed": records_count, "execution_time_seconds": (datetime.now(UTC) - start_time).total_seconds(), } except asyncio.CancelledError: task.status = "cancelled" task.phase = "cancelled" 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) raise except Exception as e: task.status = "failed" task.phase = "failed" 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 = "failed" snapshot.completed_at = datetime.now(UTC) snapshot.summary = {"error": str(e)} await db.commit() await self._publish_task_update(force=True) 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: 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 = "success" 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 == 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 = "success" snapshot.completed_at = datetime.now(UTC) snapshot.summary = { "created": created_count, "updated": updated_count, "unchanged": unchanged_count, "deleted": len(deleted_keys), } await db.commit() await self.update_progress(len(data), force=True) return records_added 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]]: async with httpx.AsyncClient(timeout=60.0) as client: response = await client.get(self.base_url, headers=self.headers) response.raise_for_status() return self.parse_response(response.json()) @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()