Files
planet/backend/app/services/collectors/base.py
linkong 8c204717cd
Some checks failed
ci / backend (push) Has been cancelled
ci / frontend (push) Has been cancelled
release / images (push) Has been cancelled
ci / delivery (push) Has been cancelled
release: bump version to 0.70.0
2026-06-04 17:16:23 +08:00

802 lines
30 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
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()