534 lines
20 KiB
Python
534 lines
20 KiB
Python
"""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:
|
|
await db.rollback()
|
|
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:
|
|
await db.rollback()
|
|
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()
|