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)
+16 -1
View File
@@ -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()
+12 -2
View File
@@ -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)