lean_ctx/core/property_graph/
edge.rs1use rusqlite::{params, Connection};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum EdgeKind {
7 Imports,
8 Calls,
9 Defines,
10 Exports,
11 TypeRef,
12}
13
14impl EdgeKind {
15 pub fn as_str(&self) -> &'static str {
16 match self {
17 EdgeKind::Imports => "imports",
18 EdgeKind::Calls => "calls",
19 EdgeKind::Defines => "defines",
20 EdgeKind::Exports => "exports",
21 EdgeKind::TypeRef => "type_ref",
22 }
23 }
24
25 pub fn parse(s: &str) -> Self {
26 match s {
27 "calls" => EdgeKind::Calls,
28 "defines" => EdgeKind::Defines,
29 "exports" => EdgeKind::Exports,
30 "type_ref" => EdgeKind::TypeRef,
31 _ => EdgeKind::Imports,
32 }
33 }
34}
35
36#[derive(Debug, Clone)]
37pub struct Edge {
38 pub id: Option<i64>,
39 pub source_id: i64,
40 pub target_id: i64,
41 pub kind: EdgeKind,
42 pub metadata: Option<String>,
43}
44
45impl Edge {
46 pub fn new(source_id: i64, target_id: i64, kind: EdgeKind) -> Self {
47 Self {
48 id: None,
49 source_id,
50 target_id,
51 kind,
52 metadata: None,
53 }
54 }
55
56 pub fn with_metadata(mut self, meta: &str) -> Self {
57 self.metadata = Some(meta.to_string());
58 self
59 }
60}
61
62pub fn upsert(conn: &Connection, edge: &Edge) -> anyhow::Result<()> {
63 conn.execute(
64 "INSERT INTO edges (source_id, target_id, kind, metadata)
65 VALUES (?1, ?2, ?3, ?4)
66 ON CONFLICT(source_id, target_id, kind) DO UPDATE SET
67 metadata = excluded.metadata",
68 params![
69 edge.source_id,
70 edge.target_id,
71 edge.kind.as_str(),
72 edge.metadata,
73 ],
74 )?;
75 Ok(())
76}
77
78pub fn from_node(conn: &Connection, node_id: i64) -> anyhow::Result<Vec<Edge>> {
79 let mut stmt = conn.prepare(
80 "SELECT id, source_id, target_id, kind, metadata
81 FROM edges WHERE source_id = ?1",
82 )?;
83 let edges = stmt
84 .query_map(params![node_id], |row| {
85 Ok(Edge {
86 id: Some(row.get(0)?),
87 source_id: row.get(1)?,
88 target_id: row.get(2)?,
89 kind: EdgeKind::parse(&row.get::<_, String>(3)?),
90 metadata: row.get(4)?,
91 })
92 })?
93 .filter_map(|r| r.ok())
94 .collect();
95 Ok(edges)
96}
97
98pub fn to_node(conn: &Connection, node_id: i64) -> anyhow::Result<Vec<Edge>> {
99 let mut stmt = conn.prepare(
100 "SELECT id, source_id, target_id, kind, metadata
101 FROM edges WHERE target_id = ?1",
102 )?;
103 let edges = stmt
104 .query_map(params![node_id], |row| {
105 Ok(Edge {
106 id: Some(row.get(0)?),
107 source_id: row.get(1)?,
108 target_id: row.get(2)?,
109 kind: EdgeKind::parse(&row.get::<_, String>(3)?),
110 metadata: row.get(4)?,
111 })
112 })?
113 .filter_map(|r| r.ok())
114 .collect();
115 Ok(edges)
116}
117
118pub fn count(conn: &Connection) -> anyhow::Result<usize> {
119 let c: i64 = conn.query_row("SELECT COUNT(*) FROM edges", [], |row| row.get(0))?;
120 Ok(c as usize)
121}