"""BFS query service for the relationship graph. The frontend force-directed graph requests a sub-graph centred on a given node with a max traversal depth. We do an iterative BFS over ``rel_relationships`` (live edges only — soft-deleted excluded) and return both nodes and edges in a shape ECharts can consume directly. """ from __future__ import annotations from collections import deque from dataclasses import dataclass, field from typing import Any from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.domains.relationships.constants import ALL_NODE_TYPES from app.domains.relationships.models import RelRelationship @dataclass class GraphNode: type: str id: str label: str attributes: dict[str, Any] = field(default_factory=dict) @dataclass class GraphEdge: from_type: str from_id: str to_type: str to_id: str rel_type: str weight: float = 1.0 attributes: dict[str, Any] = field(default_factory=dict) @dataclass class Subgraph: nodes: list[GraphNode] edges: list[GraphEdge] center: tuple[str, str] depth: int truncated: bool # Per-node fan-out cap to prevent a "popular" node (e.g. an IT contact # for 50 software entries) from blowing up the response. DEFAULT_PER_NODE_FANOUT_CAP = 30 DEFAULT_MAX_NODES = 500 class RelationshipQueryService: def __init__(self, db: AsyncSession) -> None: self.db = db async def get_subgraph( self, *, center_type: str, center_id: str, depth: int = 2, rel_types: list[str] | None = None, node_types: list[str] | None = None, per_node_fanout_cap: int = DEFAULT_PER_NODE_FANOUT_CAP, max_nodes: int = DEFAULT_MAX_NODES, ) -> Subgraph: """BFS centred on ``(center_type, center_id)``. Args: depth: max number of hops from the center. rel_types: limit to these edge labels (None = all). node_types: limit to these node types (None = all). Even if a node type is excluded, the BFS still traverses through it but doesn't return its neighbours. per_node_fanout_cap: max edges per node to traverse. Prevents stars from dominating the graph. max_nodes: hard ceiling on returned nodes. The traversal stops cleanly when reached and ``truncated=True`` is set. """ if center_type not in ALL_NODE_TYPES: raise ValueError(f"Unknown center node type {center_type!r}") if depth < 0 or depth > 5: raise ValueError(f"Depth must be in [0, 5], got {depth}") rel_filter = set(rel_types) if rel_types else None node_filter = set(node_types) if node_types else None # Track visited nodes by (type,id). Nodes/edges accumulators # preserve insertion order via dicts. nodes: dict[tuple[str, str], GraphNode] = {} edges: dict[ tuple[str, str, str, str, str], GraphEdge ] = {} center_key = (center_type, center_id) center_label = await self._lookup_label(center_type, center_id) or center_id nodes[center_key] = GraphNode( type=center_type, id=center_id, label=center_label, ) if depth == 0: return Subgraph( nodes=list(nodes.values()), edges=[], center=center_key, depth=depth, truncated=False, ) truncated = False frontier: deque[tuple[tuple[str, str], int]] = deque( [(center_key, 0)] ) while frontier: (curr_type, curr_id), curr_depth = frontier.popleft() if curr_depth >= depth: continue # Outgoing out_query = select(RelRelationship).where( RelRelationship.deleted_at.is_(None), RelRelationship.from_type == curr_type, RelRelationship.from_id == curr_id, ) in_query = select(RelRelationship).where( RelRelationship.deleted_at.is_(None), RelRelationship.to_type == curr_type, RelRelationship.to_id == curr_id, ) if rel_filter is not None: out_query = out_query.where( RelRelationship.rel_type.in_(list(rel_filter)) ) in_query = in_query.where( RelRelationship.rel_type.in_(list(rel_filter)) ) # Deterministic fan-out: when a high-degree node exceeds the # cap, pick the strongest / most-recent edges (and break ties # on id) so the rendered subgraph is stable across calls and # replicas — a bare LIMIT with no ORDER BY leaves the choice # to engine row order. _fanout_order = ( RelRelationship.weight.desc(), RelRelationship.last_seen_at.desc(), RelRelationship.id, ) out_query = out_query.order_by(*_fanout_order).limit(per_node_fanout_cap) in_query = in_query.order_by(*_fanout_order).limit(per_node_fanout_cap) out_rows = list( (await self.db.execute(out_query)).scalars().all() ) in_rows = list( (await self.db.execute(in_query)).scalars().all() ) for row in out_rows + in_rows: if len(nodes) >= max_nodes: truncated = True break neighbour_type, neighbour_id, neighbour_label = ( (row.to_type, row.to_id, row.to_label) if (row.from_type, row.from_id) == (curr_type, curr_id) else (row.from_type, row.from_id, row.from_label) ) if node_filter is not None and neighbour_type not in node_filter: # Don't include this node, but the edge to it is # also dropped to keep the graph well-formed. continue edge_key = ( row.from_type, row.from_id, row.to_type, row.to_id, row.rel_type, ) if edge_key not in edges: edges[edge_key] = GraphEdge( from_type=row.from_type, from_id=row.from_id, to_type=row.to_type, to_id=row.to_id, rel_type=row.rel_type, weight=row.weight, attributes=row.attributes or {}, ) neigh_key = (neighbour_type, neighbour_id) if neigh_key not in nodes: nodes[neigh_key] = GraphNode( type=neighbour_type, id=neighbour_id, label=neighbour_label or neighbour_id, ) if curr_depth + 1 < depth: frontier.append((neigh_key, curr_depth + 1)) if truncated: break await self._annotate_person_status(nodes) return Subgraph( nodes=list(nodes.values()), edges=list(edges.values()), center=center_key, depth=depth, truncated=truncated, ) async def _annotate_person_status( self, nodes: dict[tuple[str, str], GraphNode], ) -> None: """Stamp ``employment_status`` onto every person node's attributes so the frontend can mark a departed owner (``left``) red rather than rendering them indistinguishably from active staff. We mark, not hide — a leaver still owning a scope/domain is a governance signal worth surfacing on the graph. """ # ids are already strings — pass them straight to the IN clause. # Binding ``uuid.UUID`` objects against the String(36) PK silently # returns 0 rows (asyncmy), so keep the string form here. person_ids = [nid for (ntype, nid) in nodes if ntype == "person"] if not person_ids: return from app.identity.person_models import IdentityPerson # noqa: PLC0415 rows = ( await self.db.execute( select( IdentityPerson.id, IdentityPerson.employment_status, ).where(IdentityPerson.id.in_(person_ids)) ) ).all() status_by_id = {str(pid): status for pid, status in rows} for (ntype, nid), node in nodes.items(): if ntype == "person" and nid in status_by_id: node.attributes = { **node.attributes, "employment_status": status_by_id[nid], } async def _lookup_label(self, node_type: str, node_id: str) -> str | None: """Best-effort lookup of a human label for the centre node when it isn't yet recorded in any edge row (rare — happens for the very first /me load before any sync ran).""" import uuid from app.domains.cmdb.models import CmdbCloudScope, CmdbDomain, CmdbServer from app.identity.person_models import IdentityPerson try: uid = uuid.UUID(node_id) except (ValueError, TypeError): return None if node_type == "person": p = await self.db.get(IdentityPerson, uid) if p is not None: return p.display_name or p.work_email elif node_type == "cloud_scope": s = await self.db.get(CmdbCloudScope, uid) if s is not None: return s.display_name or str(s.id) elif node_type == "cloud_server": srv = await self.db.get(CmdbServer, uid) if srv is not None: return srv.name or str(srv.id) elif node_type == "domain": dom = await self.db.get(CmdbDomain, uid) if dom is not None: return dom.domain_name or str(dom.id) return None __all__ = ["RelationshipQueryService", "Subgraph", "GraphNode", "GraphEdge"]