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)
|
||||
|
||||
+16
-1
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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],
|
||||
queryFn: () => api<WorkloadInsight>(`/vms/${selected}/insights`),
|
||||
queryKey: ["workload-insight", selected, "summary"],
|
||||
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],
|
||||
queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights`),
|
||||
queryKey: ["workload-insight", workloadId, "summary"],
|
||||
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],
|
||||
queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights`),
|
||||
queryKey: ["workload-insight", workloadId, "full"],
|
||||
queryFn: () => api<WorkloadInsight>(`/vms/${workloadId}/insights?traffic=full`),
|
||||
enabled: Boolean(workloadId),
|
||||
});
|
||||
const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]);
|
||||
|
||||
Reference in New Issue
Block a user