feat: optimize dashboard and workload queries with database indexes and SQL aggregation for improved performance

Add database indexes on ip_addresses.address and traffic_flows columns (node_id/source_ip/destination_ip/destination_port/state with updated_at) to accelerate query performance, implement ensure_runtime_indexes to create indexes on startup with IF NOT EXISTS guards, rewrite dashboard_top_talkers to use SQL aggregation with JOIN on IpAddress/TrafficFlow instead of Python loops over all

backend/app/api/v1/router.py

@@ -4,9 +4,9 @@ import io
 from ipaddress import ip_address, ip_interface, ip_network
 from pathlib import Path
 
-from fastapi import APIRouter, Depends, Header, HTTPException, Request
+from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request
 from fastapi.responses import FileResponse, PlainTextResponse, StreamingResponse
-from sqlalchemy import func, select
+from sqlalchemy import func, or_, select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
@@ -744,29 +744,24 @@ async def policy_deployment_status(db: Session, policy: Policy) -> dict[str, obj
 
 def dashboard_top_talkers(db: Session) -> list[dict[str, int | str]]:
     totals: dict[str, int] = {}
-    workloads = db.scalars(select(Workload)).all()
+    source_rows = db.execute(
-    workloads_by_node_vmid = {(workload.node_id, workload.external_id): workload for workload in workloads}
+        select(Workload.name, func.coalesce(func.sum(TrafficFlow.bytes), 0))
-    ip_owners = {
+        .join(IpAddress, IpAddress.workload_id == Workload.id)
-        address.address: db.get(Workload, address.workload_id)
+        .join(TrafficFlow, TrafficFlow.source_ip == IpAddress.address)
-        for address in db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all()
+        .group_by(Workload.name)
-    }
+    ).all()
-
+    destination_rows = db.execute(
-    flows = db.scalars(select(TrafficFlow)).all()
+        select(Workload.name, func.coalesce(func.sum(TrafficFlow.bytes), 0))
-    for flow in flows:
+        .join(IpAddress, IpAddress.workload_id == Workload.id)
-        raw = flow.raw if isinstance(flow.raw, dict) else {}
+        .join(TrafficFlow, TrafficFlow.destination_ip == IpAddress.address)
-        candidates: list[Workload] = []
+        .group_by(Workload.name)
-        raw_vmid = str(raw.get("vmid") or "")
+    ).all()
-        if raw_vmid:
+    for name, bytes_value in [*source_rows, *destination_rows]:
-            workload = workloads_by_node_vmid.get((flow.node_id, raw_vmid))
+        totals[str(name)] = totals.get(str(name), 0) + int(bytes_value or 0)
-            if workload:
-                candidates.append(workload)
-        for owner in (ip_owners.get(flow.source_ip), ip_owners.get(flow.destination_ip)):
-            if owner and owner not in candidates:
-                candidates.append(owner)
-        for workload in candidates:
-            totals[workload.name] = totals.get(workload.name, 0) + int(flow.bytes or 0)
 
     if not totals:
+        workloads = db.scalars(select(Workload)).all()
+        workloads_by_node_vmid = {(workload.node_id, workload.external_id): workload for workload in workloads}
         agents = db.scalars(select(NodeAgent)).all()
         for agent in agents:
             payload = agent.last_payload if isinstance(agent.last_payload, dict) else {}
@@ -799,8 +794,19 @@ def dashboard_suspicious_traffic(db: Session) -> list[dict[str, int | str]]:
         address.address
         for address in db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all()
     }
+    blocked_states = ["blocked", "drop", "dropped", "reject", "rejected", "deny", "denied"]
+    candidate_flows = db.scalars(
+        select(TrafficFlow)
+        .where(
+            or_(
+                TrafficFlow.destination_port.in_(list(sensitive_ports)),
+                func.lower(func.coalesce(TrafficFlow.state, "")).in_(blocked_states),
+            )
+        )
+        .order_by(TrafficFlow.updated_at.desc())
+    ).all()
     events: dict[tuple[str, str, int], dict[str, int | str]] = {}
