Skip to main content

infigraph_core/graph/
store.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use anyhow::Result;
5use kuzu::{Connection, Database, SystemConfig};
6
7use super::schema::{CREATE_SCHEMA, MIGRATIONS};
8use super::store_util::escape;
9
10/// Persistent graph store backed by Kuzu.
11pub struct GraphStore {
12    db: Database,
13}
14
15impl GraphStore {
16    /// Open or create a Kuzu database at the given path.
17    pub fn open(path: &Path) -> Result<Self> {
18        if let Some(parent) = path.parent() {
19            std::fs::create_dir_all(parent)?;
20        }
21        let db = Database::new(path, SystemConfig::default())
22            .map_err(|e| anyhow::anyhow!("failed to open kuzu db: {e}"))?;
23        let store = Self { db };
24        store.init_schema()?;
25        Ok(store)
26    }
27
28    fn init_schema(&self) -> Result<()> {
29        let conn = self.connection()?;
30        for ddl in CREATE_SCHEMA {
31            conn.query(ddl)
32                .map_err(|e| anyhow::anyhow!("schema error: {e}\n  DDL: {ddl}"))?;
33        }
34        for migration in MIGRATIONS {
35            let _ = conn.query(migration);
36        }
37        Ok(())
38    }
39
40    pub fn connection(&self) -> Result<Connection<'_>> {
41        Connection::new(&self.db).map_err(|e| anyhow::anyhow!("failed to create connection: {e}"))
42    }
43
44    /// Remove all graph data for a deleted file.
45    pub fn remove_file(&self, file: &str) -> Result<()> {
46        let conn = self.connection()?;
47        let _ = conn.query(&format!(
48            "MATCH (f:File)-[:DEFINES]->(s:Symbol)-[:HAS_STATEMENT]->(st:Statement) WHERE f.id = '{}' DETACH DELETE st",
49            escape(file)
50        ));
51        let _ = conn.query(&format!(
52            "MATCH (s:Symbol) WHERE s.file = '{}' DETACH DELETE s",
53            escape(file)
54        ));
55        let _ = conn.query(&format!(
56            "MATCH (m:Module) WHERE m.file = '{}' DETACH DELETE m",
57            escape(file)
58        ));
59        let _ = conn.query(&format!(
60            "MATCH (f:File) WHERE f.id = '{}' DETACH DELETE f",
61            escape(file)
62        ));
63        Ok(())
64    }
65
66    /// Return map of file path -> content_hash for all indexed modules.
67    /// Used by incremental indexing to skip unchanged files.
68    pub fn get_file_hashes(&self) -> Result<HashMap<String, String>> {
69        let conn = self.connection()?;
70        let result = conn
71            .query("MATCH (m:Module) RETURN m.file, m.content_hash")
72            .map_err(|e| anyhow::anyhow!("get_file_hashes failed: {e}"))?;
73        let mut map = HashMap::new();
74        for row in result {
75            if row.len() >= 2 {
76                map.insert(row[0].to_string(), row[1].to_string());
77            }
78        }
79        Ok(map)
80    }
81
82    /// Return all symbols as (name, id, file, kind) tuples -- used by resolve_calls.
83    pub fn get_all_symbols(&self) -> Result<Vec<(String, String, String, String)>> {
84        let conn = self.connection()?;
85        let result = conn
86            .query("MATCH (s:Symbol) RETURN s.name, s.id, s.file, s.kind")
87            .map_err(|e| anyhow::anyhow!("get_all_symbols failed: {e}"))?;
88        let mut symbols = Vec::new();
89        for row in result {
90            if row.len() >= 4 {
91                symbols.push((
92                    row[0].to_string(),
93                    row[1].to_string(),
94                    row[2].to_string(),
95                    row[3].to_string(),
96                ));
97            }
98        }
99        Ok(symbols)
100    }
101
102    /// Get total counts for stats.
103    pub fn derive_tested_by_edges(&self) -> Result<usize> {
104        let conn = self.connection()?;
105        let q = super::queries::GraphQuery::new(&conn);
106        q.derive_tested_by_edges()
107    }
108
109    pub fn stats(&self) -> Result<GraphStats> {
110        let conn = self.connection()?;
111
112        let symbol_count = count_query(&conn, "MATCH (s:Symbol) RETURN count(s)")?;
113        let module_count = count_query(&conn, "MATCH (m:Module) RETURN count(m)")?;
114        let file_count = count_query(&conn, "MATCH (f:File) RETURN count(f)")?;
115        let folder_count = count_query(&conn, "MATCH (d:Folder) RETURN count(d)")?;
116        let calls_count = count_query(&conn, "MATCH ()-[r:CALLS]->() RETURN count(r)")?;
117        let inherits_count = count_query(&conn, "MATCH ()-[r:INHERITS]->() RETURN count(r)")?;
118        let contains_count = count_query(&conn, "MATCH ()-[r:CONTAINS]->() RETURN count(r)")?;
119
120        Ok(GraphStats {
121            symbols: symbol_count,
122            modules: module_count,
123            files: file_count,
124            folders: folder_count,
125            calls: calls_count,
126            inherits: inherits_count,
127            contains: contains_count,
128        })
129    }
130}
131
132#[derive(Debug)]
133pub struct GraphStats {
134    pub symbols: u64,
135    pub modules: u64,
136    pub files: u64,
137    pub folders: u64,
138    pub calls: u64,
139    pub inherits: u64,
140    pub contains: u64,
141}
142
143impl std::fmt::Display for GraphStats {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        writeln!(f, "Graph Statistics:")?;
146        writeln!(f, "  Symbols:      {}", self.symbols)?;
147        writeln!(f, "  Modules:      {}", self.modules)?;
148        writeln!(f, "  Files:        {}", self.files)?;
149        writeln!(f, "  Folders:      {}", self.folders)?;
150        writeln!(f, "  Calls edges:  {}", self.calls)?;
151        writeln!(f, "  Inherits:     {}", self.inherits)?;
152        writeln!(f, "  Contains:     {}", self.contains)
153    }
154}
155
156fn count_query(conn: &Connection, query: &str) -> Result<u64> {
157    let mut result = conn
158        .query(query)
159        .map_err(|e| anyhow::anyhow!("query failed: {e}"))?;
160    if let Some(row) = result.next() {
161        if let Some(val) = row.first() {
162            return Ok(val.to_string().parse().unwrap_or(0));
163        }
164    }
165    Ok(0)
166}