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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user