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
This commit is contained in:
2026-07-10 08:23:38 +02:00
parent 9af60945cf
commit 65bc7fad67
4 changed files with 78 additions and 40 deletions
+44 -31
View File
@@ -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()
workloads_by_node_vmid = {(workload.node_id, workload.external_id): workload for workload in workloads}
ip_owners = {
address.address: db.get(Workload, address.workload_id)
for address in db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all()
}
flows = db.scalars(select(TrafficFlow)).all()
for flow in flows:
raw = flow.raw if isinstance(flow.raw, dict) else {}
candidates: list[Workload] = []
raw_vmid = str(raw.get("vmid") or "")
if raw_vmid:
workload = workloads_by_node_vmid.get((flow.node_id, raw_vmid))
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)
source_rows = db.execute(
select(Workload.name, func.coalesce(func.sum(TrafficFlow.bytes), 0))
.join(IpAddress, IpAddress.workload_id == Workload.id)
.join(TrafficFlow, TrafficFlow.source_ip == IpAddress.address)
.group_by(Workload.name)
).all()
destination_rows = db.execute(
select(Workload.name, func.coalesce(func.sum(TrafficFlow.bytes), 0))
.join(IpAddress, IpAddress.workload_id == Workload.id)
.join(TrafficFlow, TrafficFlow.destination_ip == IpAddress.address)
.group_by(Workload.name)
).all()
for name, bytes_value in [*source_rows, *destination_rows]:
totals[str(name)] = totals.get(str(name), 0) + int(bytes_value 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 = {
address.address: db.get(Workload, address.workload_id)
for address in all_assigned_ips
}
owner_ids = {address.workload_id for address in all_assigned_ips if 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 {}
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)