1use indexmap::IndexMap;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
10#[serde(rename_all = "PascalCase")]
11pub enum NodeKind {
12 Lambda,
13 ApiGateway,
14 ApiRoute,
15 EventRule,
16 SqsQueue,
17 SnsTopic,
18 S3Bucket,
19 DynamoStream,
20 StepFunction,
21 LogGroup,
22 EcsService,
23 Ec2Instance,
24 LoadBalancer,
25 Unknown,
26}
27
28impl std::fmt::Display for NodeKind {
29 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30 write!(f, "{:?}", self)
31 }
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
39#[serde(rename_all = "PascalCase")]
40pub enum EdgeKind {
41 Triggers,
42 Invokes,
43 Consumes,
44 Publishes,
45 ReadsFrom,
46 WritesTo,
47}
48
49impl std::fmt::Display for EdgeKind {
50 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51 write!(f, "{:?}", self)
52 }
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
60pub struct Node {
61 pub id: String,
62 pub kind: NodeKind,
63 pub name: String,
64 #[serde(skip_serializing_if = "Option::is_none")]
65 pub arn: Option<String>,
66 #[serde(skip_serializing_if = "Option::is_none")]
67 pub region: Option<String>,
68 #[serde(skip_serializing_if = "Option::is_none")]
69 pub account_id: Option<String>,
70 #[serde(skip_serializing_if = "Option::is_none")]
71 pub tags: Option<IndexMap<String, String>>,
72 #[serde(default = "default_props")]
73 pub props: serde_json::Value,
74}
75
76fn default_props() -> serde_json::Value {
77 serde_json::Value::Object(serde_json::Map::new())
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct Edge {
86 pub from: String,
87 pub to: String,
88 pub kind: EdgeKind,
89 #[serde(default = "default_props")]
90 pub props: serde_json::Value,
91}
92
93#[derive(Debug, Clone, Default, Serialize, Deserialize)]
98pub struct GraphStats {
99 pub node_count: usize,
100 pub edge_count: usize,
101 pub nodes_by_kind: IndexMap<String, usize>,
102 pub edges_by_kind: IndexMap<String, usize>,
103 pub top_fan_in: Vec<(String, usize)>,
104 pub top_fan_out: Vec<(String, usize)>,
105}
106
107#[derive(Debug, Clone, Serialize, Deserialize)]
112pub struct Graph {
113 pub generated_at: String,
114 #[serde(skip_serializing_if = "Option::is_none")]
115 pub profile: Option<String>,
116 pub regions: Vec<String>,
117 pub nodes: Vec<Node>,
118 pub edges: Vec<Edge>,
119 #[serde(default)]
120 pub warnings: Vec<String>,
121 #[serde(default)]
122 pub stats: GraphStats,
123}
124
125impl Default for Graph {
126 fn default() -> Self {
127 Self::new()
128 }
129}
130
131impl Graph {
132 pub fn new() -> Self {
133 Self {
134 generated_at: chrono::Utc::now().to_rfc3339(),
135 profile: None,
136 regions: Vec::new(),
137 nodes: Vec::new(),
138 edges: Vec::new(),
139 warnings: Vec::new(),
140 stats: GraphStats::default(),
141 }
142 }
143
144 pub fn dedupe_and_sort(&mut self) {
147 let mut seen = std::collections::HashSet::new();
149 self.nodes.retain(|n| seen.insert(n.id.clone()));
150
151 let mut seen_edges = std::collections::HashSet::new();
153 self.edges.retain(|e| {
154 seen_edges.insert((e.from.clone(), e.to.clone(), format!("{:?}", e.kind)))
155 });
156
157 self.nodes.sort_by(|a, b| {
159 a.kind
160 .cmp(&b.kind)
161 .then_with(|| a.name.cmp(&b.name))
162 .then_with(|| a.id.cmp(&b.id))
163 });
164
165 self.edges.sort_by(|a, b| {
167 a.kind
168 .cmp(&b.kind)
169 .then_with(|| a.from.cmp(&b.from))
170 .then_with(|| a.to.cmp(&b.to))
171 });
172 }
173
174 pub fn compute_stats(&mut self) {
176 let id_to_name: HashMap<&str, &str> = self
178 .nodes
179 .iter()
180 .map(|n| (n.id.as_str(), n.name.as_str()))
181 .collect();
182
183 let mut nodes_by_kind: IndexMap<String, usize> = IndexMap::new();
184 for node in &self.nodes {
185 *nodes_by_kind.entry(node.kind.to_string()).or_default() += 1;
186 }
187
188 let mut edges_by_kind: IndexMap<String, usize> = IndexMap::new();
189 for edge in &self.edges {
190 *edges_by_kind.entry(edge.kind.to_string()).or_default() += 1;
191 }
192
193 let mut fan_in: HashMap<String, usize> = HashMap::new();
195 for edge in &self.edges {
196 let label = id_to_name
197 .get(edge.to.as_str())
198 .unwrap_or(&edge.to.as_str())
199 .to_string();
200 *fan_in.entry(label).or_default() += 1;
201 }
202 let mut top_fan_in: Vec<(String, usize)> = fan_in.into_iter().collect();
203 top_fan_in.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
204 top_fan_in.truncate(10);
205
206 let mut fan_out: HashMap<String, usize> = HashMap::new();
208 for edge in &self.edges {
209 let label = id_to_name
210 .get(edge.from.as_str())
211 .unwrap_or(&edge.from.as_str())
212 .to_string();
213 *fan_out.entry(label).or_default() += 1;
214 }
215 let mut top_fan_out: Vec<(String, usize)> = fan_out.into_iter().collect();
216 top_fan_out.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
217 top_fan_out.truncate(10);
218
219 self.stats = GraphStats {
220 node_count: self.nodes.len(),
221 edge_count: self.edges.len(),
222 nodes_by_kind,
223 edges_by_kind,
224 top_fan_in,
225 top_fan_out,
226 };
227 }
228
229 pub fn merge(&mut self, output: DiscoveryOutput) {
231 self.nodes.extend(output.nodes);
232 self.edges.extend(output.edges);
233 self.warnings.extend(output.warnings);
234 }
235}
236
237#[derive(Debug, Clone, Default)]
242pub struct DiscoveryOutput {
243 pub nodes: Vec<Node>,
244 pub edges: Vec<Edge>,
245 pub warnings: Vec<String>,
246}
247
248#[derive(Debug, Clone, Serialize, Deserialize)]
253pub struct GraphDiff {
254 pub added_nodes: Vec<String>,
255 pub removed_nodes: Vec<String>,
256 pub added_edges: Vec<String>,
257 pub removed_edges: Vec<String>,
258}
259
260impl GraphDiff {
261 pub fn compute(old: &Graph, new: &Graph) -> Self {
262 let old_node_ids: std::collections::HashSet<&str> =
263 old.nodes.iter().map(|n| n.id.as_str()).collect();
264 let new_node_ids: std::collections::HashSet<&str> =
265 new.nodes.iter().map(|n| n.id.as_str()).collect();
266
267 let added_nodes: Vec<String> = new_node_ids
268 .difference(&old_node_ids)
269 .map(|s| s.to_string())
270 .collect();
271 let removed_nodes: Vec<String> = old_node_ids
272 .difference(&new_node_ids)
273 .map(|s| s.to_string())
274 .collect();
275
276 let edge_key = |e: &Edge| format!("{} --{:?}--> {}", e.from, e.kind, e.to);
277 let old_edge_keys: std::collections::HashSet<String> =
278 old.edges.iter().map(edge_key).collect();
279 let new_edge_keys: std::collections::HashSet<String> =
280 new.edges.iter().map(edge_key).collect();
281
282 let added_edges: Vec<String> = new_edge_keys
283 .difference(&old_edge_keys)
284 .cloned()
285 .collect();
286 let removed_edges: Vec<String> = old_edge_keys
287 .difference(&new_edge_keys)
288 .cloned()
289 .collect();
290
291 GraphDiff {
292 added_nodes,
293 removed_nodes,
294 added_edges,
295 removed_edges,
296 }
297 }
298}
299
300#[cfg(test)]
305mod tests {
306 use super::*;
307
308 #[test]
309 fn test_dedupe_and_sort() {
310 let mut g = Graph::new();
311 g.nodes.push(Node {
312 id: "arn:aws:lambda:us-east-1:123:function:alpha".into(),
313 kind: NodeKind::Lambda,
314 name: "alpha".into(),
315 arn: Some("arn:aws:lambda:us-east-1:123:function:alpha".into()),
316 region: Some("us-east-1".into()),
317 account_id: None,
318 tags: None,
319 props: serde_json::json!({}),
320 });
321 g.nodes.push(Node {
323 id: "arn:aws:lambda:us-east-1:123:function:alpha".into(),
324 kind: NodeKind::Lambda,
325 name: "alpha".into(),
326 arn: Some("arn:aws:lambda:us-east-1:123:function:alpha".into()),
327 region: Some("us-east-1".into()),
328 account_id: None,
329 tags: None,
330 props: serde_json::json!({}),
331 });
332 g.nodes.push(Node {
333 id: "arn:aws:lambda:us-east-1:123:function:beta".into(),
334 kind: NodeKind::Lambda,
335 name: "beta".into(),
336 arn: None,
337 region: None,
338 account_id: None,
339 tags: None,
340 props: serde_json::json!({}),
341 });
342
343 g.edges.push(Edge {
344 from: "a".into(),
345 to: "b".into(),
346 kind: EdgeKind::Triggers,
347 props: serde_json::json!({}),
348 });
349 g.edges.push(Edge {
351 from: "a".into(),
352 to: "b".into(),
353 kind: EdgeKind::Triggers,
354 props: serde_json::json!({}),
355 });
356
357 g.dedupe_and_sort();
358 assert_eq!(g.nodes.len(), 2);
359 assert_eq!(g.edges.len(), 1);
360 assert_eq!(g.nodes[0].name, "alpha");
362 assert_eq!(g.nodes[1].name, "beta");
363 }
364
365 #[test]
366 fn test_compute_stats() {
367 let mut g = Graph::new();
368 g.nodes.push(Node {
369 id: "l1".into(),
370 kind: NodeKind::Lambda,
371 name: "fn1".into(),
372 arn: None,
373 region: None,
374 account_id: None,
375 tags: None,
376 props: serde_json::json!({}),
377 });
378 g.nodes.push(Node {
379 id: "l2".into(),
380 kind: NodeKind::Lambda,
381 name: "fn2".into(),
382 arn: None,
383 region: None,
384 account_id: None,
385 tags: None,
386 props: serde_json::json!({}),
387 });
388 g.nodes.push(Node {
389 id: "q1".into(),
390 kind: NodeKind::SqsQueue,
391 name: "queue1".into(),
392 arn: None,
393 region: None,
394 account_id: None,
395 tags: None,
396 props: serde_json::json!({}),
397 });
398 g.edges.push(Edge {
399 from: "q1".into(),
400 to: "l1".into(),
401 kind: EdgeKind::Triggers,
402 props: serde_json::json!({}),
403 });
404 g.edges.push(Edge {
405 from: "q1".into(),
406 to: "l2".into(),
407 kind: EdgeKind::Triggers,
408 props: serde_json::json!({}),
409 });
410 g.compute_stats();
411 assert_eq!(g.stats.node_count, 3);
412 assert_eq!(g.stats.edge_count, 2);
413 assert_eq!(g.stats.nodes_by_kind["Lambda"], 2);
414 assert_eq!(g.stats.top_fan_out[0], ("queue1".to_string(), 2));
415 }
416
417 #[test]
418 fn test_diff() {
419 let mut old = Graph::new();
420 old.nodes.push(Node {
421 id: "a".into(),
422 kind: NodeKind::Lambda,
423 name: "a".into(),
424 arn: None,
425 region: None,
426 account_id: None,
427 tags: None,
428 props: serde_json::json!({}),
429 });
430
431 let mut new = Graph::new();
432 new.nodes.push(Node {
433 id: "b".into(),
434 kind: NodeKind::Lambda,
435 name: "b".into(),
436 arn: None,
437 region: None,
438 account_id: None,
439 tags: None,
440 props: serde_json::json!({}),
441 });
442
443 let diff = GraphDiff::compute(&old, &new);
444 assert_eq!(diff.added_nodes, vec!["b".to_string()]);
445 assert_eq!(diff.removed_nodes, vec!["a".to_string()]);
446 }
447
448 #[test]
449 fn test_graph_serialization_roundtrip() {
450 let mut g = Graph::new();
451 g.nodes.push(Node {
452 id: "test".into(),
453 kind: NodeKind::Lambda,
454 name: "test-fn".into(),
455 arn: Some("arn:aws:lambda:us-east-1:123:function:test-fn".into()),
456 region: Some("us-east-1".into()),
457 account_id: Some("123".into()),
458 tags: None,
459 props: serde_json::json!({"runtime": "nodejs18.x"}),
460 });
461 g.compute_stats();
462
463 let json = serde_json::to_string_pretty(&g).unwrap();
464 let parsed: Graph = serde_json::from_str(&json).unwrap();
465 assert_eq!(parsed.nodes.len(), 1);
466 assert_eq!(parsed.nodes[0].name, "test-fn");
467 }
468}