946 lines
32 KiB
Python
946 lines
32 KiB
Python
import asyncio
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Optional
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from pydantic import BaseModel, Field
|
|
from sqlalchemy import func, or_, select, text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.time import to_iso8601_utc
|
|
from app.core.security import get_current_user
|
|
from app.core.data_sources import get_data_sources_config
|
|
from app.core.datasource_defaults import DEFAULT_DATASOURCES
|
|
from app.db.session import get_db
|
|
from app.models.collected_data import CollectedData
|
|
from app.models.data_snapshot import DataSnapshot
|
|
from app.models.datasource import DataSource
|
|
from app.models.datasource_config import DataSourceConfig
|
|
from app.models.task import CollectionTask
|
|
from app.models.user import User
|
|
from app.services.scheduler import (
|
|
cancel_running_collector_now,
|
|
get_latest_task_id_for_datasource,
|
|
run_collector_now,
|
|
sync_datasource_job,
|
|
)
|
|
|
|
router = APIRouter()
|
|
STALE_RUNNING_TASK_TIMEOUT_MINUTES = 90
|
|
|
|
PRODUCT_SOURCE_KEYWORDS: tuple[tuple[str, tuple[str, ...]], ...] = (
|
|
("vessels", ("vessel", "ais")),
|
|
("cables", ("cable", "landing", "telegeography", "arcgis", "fao")),
|
|
("satellites", ("tle", "satellite", "spacetrack", "celestrak")),
|
|
("bgp", ("bgp", "asn", "prefix_geo", "opengeofeed", "nro")),
|
|
("compute", ("top500", "gpu", "supercomputer", "compute")),
|
|
("ai", ("huggingface", "epoch_ai")),
|
|
("media", ("news", "tv", "live_stream")),
|
|
)
|
|
|
|
|
|
class DatasourceBatchTriggerRequest(BaseModel):
|
|
source_ids: list[int] = Field(default_factory=list)
|
|
force: bool = False
|
|
module: Optional[str] = None
|
|
product: Optional[str] = None
|
|
is_active: Optional[bool] = None
|
|
priority: Optional[str] = None
|
|
run_status: Optional[str] = None
|
|
collected: Optional[bool] = None
|
|
credential_status: Optional[str] = None
|
|
q: Optional[str] = None
|
|
|
|
|
|
def format_frequency_label(minutes: int) -> str:
|
|
if minutes % 1440 == 0:
|
|
return f"{minutes // 1440}d"
|
|
if minutes % 60 == 0:
|
|
return f"{minutes // 60}h"
|
|
return f"{minutes}m"
|
|
|
|
|
|
def datasource_metadata(source: str) -> dict:
|
|
info = DEFAULT_DATASOURCES.get(source, {})
|
|
return {
|
|
"display_name": info.get("display_name") or info.get("name") or source,
|
|
"is_free": bool(info.get("is_free", True)),
|
|
"requires_credentials": bool(info.get("requires_credentials", False)),
|
|
"credential_provider": info.get("credential_provider"),
|
|
"credential_status": info.get("credential_status", "none"),
|
|
}
|
|
|
|
|
|
def datasource_product_key(datasource: DataSource) -> str:
|
|
haystack = " ".join(
|
|
[
|
|
datasource.source or "",
|
|
datasource.name or "",
|
|
datasource.collector_class or "",
|
|
]
|
|
).lower()
|
|
for product, keywords in PRODUCT_SOURCE_KEYWORDS:
|
|
if any(keyword in haystack for keyword in keywords):
|
|
return product
|
|
return "other"
|
|
|
|
|
|
def is_due_for_collection(datasource: DataSource, now: datetime) -> bool:
|
|
if datasource.last_run_at is None:
|
|
return True
|
|
return datasource.last_run_at + timedelta(minutes=datasource.frequency_minutes) <= now
|
|
|
|
|
|
def _task_rank_column(order_column):
|
|
return func.row_number().over(
|
|
partition_by=CollectionTask.datasource_id,
|
|
order_by=(order_column.desc().nullslast(), CollectionTask.id.desc()),
|
|
).label("row_num")
|
|
|
|
|
|
async def _load_latest_running_tasks(
|
|
db: AsyncSession,
|
|
datasource_ids: list[int],
|
|
) -> dict[int, CollectionTask]:
|
|
if not datasource_ids:
|
|
return {}
|
|
|
|
ranked_tasks = (
|
|
select(
|
|
CollectionTask.id.label("task_id"),
|
|
_task_rank_column(CollectionTask.started_at),
|
|
)
|
|
.where(CollectionTask.datasource_id.in_(datasource_ids))
|
|
.where(CollectionTask.status == "running")
|
|
.subquery()
|
|
)
|
|
result = await db.execute(
|
|
select(CollectionTask)
|
|
.join(ranked_tasks, CollectionTask.id == ranked_tasks.c.task_id)
|
|
.where(ranked_tasks.c.row_num == 1)
|
|
)
|
|
return {task.datasource_id: task for task in result.scalars().all()}
|
|
|
|
|
|
async def _load_latest_task_ids(
|
|
db: AsyncSession,
|
|
datasource_ids: list[int],
|
|
) -> dict[int, int]:
|
|
if not datasource_ids:
|
|
return {}
|
|
|
|
ranked_tasks = (
|
|
select(
|
|
CollectionTask.id.label("task_id"),
|
|
CollectionTask.datasource_id.label("datasource_id"),
|
|
func.row_number().over(
|
|
partition_by=CollectionTask.datasource_id,
|
|
order_by=CollectionTask.id.desc(),
|
|
).label("row_num"),
|
|
)
|
|
.where(CollectionTask.datasource_id.in_(datasource_ids))
|
|
.subquery()
|
|
)
|
|
result = await db.execute(
|
|
select(ranked_tasks.c.datasource_id, ranked_tasks.c.task_id)
|
|
.where(ranked_tasks.c.row_num == 1)
|
|
)
|
|
return {datasource_id: task_id for datasource_id, task_id in result.all()}
|
|
|
|
|
|
async def _load_collected_record_counts(
|
|
db: AsyncSession,
|
|
sources: list[str],
|
|
) -> dict[str, int]:
|
|
if not sources:
|
|
return {}
|
|
|
|
result = await db.execute(
|
|
select(CollectedData.source, func.count(CollectedData.id))
|
|
.where(CollectedData.source.in_(sources))
|
|
.where(CollectedData.is_current.is_(True))
|
|
.group_by(CollectedData.source)
|
|
)
|
|
return {source: int(count or 0) for source, count in result.all()}
|
|
|
|
|
|
async def _load_datasource_endpoint_overrides(
|
|
db: AsyncSession,
|
|
sources: list[str],
|
|
) -> dict[str, str]:
|
|
if not sources:
|
|
return {}
|
|
|
|
result = await db.execute(
|
|
select(DataSourceConfig.name, DataSourceConfig.endpoint)
|
|
.where(DataSourceConfig.name.in_(sources))
|
|
.where(DataSourceConfig.is_active.is_(True))
|
|
.where(DataSourceConfig.endpoint.isnot(None))
|
|
)
|
|
return {
|
|
name: endpoint
|
|
for name, endpoint in result.all()
|
|
if endpoint
|
|
}
|
|
|
|
|
|
async def _load_datasource_list_context(
|
|
db: AsyncSession,
|
|
datasources: list[DataSource],
|
|
) -> tuple[dict[int, CollectionTask], dict[str, str]]:
|
|
datasource_ids = [datasource.id for datasource in datasources]
|
|
sources = [datasource.source for datasource in datasources]
|
|
|
|
running_tasks = await _load_latest_running_tasks(db, datasource_ids)
|
|
datasource_by_id = {datasource.id: datasource for datasource in datasources}
|
|
now = datetime.now(timezone.utc)
|
|
|
|
stale_datasource_ids: list[int] = []
|
|
for datasource_id, task in running_tasks.items():
|
|
started_at = task.started_at
|
|
if started_at is None:
|
|
continue
|
|
if started_at.tzinfo is None:
|
|
started_at = started_at.replace(tzinfo=timezone.utc)
|
|
if now - started_at > timedelta(minutes=STALE_RUNNING_TASK_TIMEOUT_MINUTES):
|
|
datasource = datasource_by_id.get(datasource_id)
|
|
if datasource is not None:
|
|
await fail_and_rollback_stale_running_task(db, datasource, task)
|
|
stale_datasource_ids.append(datasource_id)
|
|
|
|
if stale_datasource_ids:
|
|
running_tasks = await _load_latest_running_tasks(db, datasource_ids)
|
|
|
|
endpoint_overrides = await _load_datasource_endpoint_overrides(db, sources)
|
|
return running_tasks, endpoint_overrides
|
|
|
|
|
|
def _apply_datasource_query_filters(
|
|
query,
|
|
*,
|
|
module: Optional[str] = None,
|
|
is_active: Optional[bool] = None,
|
|
priority: Optional[str] = None,
|
|
run_status: Optional[str] = None,
|
|
q: Optional[str] = None,
|
|
) -> object:
|
|
if module:
|
|
query = query.where(DataSource.module == module)
|
|
if is_active is not None:
|
|
query = query.where(DataSource.is_active == is_active)
|
|
if priority:
|
|
query = query.where(DataSource.priority == priority)
|
|
if run_status and run_status not in {"running", "collected", "uncollected"}:
|
|
if run_status == "not_run":
|
|
query = query.where(DataSource.last_status.is_(None))
|
|
else:
|
|
query = query.where(DataSource.last_status == run_status)
|
|
if q:
|
|
like_value = f"%{q.strip()}%"
|
|
query = query.where(
|
|
or_(
|
|
DataSource.name.ilike(like_value),
|
|
DataSource.source.ilike(like_value),
|
|
DataSource.collector_class.ilike(like_value),
|
|
)
|
|
)
|
|
return query
|
|
|
|
|
|
def _filter_datasources_in_memory(
|
|
datasources: list[DataSource],
|
|
*,
|
|
running_tasks: dict[int, CollectionTask],
|
|
record_counts: dict[str, int],
|
|
product: Optional[str] = None,
|
|
run_status: Optional[str] = None,
|
|
collected: Optional[bool] = None,
|
|
credential_status: Optional[str] = None,
|
|
) -> list[DataSource]:
|
|
filtered: list[DataSource] = []
|
|
for datasource in datasources:
|
|
record_count = record_counts.get(datasource.source, 0)
|
|
if product and datasource_product_key(datasource) != product:
|
|
continue
|
|
if collected is not None and (record_count > 0) != collected:
|
|
continue
|
|
if credential_status:
|
|
metadata = datasource_metadata(datasource.source)
|
|
if metadata["credential_status"] != credential_status:
|
|
continue
|
|
if run_status == "running" and datasource.id not in running_tasks:
|
|
continue
|
|
if run_status == "collected" and record_count <= 0:
|
|
continue
|
|
if run_status == "uncollected" and record_count > 0:
|
|
continue
|
|
filtered.append(datasource)
|
|
return filtered
|
|
|
|
|
|
async def _trigger_datasource_batch(
|
|
db: AsyncSession,
|
|
datasources: list[DataSource],
|
|
*,
|
|
force: bool,
|
|
) -> dict:
|
|
if not datasources:
|
|
return {
|
|
"status": "noop",
|
|
"message": "No matching data sources to trigger",
|
|
"force": force,
|
|
"triggered": [],
|
|
"skipped": [],
|
|
"failed": [],
|
|
}
|
|
|
|
previous_task_ids: dict[int, Optional[int]] = {}
|
|
triggered_sources: list[dict] = []
|
|
skipped_sources: list[dict] = []
|
|
failed_sources: list[dict] = []
|
|
now = datetime.now(timezone.utc)
|
|
running_tasks = await _load_latest_running_tasks(
|
|
db,
|
|
[datasource.id for datasource in datasources],
|
|
)
|
|
|
|
for datasource in datasources:
|
|
if not datasource.is_active:
|
|
skipped_sources.append(
|
|
{
|
|
"id": datasource.id,
|
|
"source": datasource.source,
|
|
"name": datasource.name,
|
|
"reason": "disabled",
|
|
}
|
|
)
|
|
continue
|
|
|
|
running_task = running_tasks.get(datasource.id)
|
|
if running_task is not None:
|
|
if not force:
|
|
skipped_sources.append(
|
|
{
|
|
"id": datasource.id,
|
|
"source": datasource.source,
|
|
"name": datasource.name,
|
|
"reason": "already_running",
|
|
"task_id": running_task.id,
|
|
}
|
|
)
|
|
continue
|
|
cancelled = await cancel_running_collector_now(datasource.source)
|
|
if not cancelled:
|
|
await rollback_orphaned_running_task(db, datasource, running_task)
|
|
|
|
if not force and not is_due_for_collection(datasource, now):
|
|
skipped_sources.append(
|
|
{
|
|
"id": datasource.id,
|
|
"source": datasource.source,
|
|
"name": datasource.name,
|
|
"reason": "within_frequency_window",
|
|
"last_run_at": to_iso8601_utc(datasource.last_run_at),
|
|
"next_run_at": to_iso8601_utc(
|
|
datasource.last_run_at + timedelta(minutes=datasource.frequency_minutes)
|
|
),
|
|
}
|
|
)
|
|
continue
|
|
|
|
previous_task_ids[datasource.id] = None
|
|
success = run_collector_now(datasource.source)
|
|
if not success:
|
|
failed_sources.append(
|
|
{
|
|
"id": datasource.id,
|
|
"source": datasource.source,
|
|
"name": datasource.name,
|
|
"reason": "trigger_failed",
|
|
}
|
|
)
|
|
continue
|
|
|
|
triggered_sources.append(
|
|
{
|
|
"id": datasource.id,
|
|
"source": datasource.source,
|
|
"name": datasource.name,
|
|
"task_id": None,
|
|
}
|
|
)
|
|
|
|
latest_task_ids = await _load_latest_task_ids(
|
|
db,
|
|
[datasource.id for datasource in datasources],
|
|
)
|
|
for datasource_id in previous_task_ids:
|
|
previous_task_ids[datasource_id] = latest_task_ids.get(datasource_id)
|
|
|
|
for _ in range(20):
|
|
await asyncio.sleep(0.1)
|
|
pending = [item for item in triggered_sources if item["task_id"] is None]
|
|
if not pending:
|
|
break
|
|
latest_task_ids = await _load_latest_task_ids(
|
|
db,
|
|
[item["id"] for item in pending],
|
|
)
|
|
for item in pending:
|
|
task_id = latest_task_ids.get(item["id"])
|
|
if task_id is not None and task_id != previous_task_ids.get(item["id"]):
|
|
item["task_id"] = task_id
|
|
|
|
return {
|
|
"status": "triggered" if triggered_sources else "partial",
|
|
"message": f"Triggered {len(triggered_sources)} data sources",
|
|
"force": force,
|
|
"triggered": triggered_sources,
|
|
"skipped": skipped_sources,
|
|
"failed": failed_sources,
|
|
}
|
|
|
|
|
|
async def get_datasource_record(db: AsyncSession, source_id: str) -> Optional[DataSource]:
|
|
datasource = None
|
|
try:
|
|
datasource = await db.get(DataSource, int(source_id))
|
|
except ValueError:
|
|
pass
|
|
|
|
if datasource is not None:
|
|
return datasource
|
|
|
|
result = await db.execute(
|
|
select(DataSource).where(
|
|
(DataSource.source == source_id) | (DataSource.collector_class == source_id)
|
|
)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def get_running_task(db: AsyncSession, datasource_id: int) -> Optional[CollectionTask]:
|
|
result = await db.execute(
|
|
select(CollectionTask)
|
|
.where(CollectionTask.datasource_id == datasource_id)
|
|
.where(CollectionTask.status == "running")
|
|
.order_by(CollectionTask.started_at.desc())
|
|
.limit(1)
|
|
)
|
|
task = result.scalar_one_or_none()
|
|
if not task:
|
|
return None
|
|
|
|
started_at = task.started_at
|
|
if started_at is None:
|
|
return task
|
|
|
|
now = datetime.now(timezone.utc)
|
|
if started_at.tzinfo is None:
|
|
started_at = started_at.replace(tzinfo=timezone.utc)
|
|
|
|
if now - started_at <= timedelta(minutes=STALE_RUNNING_TASK_TIMEOUT_MINUTES):
|
|
return task
|
|
|
|
datasource = await db.get(DataSource, datasource_id)
|
|
if datasource is not None:
|
|
await fail_and_rollback_stale_running_task(db, datasource, task)
|
|
else:
|
|
existing_error = (task.error_message or "").strip()
|
|
stale_reason = (
|
|
f"Marked failed automatically after stale running timeout "
|
|
f"({STALE_RUNNING_TASK_TIMEOUT_MINUTES}m)"
|
|
)
|
|
task.status = "failed"
|
|
task.phase = "failed"
|
|
task.completed_at = now
|
|
task.error_message = f"{existing_error}\n{stale_reason}".strip() if existing_error else stale_reason
|
|
await db.commit()
|
|
return None
|
|
|
|
|
|
async def rollback_orphaned_running_task(
|
|
db: AsyncSession,
|
|
datasource: DataSource,
|
|
running_task: CollectionTask,
|
|
) -> None:
|
|
snapshot_result = await db.execute(
|
|
select(DataSnapshot)
|
|
.where(
|
|
DataSnapshot.datasource_id == datasource.id,
|
|
DataSnapshot.task_id == running_task.id,
|
|
)
|
|
.order_by(DataSnapshot.id.desc())
|
|
.limit(1)
|
|
)
|
|
snapshot = snapshot_result.scalar_one_or_none()
|
|
|
|
await db.execute(CollectedData.__table__.delete().where(CollectedData.task_id == running_task.id))
|
|
|
|
await db.execute(
|
|
text(
|
|
"""
|
|
UPDATE collected_data
|
|
SET is_current = FALSE
|
|
WHERE source = :source
|
|
"""
|
|
),
|
|
{"source": datasource.source},
|
|
)
|
|
|
|
if snapshot is not None:
|
|
snapshot.status = "cancelled"
|
|
snapshot.is_current = False
|
|
snapshot.completed_at = datetime.now(timezone.utc)
|
|
summary = dict(snapshot.summary or {})
|
|
summary["rollback"] = True
|
|
summary["rollback_reason"] = "orphaned_running_task_after_backend_restart"
|
|
snapshot.summary = summary
|
|
|
|
if snapshot.parent_snapshot_id is not None:
|
|
parent_snapshot = await db.get(DataSnapshot, snapshot.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": snapshot.parent_snapshot_id},
|
|
)
|
|
|
|
running_task.status = "cancelled"
|
|
running_task.phase = "cancelled"
|
|
running_task.completed_at = datetime.now(timezone.utc)
|
|
existing_error = (running_task.error_message or "").strip()
|
|
cancel_reason = "Cancelled after backend restart because the running task handle was lost; incomplete writes rolled back"
|
|
running_task.error_message = f"{existing_error}\n{cancel_reason}".strip() if existing_error else cancel_reason
|
|
datasource.last_status = "cancelled"
|
|
datasource.last_run_at = datetime.now(timezone.utc)
|
|
await db.commit()
|
|
|
|
|
|
async def fail_and_rollback_stale_running_task(
|
|
db: AsyncSession,
|
|
datasource: DataSource,
|
|
running_task: CollectionTask,
|
|
) -> None:
|
|
snapshot_result = await db.execute(
|
|
select(DataSnapshot)
|
|
.where(
|
|
DataSnapshot.datasource_id == datasource.id,
|
|
DataSnapshot.task_id == running_task.id,
|
|
)
|
|
.order_by(DataSnapshot.id.desc())
|
|
.limit(1)
|
|
)
|
|
snapshot = snapshot_result.scalar_one_or_none()
|
|
|
|
await db.execute(CollectedData.__table__.delete().where(CollectedData.task_id == running_task.id))
|
|
|
|
await db.execute(
|
|
text(
|
|
"""
|
|
UPDATE collected_data
|
|
SET is_current = FALSE
|
|
WHERE source = :source
|
|
"""
|
|
),
|
|
{"source": datasource.source},
|
|
)
|
|
|
|
if snapshot is not None:
|
|
snapshot.status = "failed"
|
|
snapshot.is_current = False
|
|
snapshot.completed_at = datetime.now(timezone.utc)
|
|
summary = dict(snapshot.summary or {})
|
|
summary["rollback"] = True
|
|
summary["rollback_reason"] = "stale_running_task_timeout"
|
|
snapshot.summary = summary
|
|
|
|
if snapshot.parent_snapshot_id is not None:
|
|
parent_snapshot = await db.get(DataSnapshot, snapshot.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": snapshot.parent_snapshot_id},
|
|
)
|
|
|
|
existing_error = (running_task.error_message or "").strip()
|
|
stale_reason = (
|
|
f"Marked failed automatically after stale running timeout "
|
|
f"({STALE_RUNNING_TASK_TIMEOUT_MINUTES}m); incomplete writes rolled back"
|
|
)
|
|
running_task.status = "failed"
|
|
running_task.phase = "failed"
|
|
running_task.completed_at = datetime.now(timezone.utc)
|
|
running_task.error_message = f"{existing_error}\n{stale_reason}".strip() if existing_error else stale_reason
|
|
datasource.last_status = "failed"
|
|
datasource.last_run_at = datetime.now(timezone.utc)
|
|
await db.commit()
|
|
|
|
|
|
@router.get("")
|
|
async def list_datasources(
|
|
module: Optional[str] = None,
|
|
is_active: Optional[bool] = None,
|
|
priority: Optional[str] = None,
|
|
product: Optional[str] = None,
|
|
run_status: Optional[str] = None,
|
|
collected: Optional[bool] = None,
|
|
credential_status: Optional[str] = None,
|
|
q: Optional[str] = None,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
query = select(DataSource).order_by(DataSource.module, DataSource.id)
|
|
query = _apply_datasource_query_filters(
|
|
query,
|
|
module=module,
|
|
is_active=is_active,
|
|
priority=priority,
|
|
run_status=run_status,
|
|
q=q,
|
|
)
|
|
|
|
result = await db.execute(query)
|
|
datasources = result.scalars().all()
|
|
|
|
collector_list = []
|
|
config = get_data_sources_config()
|
|
running_tasks, endpoint_overrides = await _load_datasource_list_context(db, datasources)
|
|
record_counts = await _load_collected_record_counts(db, [datasource.source for datasource in datasources])
|
|
datasources = _filter_datasources_in_memory(
|
|
datasources,
|
|
running_tasks=running_tasks,
|
|
record_counts=record_counts,
|
|
product=product,
|
|
run_status=run_status,
|
|
collected=collected,
|
|
credential_status=credential_status,
|
|
)
|
|
for datasource in datasources:
|
|
running_task = running_tasks.get(datasource.id)
|
|
endpoint = endpoint_overrides.get(datasource.source) or config.get_yaml_url(datasource.source)
|
|
last_run_at = datasource.last_run_at
|
|
last_status = datasource.last_status
|
|
collected_records = record_counts.get(datasource.source, 0)
|
|
|
|
collector_list.append(
|
|
{
|
|
"id": datasource.id,
|
|
"source": datasource.source,
|
|
"name": datasource.name,
|
|
**datasource_metadata(datasource.source),
|
|
"product": datasource_product_key(datasource),
|
|
"module": datasource.module,
|
|
"priority": datasource.priority,
|
|
"frequency": format_frequency_label(datasource.frequency_minutes),
|
|
"frequency_minutes": datasource.frequency_minutes,
|
|
"is_active": datasource.is_active,
|
|
"collector_class": datasource.collector_class,
|
|
"endpoint": endpoint,
|
|
"last_run": to_iso8601_utc(last_run_at),
|
|
"last_run_at": to_iso8601_utc(last_run_at),
|
|
"last_status": last_status,
|
|
"is_running": running_task is not None,
|
|
"task_id": running_task.id if running_task else None,
|
|
"progress": running_task.progress if running_task else None,
|
|
"phase": running_task.phase if running_task else None,
|
|
"phase_progress": running_task.phase_progress if running_task else None,
|
|
"phase_message": running_task.phase_message if running_task else None,
|
|
"phase_current": running_task.phase_current if running_task else None,
|
|
"phase_total": running_task.phase_total if running_task else None,
|
|
"phase_unit": running_task.phase_unit if running_task else None,
|
|
"records_processed": running_task.records_processed if running_task else None,
|
|
"total_records": running_task.total_records if running_task else None,
|
|
"collected_records": collected_records,
|
|
"has_collected_data": collected_records > 0,
|
|
}
|
|
)
|
|
|
|
return {"total": len(collector_list), "data": collector_list}
|
|
|
|
|
|
@router.post("/trigger-all")
|
|
async def trigger_all_datasources(
|
|
force: bool = Query(False),
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
result = await db.execute(
|
|
select(DataSource)
|
|
.where(DataSource.is_active.is_(True))
|
|
.order_by(DataSource.module, DataSource.id)
|
|
)
|
|
datasources = result.scalars().all()
|
|
return await _trigger_datasource_batch(db, datasources, force=force)
|
|
|
|
|
|
@router.post("/trigger-batch")
|
|
async def trigger_datasource_batch(
|
|
payload: DatasourceBatchTriggerRequest,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
query = select(DataSource).order_by(DataSource.module, DataSource.id)
|
|
if payload.source_ids:
|
|
query = query.where(DataSource.id.in_(payload.source_ids))
|
|
else:
|
|
query = _apply_datasource_query_filters(
|
|
query,
|
|
module=payload.module,
|
|
is_active=payload.is_active,
|
|
priority=payload.priority,
|
|
run_status=payload.run_status,
|
|
q=payload.q,
|
|
)
|
|
|
|
result = await db.execute(query)
|
|
datasources = result.scalars().all()
|
|
running_tasks, _ = await _load_datasource_list_context(db, datasources)
|
|
record_counts = await _load_collected_record_counts(db, [datasource.source for datasource in datasources])
|
|
datasources = _filter_datasources_in_memory(
|
|
datasources,
|
|
running_tasks=running_tasks,
|
|
record_counts=record_counts,
|
|
product=None if payload.source_ids else payload.product,
|
|
run_status=None if payload.source_ids else payload.run_status,
|
|
collected=None if payload.source_ids else payload.collected,
|
|
credential_status=None if payload.source_ids else payload.credential_status,
|
|
)
|
|
return await _trigger_datasource_batch(db, datasources, force=payload.force)
|
|
|
|
|
|
@router.get("/{source_id}")
|
|
async def get_datasource(
|
|
source_id: str,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
config = get_data_sources_config()
|
|
endpoint = await config.get_url(datasource.source, db)
|
|
|
|
return {
|
|
"id": datasource.id,
|
|
"name": datasource.name,
|
|
**datasource_metadata(datasource.source),
|
|
"module": datasource.module,
|
|
"priority": datasource.priority,
|
|
"frequency": format_frequency_label(datasource.frequency_minutes),
|
|
"frequency_minutes": datasource.frequency_minutes,
|
|
"collector_class": datasource.collector_class,
|
|
"source": datasource.source,
|
|
"endpoint": endpoint,
|
|
"is_active": datasource.is_active,
|
|
}
|
|
|
|
|
|
@router.post("/{source_id}/enable")
|
|
async def enable_datasource(
|
|
source_id: str,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
datasource.is_active = True
|
|
await db.commit()
|
|
await sync_datasource_job(datasource.id)
|
|
return {"status": "enabled", "source_id": datasource.id}
|
|
|
|
|
|
@router.post("/{source_id}/disable")
|
|
async def disable_datasource(
|
|
source_id: str,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
datasource.is_active = False
|
|
await db.commit()
|
|
await sync_datasource_job(datasource.id)
|
|
return {"status": "disabled", "source_id": datasource.id}
|
|
|
|
|
|
@router.get("/{source_id}/stats")
|
|
async def get_datasource_stats(
|
|
source_id: str,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
result = await db.execute(
|
|
select(func.count(CollectedData.id)).where(CollectedData.source == datasource.source)
|
|
)
|
|
total = result.scalar() or 0
|
|
|
|
return {
|
|
"source_id": datasource.id,
|
|
"collector_name": datasource.collector_class,
|
|
"name": datasource.name,
|
|
"total_records": total,
|
|
}
|
|
|
|
|
|
@router.post("/{source_id}/trigger")
|
|
async def trigger_datasource(
|
|
source_id: str,
|
|
force: bool = Query(False),
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
if not datasource.is_active:
|
|
raise HTTPException(status_code=400, detail="Data source is disabled")
|
|
|
|
running_task = await get_running_task(db, datasource.id)
|
|
if running_task is not None and not force:
|
|
raise HTTPException(
|
|
status_code=409,
|
|
detail={
|
|
"reason": "running_task_in_progress",
|
|
"message": "当前采集任务尚未完成，重新触发会丢失本次未完成进度。是否强制重新采集？",
|
|
"task_id": running_task.id,
|
|
"phase": running_task.phase,
|
|
"phase_progress": running_task.phase_progress,
|
|
"phase_message": running_task.phase_message,
|
|
"phase_current": running_task.phase_current,
|
|
"phase_total": running_task.phase_total,
|
|
"phase_unit": running_task.phase_unit,
|
|
"progress": running_task.progress,
|
|
"records_processed": running_task.records_processed,
|
|
"total_records": running_task.total_records,
|
|
},
|
|
)
|
|
|
|
if running_task is not None and force:
|
|
cancelled = await cancel_running_collector_now(datasource.source)
|
|
if not cancelled:
|
|
await rollback_orphaned_running_task(db, datasource, running_task)
|
|
|
|
previous_task_id = await get_latest_task_id_for_datasource(datasource.id)
|
|
success = run_collector_now(datasource.source)
|
|
if not success:
|
|
raise HTTPException(status_code=500, detail=f"Failed to trigger collector '{datasource.source}'")
|
|
|
|
task_id = None
|
|
for _ in range(20):
|
|
await asyncio.sleep(0.1)
|
|
task_id = await get_latest_task_id_for_datasource(datasource.id)
|
|
if task_id is not None and task_id != previous_task_id:
|
|
break
|
|
if task_id == previous_task_id:
|
|
task_id = None
|
|
|
|
return {
|
|
"status": "triggered",
|
|
"source_id": datasource.id,
|
|
"task_id": task_id,
|
|
"collector_name": datasource.source,
|
|
"force": force,
|
|
"message": f"Collector '{datasource.source}' has been triggered",
|
|
}
|
|
|
|
|
|
@router.delete("/{source_id}/data")
|
|
async def clear_datasource_data(
|
|
source_id: str,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
result = await db.execute(
|
|
select(func.count(CollectedData.id)).where(CollectedData.source == datasource.source)
|
|
)
|
|
count = result.scalar() or 0
|
|
|
|
if count == 0:
|
|
return {"status": "success", "message": "No data to clear", "deleted_count": 0}
|
|
|
|
delete_query = CollectedData.__table__.delete().where(CollectedData.source == datasource.source)
|
|
await db.execute(delete_query)
|
|
await db.commit()
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": f"Cleared {count} records for data source '{datasource.name}'",
|
|
"deleted_count": count,
|
|
}
|
|
|
|
|
|
@router.get("/{source_id}/task-status")
|
|
async def get_task_status(
|
|
source_id: str,
|
|
task_id: Optional[int] = None,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
datasource = await get_datasource_record(db, source_id)
|
|
if not datasource:
|
|
raise HTTPException(status_code=404, detail="Data source not found")
|
|
|
|
if task_id is not None:
|
|
task = await db.get(CollectionTask, task_id)
|
|
if not task or task.datasource_id != datasource.id:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
else:
|
|
task = await get_running_task(db, datasource.id)
|
|
|
|
if not task:
|
|
return {
|
|
"is_running": False,
|
|
"task_id": None,
|
|
"progress": None,
|
|
"phase": None,
|
|
"phase_progress": None,
|
|
"phase_message": None,
|
|
"phase_current": None,
|
|
"phase_total": None,
|
|
"phase_unit": None,
|
|
"status": "idle",
|
|
}
|
|
|
|
return {
|
|
"is_running": task.status == "running",
|
|
"task_id": task.id,
|
|
"progress": task.progress,
|
|
"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,
|
|
"records_processed": task.records_processed,
|
|
"total_records": task.total_records,
|
|
"status": task.status,
|
|
}
|