fix: optimize visualization hot paths and refine bgp brief workspace

backend/app/api/v1/datasources.py

@@ -13,6 +13,7 @@ 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 (
@@ -40,6 +41,156 @@ def is_due_for_collection(datasource: DataSource, now: datetime) -> bool:
     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_completed_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.completed_at),
+        )
+        .where(CollectionTask.datasource_id.in_(datasource_ids))
+        .where(CollectionTask.completed_at.isnot(None))
+        .where(CollectionTask.status.in_(("success", "failed", "cancelled")))
+        .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_datasource_data_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))
+        .group_by(CollectedData.source)
+    )
+    return {source: count 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[int, CollectionTask], dict[str, int], 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)
+
+    completed_tasks = await _load_latest_completed_tasks(db, datasource_ids)
+    data_counts = await _load_datasource_data_counts(db, sources)
+    endpoint_overrides = await _load_datasource_endpoint_overrides(db, sources)
+    return running_tasks, completed_tasks, data_counts, endpoint_overrides
+
+
 async def get_datasource_record(db: AsyncSession, source_id: str) -> Optional[DataSource]:
     datasource = None
     try:
@@ -58,18 +209,6 @@ async def get_datasource_record(db: AsyncSession, source_id: str) -> Optional[Da
     return result.scalar_one_or_none()
 
 
-async def get_last_completed_task(db: AsyncSession, datasource_id: int) -> Optional[CollectionTask]:
-    result = await db.execute(
-        select(CollectionTask)
-        .where(CollectionTask.datasource_id == datasource_id)
-        .where(CollectionTask.completed_at.isnot(None))
-        .where(CollectionTask.status.in_(("success", "failed", "cancelled")))
-        .order_by(CollectionTask.completed_at.desc())
-        .limit(1)
-    )
-    return result.scalar_one_or_none()
-
-
 async def get_running_task(db: AsyncSession, datasource_id: int) -> Optional[CollectionTask]:
     result = await db.execute(
         select(CollectionTask)
@@ -262,14 +401,17 @@ async def list_datasources(
 
     collector_list = []
     config = get_data_sources_config()
+    running_tasks, completed_tasks, data_counts, endpoint_overrides = await _load_datasource_list_context(
+        db,
+        datasources,
+    )
     for datasource in datasources:
-        running_task = await get_running_task(db, datasource.id)
-        last_task = await get_last_completed_task(db, datasource.id)
-        endpoint = await config.get_url(datasource.source, db)
-        data_count_result = await db.execute(
-            select(func.count(CollectedData.id)).where(CollectedData.source == datasource.source)
+        running_task = running_tasks.get(datasource.id)
+        last_task = completed_tasks.get(datasource.id)
+        endpoint = endpoint_overrides.get(datasource.source) or config.get_yaml_url(
+            datasource.source,
         )
-        data_count = data_count_result.scalar() or 0
+        data_count = data_counts.get(datasource.source, 0)
 
         last_run_at = datasource.last_run_at or (last_task.completed_at if last_task else None)
         last_run = to_iso8601_utc(last_run_at)
@@ -330,9 +472,13 @@ async def trigger_all_datasources(
     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:
-        running_task = await get_running_task(db, datasource.id)
+        running_task = running_tasks.get(datasource.id)
         if running_task is not None:
             skipped_sources.append(
                 {
@@ -360,7 +506,7 @@ async def trigger_all_datasources(
             )
             continue
 
-        previous_task_ids[datasource.id] = await get_latest_task_id_for_datasource(datasource.id)
+        previous_task_ids[datasource.id] = None
         success = run_collector_now(datasource.source)
         if not success:
             failed_sources.append(
@@ -382,13 +528,24 @@ async def trigger_all_datasources(
             }
         )
 
+    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 = await get_latest_task_id_for_datasource(item["id"])
+            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