-    for flow in db.scalars(select(TrafficFlow).order_by(TrafficFlow.updated_at.desc())).all():
+    for flow in candidate_flows:
         port = flow.destination_port or 0
         source_internal = bool(subnet_label_for_ip(subnets, flow.source_ip))
         destination_internal = bool(subnet_label_for_ip(subnets, flow.destination_ip)) or flow.destination_ip in workload_ips
@@ -1507,7 +1513,12 @@ def workloads(_: CurrentUser, db: Session = Depends(get_db)) -> list[Workload]:
 
 
 @api_router.get("/vms/{workload_id}/insights", response_model=WorkloadInsight)
-async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depends(get_db)) -> WorkloadInsight:
+async def workload_insights(
+    workload_id: str,
+    _: CurrentUser,
+    db: Session = Depends(get_db),
+    traffic: str = Query(default="summary", pattern="^(summary|full)$"),
+) -> WorkloadInsight:
     workload = db.get(Workload, workload_id)
     if not workload:
         raise HTTPException(status_code=404, detail="Workload not found")
@@ -1517,10 +1528,9 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe
     assigned_ips = db.scalars(select(IpAddress).where(IpAddress.workload_id == workload.id).order_by(IpAddress.address)).all()
     workload_ips = [address.address for address in assigned_ips]
     all_assigned_ips = db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all()
-    ip_owners = {
+    owner_ids = {address.workload_id for address in all_assigned_ips if address.workload_id}
-        address.address: db.get(Workload, address.workload_id)
+    owners = {workload.id: workload for workload in db.scalars(select(Workload).where(Workload.id.in_(owner_ids))).all()} if owner_ids else {}
-        for address in all_assigned_ips
+    ip_owners = {address.address: owners.get(address.workload_id) for address in all_assigned_ips if address.workload_id}
-    }
     workload_ips_by_id: dict[str, set[str]] = {}
     for address in all_assigned_ips:
         if address.workload_id:
@@ -1530,11 +1540,14 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe
     traffic = []
     if workload_ips:
         workload_ip_set = set(workload_ips)
-        flows = db.scalars(
+        flow_query = (
             select(TrafficFlow)
             .where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips)))
             .order_by((TrafficFlow.state == "blocked").desc(), TrafficFlow.updated_at.desc())
-        ).all()
+        )
+        if traffic == "summary":
+            flow_query = flow_query.limit(250)
+        flows = db.scalars(flow_query).all()
         for flow in flows:
             source_owner = ip_owners.get(flow.source_ip)
             destination_owner = ip_owners.get(flow.destination_ip)

backend/app/main.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from sqlalchemy import text
 
 from app.api.v1.router import api_router
 from app.core.config import get_settings
@@ -7,6 +8,20 @@ from app.db.session import Base, SessionLocal, engine
 from app.seed.demo import seed_demo_data
 
 
