lean_ctx/core/property_graph/
node.rs1use rusqlite::{params, Connection, OptionalExtension};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum NodeKind {
7 File,
8 Symbol,
9 Module,
10}
11
12impl NodeKind {
13 pub fn as_str(&self) -> &'static str {
14 match self {
15 NodeKind::File => "file",
16 NodeKind::Symbol => "symbol",
17 NodeKind::Module => "module",
18 }
19 }
20
21 pub fn parse(s: &str) -> Self {
22 match s {
23 "symbol" => NodeKind::Symbol,
24 "module" => NodeKind::Module,
25 _ => NodeKind::File,
26 }
27 }
28}
29
30#[derive(Debug, Clone)]
31pub struct Node {
32 pub id: Option<i64>,
33 pub kind: NodeKind,
34 pub name: String,
35 pub file_path: String,
36 pub line_start: Option<usize>,
37 pub line_end: Option<usize>,
38 pub metadata: Option<String>,
39}
40
41impl Node {
42 pub fn file(path: &str) -> Self {
43 Self {
44 id: None,
45 kind: NodeKind::File,
46 name: path.to_string(),
47 file_path: path.to_string(),
48 line_start: None,
49 line_end: None,
50 metadata: None,
51 }
52 }
53
54 pub fn symbol(name: &str, file_path: &str, kind: NodeKind) -> Self {
55 Self {
56 id: None,
57 kind,
58 name: name.to_string(),
59 file_path: file_path.to_string(),
60 line_start: None,
61 line_end: None,
62 metadata: None,
63 }
64 }
65
66 pub fn with_lines(mut self, start: usize, end: usize) -> Self {
67 self.line_start = Some(start);
68 self.line_end = Some(end);
69 self
70 }
71
72 pub fn with_metadata(mut self, meta: &str) -> Self {
73 self.metadata = Some(meta.to_string());
74 self
75 }
76}
77
78pub fn upsert(conn: &Connection, node: &Node) -> anyhow::Result<i64> {
79 conn.execute(
80 "INSERT INTO nodes (kind, name, file_path, line_start, line_end, metadata)
81 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
82 ON CONFLICT(kind, name, file_path) DO UPDATE SET
83 line_start = excluded.line_start,
84 line_end = excluded.line_end,
85 metadata = excluded.metadata",
86 params![
87 node.kind.as_str(),
88 node.name,
89 node.file_path,
90 node.line_start.map(|v| v as i64),
91 node.line_end.map(|v| v as i64),
92 node.metadata,
93 ],
94 )?;
95
96 let id: i64 = conn.query_row(
97 "SELECT id FROM nodes WHERE kind = ?1 AND name = ?2 AND file_path = ?3",
98 params![node.kind.as_str(), node.name, node.file_path],
99 |row| row.get(0),
100 )?;
101
102 Ok(id)
103}
104
105pub fn get_by_path(conn: &Connection, file_path: &str) -> anyhow::Result<Option<Node>> {
106 let result = conn
107 .query_row(
108 "SELECT id, kind, name, file_path, line_start, line_end, metadata
109 FROM nodes WHERE kind = 'file' AND file_path = ?1",
110 params![file_path],
111 |row| {
112 Ok(Node {
113 id: Some(row.get(0)?),
114 kind: NodeKind::parse(&row.get::<_, String>(1)?),
115 name: row.get(2)?,
116 file_path: row.get(3)?,
117 line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
118 line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
119 metadata: row.get(6)?,
120 })
121 },
122 )
123 .optional()?;
124 Ok(result)
125}
126
127pub fn get_by_symbol(
128 conn: &Connection,
129 name: &str,
130 file_path: &str,
131) -> anyhow::Result<Option<Node>> {
132 let result = conn
133 .query_row(
134 "SELECT id, kind, name, file_path, line_start, line_end, metadata
135 FROM nodes WHERE name = ?1 AND file_path = ?2 AND kind != 'file'",
136 params![name, file_path],
137 |row| {
138 Ok(Node {
139 id: Some(row.get(0)?),
140 kind: NodeKind::parse(&row.get::<_, String>(1)?),
141 name: row.get(2)?,
142 file_path: row.get(3)?,
143 line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
144 line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
145 metadata: row.get(6)?,
146 })
147 },
148 )
149 .optional()?;
150 Ok(result)
151}
152
153pub fn remove_by_file(conn: &Connection, file_path: &str) -> anyhow::Result<()> {
154 conn.execute(
155 "DELETE FROM edges WHERE source_id IN (SELECT id FROM nodes WHERE file_path = ?1)
156 OR target_id IN (SELECT id FROM nodes WHERE file_path = ?1)",
157 params![file_path],
158 )?;
159 conn.execute("DELETE FROM nodes WHERE file_path = ?1", params![file_path])?;
160 Ok(())
161}
162
163pub fn count(conn: &Connection) -> anyhow::Result<usize> {
164 let c: i64 = conn.query_row("SELECT COUNT(*) FROM nodes", [], |row| row.get(0))?;
165 Ok(c as usize)
166}