release: bump version to 0.66.0

backend/app/services/data_jobs.py

@@ -0,0 +1,597 @@
+"""Kafka-ready datasource job queue backed by PostgreSQL for v1."""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import UTC, datetime, timedelta
+from uuid import uuid4
+from typing import Any
+
+from sqlalchemy import bindparam, select, text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.cache import cache
+from app.core.config import settings
+from app.core.logging import get_logger
+from app.core.time import to_iso8601_utc
+from app.core.websocket.broadcaster import broadcaster
+from app.db.session import async_session_factory
+from app.models.collected_data import CollectedData
+from app.models.data_snapshot import DataSnapshot
+from app.models.datasource import DataSource
+from app.models.task import CollectionTask
+from app.services.collectors.registry import collector_registry
+from app.services.datasource_connectivity import (
+    build_builtin_connectivity_checksum,
+    get_builtin_effective_candidate,
+    save_connectivity_success,
+)
+from app.services.earth_layer_adapters import (
+    clear_derived_datasource_data,
+    get_earth_refresh_strategy_for_change,
+    get_earth_update_layers_for_source,
+)
+from app.services.earth_layer_cache import invalidate_earth_layer_cache_for_source
+from app.services.scheduler import sync_datasource_job
+
+logger = get_logger(__name__)
+
+JOB_TYPE_COLLECT = "collect"
+JOB_TYPE_CLEAR_DATA = "clear_data"
+JOB_TYPE_CLEAR_CACHE = "clear_cache"
+JOB_TYPE_EARTH_REFRESH = "earth_refresh"
+
+JOB_STATUS_QUEUED = "queued"
+JOB_STATUS_RUNNING = "running"
+JOB_STATUS_CANCELLING = "cancelling"
+JOB_STATUS_SUCCESS = "success"
+JOB_STATUS_FAILED = "failed"
+JOB_STATUS_CANCELLED = "cancelled"
+
+ACTIVE_JOB_STATUSES = (JOB_STATUS_QUEUED, JOB_STATUS_RUNNING, JOB_STATUS_CANCELLING)
+TERMINAL_JOB_STATUSES = (JOB_STATUS_SUCCESS, JOB_STATUS_FAILED, JOB_STATUS_CANCELLED)
+DATA_WRITE_JOB_TYPES = (JOB_TYPE_COLLECT, JOB_TYPE_CLEAR_DATA, JOB_TYPE_CLEAR_CACHE)
+SOURCE_LOCK_JOB_STATUSES = (JOB_STATUS_RUNNING, JOB_STATUS_CANCELLING)
+QUEUE_POLL_SECONDS = 0.35
+JOB_STALE_LOCK_MINUTES = 90
+DEFAULT_WORKER_CONCURRENCY = 2
+
+RUNNING_DATA_JOB_TASKS: dict[int, asyncio.Task[Any]] = {}
+
+
+def _utcnow() -> datetime:
+    return datetime.now(UTC)
+
+
+def _job_worker_id() -> str:
+    return f"{settings.PROJECT_NAME}:data-job-worker:{uuid4().hex[:8]}"
+
+
+def is_terminal_job_status(status: str | None) -> bool:
+    return status in TERMINAL_JOB_STATUSES
+
+
+async def enqueue_datasource_job(
+    db: AsyncSession,
+    datasource: DataSource,
+    task_type: str,
+    *,
+    payload: dict[str, Any] | None = None,
+    rollback_policy: str = "keep_committed_batches",
+    dedupe_key: str | None = None,
+) -> CollectionTask:
+    if dedupe_key:
+        existing = await _get_active_job_by_dedupe_key(db, dedupe_key)
+        if existing is not None:
+            return existing
+
+    task = CollectionTask(
+        datasource_id=datasource.id,
+        source=datasource.source,
+        task_type=task_type,
+        status=JOB_STATUS_QUEUED,
+        phase="queued",
+        phase_message="任务已进入队列",
+        payload=payload or {},
+        rollback_policy=rollback_policy,
+        dedupe_key=dedupe_key,
+    )
+    db.add(task)
+    await db.commit()
+    await db.refresh(task)
+    await _broadcast_task_update(task)
+    return task
+
+
+async def enqueue_earth_refresh_job(
+    db: AsyncSession,
+    *,
+    source: str,
+    payload: dict[str, Any] | None = None,
+) -> CollectionTask | None:
+    layers = list((payload or {}).get("layers") or get_earth_update_layers_for_source(source))
+    if not layers:
+        return None
+
+    datasource = await _get_or_create_virtual_datasource(db, source)
+    refresh_payload = {
+        "source": source,
+        "layers": layers,
+        "refresh_strategy": (payload or {}).get("refresh_strategy")
+        or get_earth_refresh_strategy_for_change((payload or {}).get("table"), source)
+        or "clear_then_reload",
+        **(payload or {}),
+    }
+    return await enqueue_datasource_job(
+        db,
+        datasource,
+        JOB_TYPE_EARTH_REFRESH,
+        payload=refresh_payload,
+        dedupe_key=f"earth_refresh:{source}",
+    )
+
+
+async def enqueue_earth_refresh_from_update(payload: dict[str, Any]) -> None:
+    source = str(payload.get("source") or "").strip()
+    if not source:
+        return
+    async with async_session_factory() as db:
+        await enqueue_earth_refresh_job(db, source=source, payload=payload)
+
+
+async def request_cancel_datasource_task(
+    db: AsyncSession,
+    task: CollectionTask,
+    *,
+    reason: str = "cancelled_by_operator",
+) -> CollectionTask:
+    if is_terminal_job_status(task.status):
+        return task
+
+    running_task = RUNNING_DATA_JOB_TASKS.get(task.id)
+    if task.status == JOB_STATUS_QUEUED or (
+        running_task is None
+        and (
+            task.status == JOB_STATUS_CANCELLING
+            or (task.status == JOB_STATUS_RUNNING and task.task_type != JOB_TYPE_COLLECT)
+        )
+    ):
+        return await _cancel_task_without_runner(db, task, reason=reason)
+
+    task.status = JOB_STATUS_CANCELLING
+    task.phase = JOB_STATUS_CANCELLING
+    task.phase_message = "正在停止任务"
+    task.requested_cancel_at = _utcnow()
+    task.cancel_reason = reason
+    await db.commit()
+    await db.refresh(task)
+
+    if running_task is not None and not running_task.done():
+        running_task.cancel()
+
+    await _broadcast_task_update(task)
+    return task
+
+
+async def _cancel_task_without_runner(
+    db: AsyncSession,
+    task: CollectionTask,
+    *,
+    reason: str,
+) -> CollectionTask:
+    if task.task_type == JOB_TYPE_COLLECT:
+        await db.execute(CollectedData.__table__.delete().where(CollectedData.task_id == task.id))
+        snapshot_result = await db.execute(select(DataSnapshot).where(DataSnapshot.task_id == task.id))
+        for snapshot in snapshot_result.scalars().all():
+            snapshot.status = JOB_STATUS_CANCELLED
+            snapshot.completed_at = _utcnow()
+            snapshot.error_message = "Cancelled after operator stop request; no active worker handle remained"
+        datasource = await db.get(DataSource, task.datasource_id)
+        if datasource is not None:
+            datasource.last_status = JOB_STATUS_CANCELLED
+
+    task.status = JOB_STATUS_CANCELLED
+    task.phase = JOB_STATUS_CANCELLED
+    task.phase_message = "任务已停止"
+    task.completed_at = _utcnow()
+    task.requested_cancel_at = task.requested_cancel_at or _utcnow()
+    task.cancel_reason = reason
+    await db.commit()
+    await db.refresh(task)
+    await _broadcast_task_update(task)
+    return task
+
+
+async def get_active_datasource_job(
+    db: AsyncSession,
+    datasource_id: int,
+    *,
+    task_types: tuple[str, ...] = DATA_WRITE_JOB_TYPES,
+) -> CollectionTask | None:
+    result = await db.execute(
+        select(CollectionTask)
+        .where(CollectionTask.datasource_id == datasource_id)
+        .where(CollectionTask.task_type.in_(task_types))
+        .where(CollectionTask.status.in_(ACTIVE_JOB_STATUSES))
+        .order_by(CollectionTask.created_at.desc().nullslast(), CollectionTask.id.desc())
+        .limit(1)
+    )
+    return result.scalar_one_or_none()
+
+
+async def _get_active_job_by_dedupe_key(db: AsyncSession, dedupe_key: str) -> CollectionTask | None:
+    result = await db.execute(
+        select(CollectionTask)
+        .where(CollectionTask.dedupe_key == dedupe_key)
+        .where(CollectionTask.status.in_(ACTIVE_JOB_STATUSES))
+        .order_by(CollectionTask.created_at.desc().nullslast(), CollectionTask.id.desc())
+        .limit(1)
+    )
+    return result.scalar_one_or_none()
+
+
+async def _get_or_create_virtual_datasource(db: AsyncSession, source: str) -> DataSource:
+    result = await db.execute(select(DataSource).where(DataSource.source == source))
+    datasource = result.scalar_one_or_none()
+    if datasource is not None:
+        return datasource
+
+    datasource = DataSource(
+        name=f"Earth refresh: {source}",
+        source=source,
+        module="SYS",
+        collector_class="EarthRefreshJob",
+        is_active=True,
+    )
+    db.add(datasource)
+    await db.commit()
+    await db.refresh(datasource)
+    return datasource
+
+
+async def _broadcast_task_update(task: CollectionTask) -> None:
+    await broadcaster.broadcast_datasource_task_update(
+        {
+            "datasource_id": task.datasource_id,
+            "collector_name": task.source,
+            "task_id": task.id,
+            "task_type": task.task_type,
+            "status": task.status,
+            "phase": task.phase,
+            "phase_progress": task.phase_progress,
+            "phase_message": task.phase_message,
+            "phase_current": task.phase_current,
+            "phase_total": task.phase_total,
+            "phase_unit": task.phase_unit,
+            "progress": task.progress,
+            "records_processed": task.records_processed,
+            "total_records": task.total_records,
+            "started_at": to_iso8601_utc(task.started_at),
+            "completed_at": to_iso8601_utc(task.completed_at),
+            "requested_cancel_at": to_iso8601_utc(task.requested_cancel_at),
+            "error_message": task.error_message,
+        }
+    )
+
+
+class DataJobWorker:
+    def __init__(self, *, concurrency: int = DEFAULT_WORKER_CONCURRENCY) -> None:
+        self.worker_id = _job_worker_id()
+        self.concurrency = max(1, concurrency)
+        self._task: asyncio.Task[None] | None = None
+        self._stop_event: asyncio.Event | None = None
+        self._running: set[asyncio.Task[Any]] = set()
+
+    def start(self) -> None:
+        if self._task and not self._task.done():
+            return
+        self._stop_event = asyncio.Event()
+        self._task = asyncio.create_task(self._run(), name="data-job-worker")
+
+    async def stop(self) -> None:
+        if self._stop_event:
+            self._stop_event.set()
+        for task in list(self._running):
+            task.cancel()
+        if self._task:
+            await asyncio.gather(self._task, return_exceptions=True)
+        if self._running:
+            await asyncio.gather(*self._running, return_exceptions=True)
+
+    async def _run(self) -> None:
+        assert self._stop_event is not None
+        await self._recover_stale_running_jobs()
+        while not self._stop_event.is_set():
+            self._running = {task for task in self._running if not task.done()}
+            if len(self._running) >= self.concurrency:
+                await asyncio.sleep(QUEUE_POLL_SECONDS)
+                continue
+
+            task_id = await self._claim_next_job()
+            if task_id is None:
+                await asyncio.sleep(QUEUE_POLL_SECONDS)
+                continue
+
+            runner = asyncio.create_task(self._run_claimed_job(task_id), name=f"data-job:{task_id}")
+            self._running.add(runner)
+
+    async def _recover_stale_running_jobs(self) -> None:
+        cutoff = _utcnow() - timedelta(minutes=JOB_STALE_LOCK_MINUTES)
+        async with async_session_factory() as db:
+            result = await db.execute(
+                select(CollectionTask)
+                .where(CollectionTask.status.in_((JOB_STATUS_RUNNING, JOB_STATUS_CANCELLING)))
+                .where(CollectionTask.locked_at.is_not(None))
+                .where(CollectionTask.locked_at < cutoff)
+            )
+            stale_jobs = list(result.scalars().all())
+            for job in stale_jobs:
+                job.status = JOB_STATUS_FAILED
+                job.phase = JOB_STATUS_FAILED
+                job.completed_at = _utcnow()
+                job.error_message = "Marked failed after stale data job lock timeout"
+            if stale_jobs:
+                await db.commit()
+
+    async def _claim_next_job(self) -> int | None:
+        async with async_session_factory() as db:
+            row = await db.execute(
+                text(
+                    """
+                    SELECT queued.id
+                    FROM collection_tasks AS queued
+                    WHERE queued.status = :queued_status
+                      AND NOT EXISTS (
+                        SELECT 1
+                        FROM collection_tasks AS active
+                        WHERE active.source = queued.source
+                          AND active.id <> queued.id
+                          AND active.status IN :active_statuses
+                      )
+                    ORDER BY queued.created_at ASC NULLS FIRST, queued.id ASC
+                    LIMIT 1
+                    FOR UPDATE SKIP LOCKED
+                    """
+                ).bindparams(bindparam("active_statuses", expanding=True)),
+                {
+                    "queued_status": JOB_STATUS_QUEUED,
+                    "active_statuses": SOURCE_LOCK_JOB_STATUSES,
+                },
+            )
+            task_id = row.scalar_one_or_none()
+            if task_id is None:
+                return None
+
+            task = await db.get(CollectionTask, int(task_id))
+            if task is None:
+                return None
+            task.status = JOB_STATUS_RUNNING
+            task.phase = "starting"
+            task.phase_message = "任务开始执行"
+            task.started_at = task.started_at or _utcnow()
+            task.worker_id = self.worker_id
+            task.locked_at = _utcnow()
+            await db.commit()
+            await _broadcast_task_update(task)
+            return int(task_id)
+
+    async def _run_claimed_job(self, task_id: int) -> None:
+        current_task = asyncio.current_task()
+        if current_task is not None:
+            RUNNING_DATA_JOB_TASKS[task_id] = current_task
+        try:
+            async with async_session_factory() as db:
+                task = await db.get(CollectionTask, task_id)
+                if task is None:
+                    return
+                await self._execute_job(db, task)
+        except asyncio.CancelledError:
+            async with async_session_factory() as db:
+                task = await db.get(CollectionTask, task_id)
+                if task is not None and not is_terminal_job_status(task.status):
+                    task.status = JOB_STATUS_CANCELLED
+                    task.phase = JOB_STATUS_CANCELLED
+                    task.phase_message = "任务已停止"
+                    task.completed_at = _utcnow()
+                    await db.commit()
+                    await _broadcast_task_update(task)
+            raise
+        except Exception as exc:
+            logger.exception_event(
+                "Data job failed",
+                event="data_jobs.job_failed",
+                context={"task_id": task_id, "error": str(exc)},
+            )
+            async with async_session_factory() as db:
+                task = await db.get(CollectionTask, task_id)
+                if task is not None:
+                    task.status = JOB_STATUS_FAILED
+                    task.phase = JOB_STATUS_FAILED
+                    task.phase_message = str(exc)
+                    task.error_message = str(exc)
+                    task.completed_at = _utcnow()
+                    await db.commit()
+                    await _broadcast_task_update(task)
+        finally:
+            RUNNING_DATA_JOB_TASKS.pop(task_id, None)
+
+    async def _execute_job(self, db: AsyncSession, task: CollectionTask) -> None:
+        if task.task_type == JOB_TYPE_COLLECT:
+            await _run_collect_job(db, task)
+        elif task.task_type == JOB_TYPE_CLEAR_DATA:
+            await _run_clear_data_job(db, task)
+        elif task.task_type == JOB_TYPE_CLEAR_CACHE:
+            await _run_clear_cache_job(db, task)
+        elif task.task_type == JOB_TYPE_EARTH_REFRESH:
+            await _run_earth_refresh_job(db, task)
+        else:
+            raise RuntimeError(f"Unsupported data job type: {task.task_type}")
+
+
+async def _run_collect_job(db: AsyncSession, task: CollectionTask) -> None:
+    datasource = await db.get(DataSource, task.datasource_id)
+    if datasource is None:
+        raise RuntimeError("Data source not found")
+
+    collector = collector_registry.get(datasource.source)
+    if collector is None:
+        raise RuntimeError(f"Collector '{datasource.source}' not found")
+    if not datasource.is_active:
+        raise RuntimeError("Data source is disabled")
+
+    collector._datasource_id = datasource.id
+    collector._current_task = task
+    collector._db_session = db
+    result = await collector.run(db)
+
+    datasource.last_run_at = _utcnow()
+    datasource.last_status = result.get("status")
+    if datasource.last_status == JOB_STATUS_SUCCESS:
+        effective_candidate = await get_builtin_effective_candidate(db, datasource.source)
+        checksum, _credential_context = await build_builtin_connectivity_checksum(
+            datasource.source,
+            effective_candidate["endpoint"],
+            effective_candidate["auth_type"],
+            effective_candidate["headers"],
+            effective_candidate["config"],
+            db,
+        )
+        await save_connectivity_success(
+            db,
+            datasource.source,
+            checksum,
+            {"status_code": None},
+            connected_by="collection",
+        )
+    await db.commit()
+    await sync_datasource_job(datasource.id)
+
+
+async def _run_clear_data_job(db: AsyncSession, task: CollectionTask) -> None:
+    source = str(task.source or (task.payload or {}).get("source") or "").strip()
+    if not source:
+        raise RuntimeError("Clear data job has no source")
+
+    task.phase = "clearing_data"
+    task.phase_message = "正在删除数据库数据"
+    await db.commit()
+    await _broadcast_task_update(task)
+
+    count_result = await db.execute(
+        select(CollectedData.id).where(CollectedData.source == source)
+    )
+    collected_ids = [row[0] for row in count_result.all()]
+    derived_deleted_counts = await clear_derived_datasource_data(db, source)
+    if collected_ids:
+        await db.execute(CollectedData.__table__.delete().where(CollectedData.id.in_(collected_ids)))
+    deleted_count = len(collected_ids)
+    derived_deleted_count = sum(derived_deleted_counts.values())
+
+    task.records_processed = deleted_count + derived_deleted_count
+    task.total_records = task.records_processed
+    task.progress = 100.0
+    task.phase_progress = 100.0
+    task.phase_current = task.records_processed
+    task.phase_total = task.records_processed
+    task.phase_unit = "records"
+    task.payload = {
+        **(task.payload or {}),
+        "deleted_count": deleted_count,
+        "derived_deleted_count": derived_deleted_count,
+        "derived_deleted_counts": derived_deleted_counts,
+    }
+    task.status = JOB_STATUS_SUCCESS
+    task.phase = "completed"
+    task.phase_message = "数据库数据已清理"
+    task.completed_at = _utcnow()
+    await db.commit()
+    await _broadcast_task_update(task)
+
+
+async def _run_clear_cache_job(db: AsyncSession, task: CollectionTask) -> None:
+    source = str(task.source or (task.payload or {}).get("source") or "").strip()
+    if not source:
+        raise RuntimeError("Clear cache job has no source")
+
+    earth_deleted_count = invalidate_earth_layer_cache_for_source(source)
+    dashboard_deleted_count = int(cache.delete("dashboard:stats")) + int(cache.delete("dashboard:summary"))
+    deleted_count = earth_deleted_count + dashboard_deleted_count
+
+    task.records_processed = deleted_count
+    task.total_records = deleted_count
+    task.progress = 100.0
+    task.phase_progress = 100.0
+    task.phase = "completed"
+    task.phase_message = "缓存已清理"
+    task.payload = {
+        **(task.payload or {}),
+        "earth_layer_deleted_count": earth_deleted_count,
+        "dashboard_deleted_count": dashboard_deleted_count,
+    }
+    task.status = JOB_STATUS_SUCCESS
+    task.completed_at = _utcnow()
+    await db.commit()
+    await _broadcast_task_update(task)
+    await enqueue_earth_refresh_job(db, source=source, payload={"operation": "CACHE_INVALIDATED"})
+
+
+async def _run_earth_refresh_job(db: AsyncSession, task: CollectionTask) -> None:
+    payload = task.payload or {}
+    source = str(payload.get("source") or task.source or "").strip()
+    layers = list(payload.get("layers") or get_earth_update_layers_for_source(source))
+    if not source or not layers:
+        task.status = JOB_STATUS_SUCCESS
+        task.phase = "completed"
+        task.phase_message = "没有需要刷新的 Earth 图层"
+        task.completed_at = _utcnow()
+        await db.commit()
+        await _broadcast_task_update(task)
+        return
+
+    deleted_cache_entries = invalidate_earth_layer_cache_for_source(source)
+    update_payload = {
+        "event": "earth.layer.changed",
+        "action": "database_changed",
+        "source": source,
+        "table": payload.get("table"),
+        "data_type": source,
+        "layers": layers,
+        "refresh_strategy": payload.get("refresh_strategy") or "clear_then_reload",
+        "records_processed": payload.get("records_processed", 0),
+        "operations": payload.get("operations") or [payload.get("operation") or "CHANGE"],
+        "operation": payload.get("operation"),
+        "cache_entries_invalidated": deleted_cache_entries,
+        "timestamp": to_iso8601_utc(_utcnow()),
+    }
+    if payload.get("entity") == "interactable":
+        update_payload.update(
+            {
+                "entity": "interactable",
+                "action": payload.get("action") or "changed",
+                "ids": payload.get("ids") or payload.get("entity_keys") or [],
+                "item": payload.get("item"),
+            }
+        )
+    await broadcaster.broadcast_earth_update(update_payload)
+
+    task.records_processed = int(payload.get("records_processed") or 0)
+    task.progress = 100.0
+    task.phase_progress = 100.0
+    task.phase = "completed"
+    task.phase_message = "Earth 图层刷新通知已发送"
+    task.status = JOB_STATUS_SUCCESS
+    task.completed_at = _utcnow()
+    task.payload = {**payload, "cache_entries_invalidated": deleted_cache_entries}
+    await db.commit()
+    await _broadcast_task_update(task)
+
+
+_worker = DataJobWorker()
+
+
+def start_data_job_worker() -> None:
+    _worker.start()
+
+
+async def stop_data_job_worker() -> None:
+    await _worker.stop()