+def ensure_runtime_indexes() -> None:
+    index_statements = [
+        "CREATE INDEX IF NOT EXISTS ix_ip_addresses_address ON ip_addresses (address)",
+        "CREATE INDEX IF NOT EXISTS ix_traffic_flows_node_updated ON traffic_flows (node_id, updated_at)",
+        "CREATE INDEX IF NOT EXISTS ix_traffic_flows_source_updated ON traffic_flows (source_ip, updated_at)",
+        "CREATE INDEX IF NOT EXISTS ix_traffic_flows_destination_updated ON traffic_flows (destination_ip, updated_at)",
+        "CREATE INDEX IF NOT EXISTS ix_traffic_flows_destination_port_updated ON traffic_flows (destination_port, updated_at)",
+        "CREATE INDEX IF NOT EXISTS ix_traffic_flows_state_updated ON traffic_flows (state, updated_at)",
+    ]
+    with engine.begin() as connection:
+        for statement in index_statements:
+            connection.execute(text(statement))
+
+
 def create_app() -> FastAPI:
     settings = get_settings()
     app = FastAPI(
@@ -27,6 +42,7 @@ def create_app() -> FastAPI:
     @app.on_event("startup")
     def startup() -> None:
         Base.metadata.create_all(bind=engine)
+        ensure_runtime_indexes()
         with SessionLocal() as db:
             seed_demo_data(db)
 
@@ -39,4 +55,3 @@ def create_app() -> FastAPI:
 
 
 app = create_app()
-

backend/app/models/domain.py

@@ -2,7 +2,7 @@ from datetime import datetime
 from enum import StrEnum
 from uuid import uuid4
 
-from sqlalchemy import JSON, BigInteger, Boolean, DateTime, Enum, ForeignKey, Integer, String, Text, UniqueConstraint
+from sqlalchemy import JSON, BigInteger, Boolean, DateTime, Enum, ForeignKey, Index, Integer, String, Text, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from app.db.session import Base
@@ -185,7 +185,10 @@ class Subnet(Base, TimestampMixin):
 
 class IpAddress(Base, TimestampMixin):
     __tablename__ = "ip_addresses"
-    __table_args__ = (UniqueConstraint("subnet_id", "address"),)
+    __table_args__ = (
+        UniqueConstraint("subnet_id", "address"),
+        Index("ix_ip_addresses_address", "address"),
+    )
 
     id: Mapped[str] = mapped_column(String, primary_key=True, default=new_id)
     subnet_id: Mapped[str] = mapped_column(ForeignKey("subnets.id"), index=True)
@@ -197,6 +200,13 @@ class IpAddress(Base, TimestampMixin):
 
 class TrafficFlow(Base, TimestampMixin):
     __tablename__ = "traffic_flows"
+    __table_args__ = (
+        Index("ix_traffic_flows_node_updated", "node_id", "updated_at"),
+        Index("ix_traffic_flows_source_updated", "source_ip", "updated_at"),
+        Index("ix_traffic_flows_destination_updated", "destination_ip", "updated_at"),
+        Index("ix_traffic_flows_destination_port_updated", "destination_port", "updated_at"),
+        Index("ix_traffic_flows_state_updated", "state", "updated_at"),
+    )
 
     id: Mapped[str] = mapped_column(String, primary_key=True, default=new_id)
     node_id: Mapped[str] = mapped_column(ForeignKey("nodes.id"), index=True)

frontend/src/pages/Workloads.tsx

@@ -452,8 +452,8 @@ export function Workloads() {
   const [selectedId, setSelectedId] = useState("");
   const selected = selectedId || workloads.data?.[0]?.id || "";
   const insight = useQuery({
-    queryKey: ["workload-insight", selected],
+    queryKey: ["workload-insight", selected, "summary"],
-    queryFn: () => api<WorkloadInsight>(`/vms/${selected}/insights`),
+    queryFn: () => api<WorkloadInsight>(`/vms/${selected}/insights?traffic=summary`),
     enabled: Boolean(selected),
   });
   const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]);
@@ -540,8 +540,8 @@ export function Workloads() {
 export function WorkloadDetail() {
   const { workloadId } = useParams();
   const insight = useQuery({
-    queryKey: ["workload-insight", workloadId],
+    queryKey: ["workload-insight", workloadId, "summary"],
-    queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights`),
+    queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights?traffic=summary`),
     enabled: Boolean(workloadId),
   });
   const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]);
@@ -638,8 +638,8 @@ export function WorkloadFlows() {
   const [port, setPort] = useState("");
   const [page, setPage] = useState(1);
   const insight = useQuery({
-    queryKey: ["workload-insight", workloadId],
+    queryKey: ["workload-insight", workloadId, "full"],
-    queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights`),
+    queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights?traffic=full`),
     enabled: Boolean(workloadId),
   });
   const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]);