1use rusqlite::{params, Connection, OptionalExtension};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum NodeKind {
7 File,
8 Symbol,
9 Module,
10 Commit,
11 Test,
12 CIRun,
13 Knowledge,
14 Issue,
15}
16
17impl NodeKind {
18 pub fn as_str(&self) -> &'static str {
19 match self {
20 Self::File => "file",
21 Self::Symbol => "symbol",
22 Self::Module => "module",
23 Self::Commit => "commit",
24 Self::Test => "test",
25 Self::CIRun => "ci_run",
26 Self::Knowledge => "knowledge",
27 Self::Issue => "issue",
28 }
29 }
30
31 pub fn parse(s: &str) -> Self {
32 match s {
33 "symbol" => Self::Symbol,
34 "module" => Self::Module,
35 "commit" => Self::Commit,
36 "test" => Self::Test,
37 "ci_run" => Self::CIRun,
38 "knowledge" => Self::Knowledge,
39 "issue" => Self::Issue,
40 _ => Self::File,
41 }
42 }
43}
44
45#[derive(Debug, Clone)]
46pub struct Node {
47 pub id: Option<i64>,
48 pub kind: NodeKind,
49 pub name: String,
50 pub file_path: String,
51 pub line_start: Option<usize>,
52 pub line_end: Option<usize>,
53 pub metadata: Option<String>,
54}
55
56impl Node {
57 pub fn file(path: &str) -> Self {
58 Self {
59 id: None,
60 kind: NodeKind::File,
61 name: path.to_string(),
62 file_path: path.to_string(),
63 line_start: None,
64 line_end: None,
65 metadata: None,
66 }
67 }
68
69 pub fn symbol(name: &str, file_path: &str, kind: NodeKind) -> Self {
70 Self {
71 id: None,
72 kind,
73 name: name.to_string(),
74 file_path: file_path.to_string(),
75 line_start: None,
76 line_end: None,
77 metadata: None,
78 }
79 }
80
81 pub fn with_lines(mut self, start: usize, end: usize) -> Self {
82 self.line_start = Some(start);
83 self.line_end = Some(end);
84 self
85 }
86
87 pub fn with_metadata(mut self, meta: &str) -> Self {
88 self.metadata = Some(meta.to_string());
89 self
90 }
91
92 pub fn commit(hash: &str, message: &str) -> Self {
93 Self {
94 id: None,
95 kind: NodeKind::Commit,
96 name: hash.to_string(),
97 file_path: String::new(),
98 line_start: None,
99 line_end: None,
100 metadata: Some(message.to_string()),
101 }
102 }
103
104 pub fn test(path: &str, test_name: &str) -> Self {
105 Self {
106 id: None,
107 kind: NodeKind::Test,
108 name: test_name.to_string(),
109 file_path: path.to_string(),
110 line_start: None,
111 line_end: None,
112 metadata: None,
113 }
114 }
115
116 pub fn knowledge(id: &str, summary: &str) -> Self {
117 Self {
118 id: None,
119 kind: NodeKind::Knowledge,
120 name: id.to_string(),
121 file_path: String::new(),
122 line_start: None,
123 line_end: None,
124 metadata: Some(summary.to_string()),
125 }
126 }
127
128 pub fn issue(id: &str, title: &str) -> Self {
129 Self {
130 id: None,
131 kind: NodeKind::Issue,
132 name: id.to_string(),
133 file_path: String::new(),
134 line_start: None,
135 line_end: None,
136 metadata: Some(title.to_string()),
137 }
138 }
139}
140
141pub(super) fn upsert(conn: &Connection, node: &Node) -> anyhow::Result<i64> {
142 conn.execute(
143 "INSERT INTO nodes (kind, name, file_path, line_start, line_end, metadata)
144 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
145 ON CONFLICT(kind, name, file_path) DO UPDATE SET
146 line_start = excluded.line_start,
147 line_end = excluded.line_end,
148 metadata = excluded.metadata",
149 params![
150 node.kind.as_str(),
151 node.name,
152 node.file_path,
153 node.line_start.map(|v| v as i64),
154 node.line_end.map(|v| v as i64),
155 node.metadata,
156 ],
157 )?;
158
159 let id: i64 = conn.query_row(
160 "SELECT id FROM nodes WHERE kind = ?1 AND name = ?2 AND file_path = ?3",
161 params![node.kind.as_str(), node.name, node.file_path],
162 |row| row.get(0),
163 )?;
164
165 Ok(id)
166}
167
168pub(super) fn get_by_path(conn: &Connection, file_path: &str) -> anyhow::Result<Option<Node>> {
169 let result = conn
170 .query_row(
171 "SELECT id, kind, name, file_path, line_start, line_end, metadata
172 FROM nodes WHERE kind = 'file' AND file_path = ?1",
173 params![file_path],
174 |row| {
175 Ok(Node {
176 id: Some(row.get(0)?),
177 kind: NodeKind::parse(&row.get::<_, String>(1)?),
178 name: row.get(2)?,
179 file_path: row.get(3)?,
180 line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
181 line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
182 metadata: row.get(6)?,
183 })
184 },
185 )
186 .optional()?;
187 Ok(result)
188}
189
190pub(super) fn get_by_symbol(
191 conn: &Connection,
192 name: &str,
193 file_path: &str,
194) -> anyhow::Result<Option<Node>> {
195 let result = conn
196 .query_row(
197 "SELECT id, kind, name, file_path, line_start, line_end, metadata
198 FROM nodes WHERE name = ?1 AND file_path = ?2 AND kind != 'file'",
199 params![name, file_path],
200 |row| {
201 Ok(Node {
202 id: Some(row.get(0)?),
203 kind: NodeKind::parse(&row.get::<_, String>(1)?),
204 name: row.get(2)?,
205 file_path: row.get(3)?,
206 line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
207 line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
208 metadata: row.get(6)?,
209 })
210 },
211 )
212 .optional()?;
213 Ok(result)
214}
215
216pub(super) fn remove_by_file(conn: &Connection, file_path: &str) -> anyhow::Result<()> {
217 conn.execute(
218 "DELETE FROM edges WHERE source_id IN (SELECT id FROM nodes WHERE file_path = ?1)
219 OR target_id IN (SELECT id FROM nodes WHERE file_path = ?1)",
220 params![file_path],
221 )?;
222 conn.execute("DELETE FROM nodes WHERE file_path = ?1", params![file_path])?;
223 Ok(())
224}
225
226pub(super) fn find_symbols(
227 conn: &Connection,
228 name: &str,
229 file_filter: Option<&str>,
230 kind_filter: Option<&str>,
231) -> anyhow::Result<Vec<Node>> {
232 let name_lower = name.to_lowercase();
233 let mut sql = String::from(
234 "SELECT id, kind, name, file_path, line_start, line_end, metadata
235 FROM nodes WHERE kind != 'file'
236 AND LOWER(name) LIKE '%' || ?1 || '%'",
237 );
238 let mut param_idx = 2;
239 if file_filter.is_some() {
240 sql.push_str(&format!(" AND file_path LIKE '%' || ?{param_idx} || '%'"));
241 param_idx += 1;
242 }
243 if kind_filter.is_some() {
244 sql.push_str(&format!(" AND kind = ?{param_idx}"));
245 }
246 sql.push_str(" ORDER BY file_path, line_start LIMIT 100");
247
248 let mut stmt = conn.prepare(&sql)?;
249
250 let params_vec: Vec<Box<dyn rusqlite::types::ToSql>> = {
251 let mut v: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(name_lower)];
252 if let Some(f) = file_filter {
253 v.push(Box::new(f.to_string()));
254 }
255 if let Some(k) = kind_filter {
256 v.push(Box::new(k.to_string()));
257 }
258 v
259 };
260 let refs: Vec<&dyn rusqlite::types::ToSql> =
261 params_vec.iter().map(std::convert::AsRef::as_ref).collect();
262
263 let rows = stmt.query_map(refs.as_slice(), |row| {
264 Ok(Node {
265 id: Some(row.get(0)?),
266 kind: NodeKind::parse(&row.get::<_, String>(1)?),
267 name: row.get(2)?,
268 file_path: row.get(3)?,
269 line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
270 line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
271 metadata: row.get(6)?,
272 })
273 })?;
274
275 let mut results = Vec::new();
276 for r in rows {
277 results.push(r?);
278 }
279 Ok(results)
280}
281
282pub(super) fn symbol_count(conn: &Connection) -> anyhow::Result<usize> {
283 let c: i64 = conn.query_row(
284 "SELECT COUNT(*) FROM nodes WHERE kind != 'file'",
285 [],
286 |row| row.get(0),
287 )?;
288 Ok(c as usize)
289}
290
291pub(super) fn all_edges_flat(
292 conn: &Connection,
293) -> anyhow::Result<Vec<(String, String, String, f64)>> {
294 let mut stmt = conn.prepare(
295 "SELECT n1.file_path, n2.file_path, e.kind, e.weight
296 FROM edges e
297 JOIN nodes n1 ON e.source_id = n1.id
298 JOIN nodes n2 ON e.target_id = n2.id
299 WHERE n1.kind = 'file' AND n2.kind = 'file'",
300 )?;
301 let rows = stmt.query_map([], |row| {
302 Ok((
303 row.get::<_, String>(0)?,
304 row.get::<_, String>(1)?,
305 row.get::<_, String>(2)?,
306 row.get::<_, f64>(3)?,
307 ))
308 })?;
309 let mut result = Vec::new();
310 for r in rows {
311 result.push(r?);
312 }
313 Ok(result)
314}
315
316pub(super) fn count(conn: &Connection) -> anyhow::Result<usize> {
317 let c: i64 = conn.query_row("SELECT COUNT(*) FROM nodes", [], |row| row.get(0))?;
318 Ok(c as usize)
319}