keel_core/
sqlite.rs

1use rusqlite::{params, Connection, Result as SqlResult};
2
3use crate::types::{ExternalEndpoint, GraphError, GraphNode, NodeKind};
4
5const SCHEMA_VERSION: u32 = 3;
6
7/// SQLite-backed implementation of the GraphStore trait.
8pub struct SqliteGraphStore {
9    pub(crate) conn: Connection,
10}
11
12impl SqliteGraphStore {
13    /// Open or create a graph database at the given path.
14    pub fn open(path: &str) -> Result<Self, GraphError> {
15        let conn = Connection::open(path)?;
16        Self::set_performance_pragmas(&conn)?;
17        let store = SqliteGraphStore { conn };
18        store.initialize_schema()?;
19        Ok(store)
20    }
21
22    /// Create an in-memory graph database (for testing).
23    pub fn in_memory() -> Result<Self, GraphError> {
24        let conn = Connection::open_in_memory()?;
25        Self::set_performance_pragmas(&conn)?;
26        let store = SqliteGraphStore { conn };
27        store.initialize_schema()?;
28        Ok(store)
29    }
30
31    /// Apply SQLite performance pragmas for faster reads and writes.
32    fn set_performance_pragmas(conn: &Connection) -> Result<(), GraphError> {
33        conn.execute_batch(
34            "
35            PRAGMA journal_mode = WAL;
36            PRAGMA synchronous = NORMAL;
37            PRAGMA cache_size = -8000;
38            PRAGMA temp_store = MEMORY;
39            PRAGMA mmap_size = 268435456;
40            PRAGMA foreign_keys = ON;
41            ",
42        )?;
43        Ok(())
44    }
45
46    /// Temporarily disable foreign key enforcement (for bulk re-map operations).
47    /// Returns the actual FK state after the change (for verification).
48    pub fn set_foreign_keys(&self, enabled: bool) -> Result<bool, GraphError> {
49        let val = if enabled { "ON" } else { "OFF" };
50        self.conn
51            .execute_batch(&format!("PRAGMA foreign_keys = {};", val))?;
52        // Verify the change took effect
53        let actual: i32 = self
54            .conn
55            .pragma_query_value(None, "foreign_keys", |row| row.get(0))
56            .unwrap_or(if enabled { 1 } else { 0 });
57        Ok(actual != 0)
58    }
59
60    fn initialize_schema(&self) -> Result<(), GraphError> {
61        self.conn.execute_batch(
62            "
63            -- Schema version tracking
64            CREATE TABLE IF NOT EXISTS keel_meta (
65                key TEXT PRIMARY KEY,
66                value TEXT NOT NULL
67            );
68
69            -- Nodes
70            CREATE TABLE IF NOT EXISTS nodes (
71                id INTEGER PRIMARY KEY,
72                hash TEXT NOT NULL UNIQUE,
73                kind TEXT NOT NULL CHECK (kind IN ('module', 'class', 'function')),
74                name TEXT NOT NULL,
75                signature TEXT NOT NULL DEFAULT '',
76                file_path TEXT NOT NULL,
77                line_start INTEGER NOT NULL,
78                line_end INTEGER NOT NULL,
79                docstring TEXT,
80                is_public INTEGER NOT NULL DEFAULT 0,
81                type_hints_present INTEGER NOT NULL DEFAULT 0,
82                has_docstring INTEGER NOT NULL DEFAULT 0,
83                module_id INTEGER REFERENCES nodes(id),
84                package TEXT DEFAULT NULL,
85                resolution_tier TEXT NOT NULL DEFAULT '',
86                created_at TEXT NOT NULL DEFAULT (datetime('now')),
87                updated_at TEXT NOT NULL DEFAULT (datetime('now'))
88            );
89            CREATE INDEX IF NOT EXISTS idx_nodes_hash ON nodes(hash);
90            CREATE INDEX IF NOT EXISTS idx_nodes_file ON nodes(file_path);
91            CREATE INDEX IF NOT EXISTS idx_nodes_module ON nodes(module_id);
92            CREATE INDEX IF NOT EXISTS idx_nodes_kind ON nodes(kind);
93            CREATE INDEX IF NOT EXISTS idx_nodes_name_kind ON nodes(name, kind);
94
95            -- Previous hashes for rename tracking
96            CREATE TABLE IF NOT EXISTS previous_hashes (
97                node_id INTEGER NOT NULL REFERENCES nodes(id) ON DELETE CASCADE,
98                hash TEXT NOT NULL,
99                created_at TEXT NOT NULL DEFAULT (datetime('now')),
100                PRIMARY KEY (node_id, hash)
101            );
102
103            -- External endpoints
104            CREATE TABLE IF NOT EXISTS external_endpoints (
105                id INTEGER PRIMARY KEY,
106                node_id INTEGER NOT NULL REFERENCES nodes(id) ON DELETE CASCADE,
107                kind TEXT NOT NULL,
108                method TEXT NOT NULL DEFAULT '',
109                path TEXT NOT NULL,
110                direction TEXT NOT NULL CHECK (direction IN ('serves', 'calls'))
111            );
112            CREATE INDEX IF NOT EXISTS idx_endpoints_node ON external_endpoints(node_id);
113
114            -- Edges
115            CREATE TABLE IF NOT EXISTS edges (
116                id INTEGER PRIMARY KEY,
117                source_id INTEGER NOT NULL REFERENCES nodes(id) ON DELETE CASCADE,
118                target_id INTEGER NOT NULL REFERENCES nodes(id) ON DELETE CASCADE,
119                kind TEXT NOT NULL CHECK (kind IN ('calls', 'imports', 'inherits', 'contains')),
120                confidence REAL NOT NULL DEFAULT 1.0,
121                file_path TEXT NOT NULL,
122                line INTEGER NOT NULL,
123                UNIQUE(source_id, target_id, kind, file_path, line)
124            );
125            CREATE INDEX IF NOT EXISTS idx_edges_source ON edges(source_id);
126            CREATE INDEX IF NOT EXISTS idx_edges_target ON edges(target_id);
127            CREATE INDEX IF NOT EXISTS idx_edges_source_kind ON edges(source_id, kind);
128
129            -- Module profiles
130            CREATE TABLE IF NOT EXISTS module_profiles (
131                module_id INTEGER PRIMARY KEY REFERENCES nodes(id) ON DELETE CASCADE,
132                path TEXT NOT NULL,
133                function_count INTEGER NOT NULL DEFAULT 0,
134                class_count INTEGER NOT NULL DEFAULT 0,
135                line_count INTEGER NOT NULL DEFAULT 0,
136                function_name_prefixes TEXT NOT NULL DEFAULT '[]',
137                primary_types TEXT NOT NULL DEFAULT '[]',
138                import_sources TEXT NOT NULL DEFAULT '[]',
139                export_targets TEXT NOT NULL DEFAULT '[]',
140                external_endpoint_count INTEGER NOT NULL DEFAULT 0,
141                responsibility_keywords TEXT NOT NULL DEFAULT '[]'
142            );
143
144            -- Resolution cache
145            CREATE TABLE IF NOT EXISTS resolution_cache (
146                call_site_hash TEXT PRIMARY KEY,
147                resolved_node_id INTEGER REFERENCES nodes(id),
148                confidence REAL NOT NULL,
149                resolution_tier TEXT NOT NULL,
150                cached_at TEXT NOT NULL DEFAULT (datetime('now'))
151            );
152
153            -- Circuit breaker state
154            CREATE TABLE IF NOT EXISTS circuit_breaker (
155                error_code TEXT NOT NULL,
156                hash TEXT NOT NULL,
157                consecutive_failures INTEGER NOT NULL DEFAULT 0,
158                last_failure_at TEXT NOT NULL DEFAULT (datetime('now')),
159                downgraded INTEGER NOT NULL DEFAULT 0,
160                PRIMARY KEY (error_code, hash)
161            );
162            ",
163        )?;
164
165        // Set schema version if not present (new databases get current version)
166        self.conn.execute(
167            "INSERT OR IGNORE INTO keel_meta (key, value) VALUES ('schema_version', ?1)",
168            params![SCHEMA_VERSION.to_string()],
169        )?;
170
171        // Run migrations for existing databases
172        self.run_migrations()?;
173
174        // Create indexes that depend on columns added by migrations.
175        // These use IF NOT EXISTS so they're safe to run on every open.
176        let _ = self
177            .conn
178            .execute_batch("CREATE INDEX IF NOT EXISTS idx_nodes_package ON nodes(package)");
179
180        Ok(())
181    }
182
183    /// Run schema migrations from current version to SCHEMA_VERSION.
184    fn run_migrations(&self) -> Result<(), GraphError> {
185        let current = self.schema_version()?;
186        if current >= SCHEMA_VERSION {
187            return Ok(());
188        }
189        if current < 2 {
190            self.migrate_v1_to_v2()?;
191        }
192        if current < 3 {
193            self.migrate_v2_to_v3()?;
194        }
195        Ok(())
196    }
197
198    /// Migrate from schema v1 to v2: add resolution_tier to nodes, confidence to edges.
199    fn migrate_v1_to_v2(&self) -> Result<(), GraphError> {
200        // Add resolution_tier column to nodes (ignore if already exists)
201        let _ = self.conn.execute_batch(
202            "ALTER TABLE nodes ADD COLUMN resolution_tier TEXT NOT NULL DEFAULT ''",
203        );
204        // Add confidence column to edges (ignore if already exists)
205        let _ = self.conn.execute_batch(
206            "ALTER TABLE edges ADD COLUMN confidence REAL NOT NULL DEFAULT 1.0",
207        );
208        // Update schema version to 2
209        self.conn.execute(
210            "UPDATE keel_meta SET value = '2' WHERE key = 'schema_version'",
211            [],
212        )?;
213        Ok(())
214    }
215
216    /// Migrate from schema v2 to v3: add package column to nodes.
217    fn migrate_v2_to_v3(&self) -> Result<(), GraphError> {
218        let _ = self
219            .conn
220            .execute_batch("ALTER TABLE nodes ADD COLUMN package TEXT DEFAULT NULL");
221        let _ = self
222            .conn
223            .execute_batch("CREATE INDEX IF NOT EXISTS idx_nodes_package ON nodes(package)");
224        self.conn.execute(
225            "UPDATE keel_meta SET value = '3' WHERE key = 'schema_version'",
226            [],
227        )?;
228        Ok(())
229    }
230
231    /// Get the current schema version.
232    pub fn schema_version(&self) -> Result<u32, GraphError> {
233        let version: String = self.conn.query_row(
234            "SELECT value FROM keel_meta WHERE key = 'schema_version'",
235            [],
236            |row| row.get(0),
237        )?;
238        version
239            .parse()
240            .map_err(|e| GraphError::Internal(format!("Invalid schema version: {}", e)))
241    }
242
243    /// Remove edges whose source or target node no longer exists.
244    pub fn cleanup_orphaned_edges(&self) -> Result<u64, GraphError> {
245        let deleted = self.conn.execute(
246            "DELETE FROM edges WHERE source_id NOT IN (SELECT id FROM nodes) OR target_id NOT IN (SELECT id FROM nodes)",
247            [],
248        )?;
249        Ok(deleted as u64)
250    }
251
252    /// Clear all graph data (nodes, edges, etc.) for a full re-map.
253    /// Preserves schema and metadata.
254    pub fn clear_all(&mut self) -> Result<(), GraphError> {
255        self.conn.execute_batch(
256            "
257            DELETE FROM edges;
258            DELETE FROM resolution_cache;
259            DELETE FROM circuit_breaker;
260            DELETE FROM module_profiles;
261            DELETE FROM external_endpoints;
262            DELETE FROM previous_hashes;
263            DELETE FROM nodes;
264            ",
265        )?;
266        Ok(())
267    }
268
269    pub(crate) fn row_to_node(row: &rusqlite::Row) -> SqlResult<GraphNode> {
270        let kind_str: String = row.get("kind")?;
271        let kind = match kind_str.as_str() {
272            "module" => NodeKind::Module,
273            "class" => NodeKind::Class,
274            "function" => NodeKind::Function,
275            _ => NodeKind::Function, // fallback
276        };
277        Ok(GraphNode {
278            id: row.get("id")?,
279            hash: row.get("hash")?,
280            kind,
281            name: row.get("name")?,
282            signature: row.get("signature")?,
283            file_path: row.get("file_path")?,
284            line_start: row.get("line_start")?,
285            line_end: row.get("line_end")?,
286            docstring: row.get("docstring")?,
287            is_public: row.get::<_, i32>("is_public")? != 0,
288            type_hints_present: row.get::<_, i32>("type_hints_present")? != 0,
289            has_docstring: row.get::<_, i32>("has_docstring")? != 0,
290            external_endpoints: Vec::new(), // loaded separately
291            previous_hashes: Vec::new(),    // loaded separately
292            module_id: row.get::<_, Option<u64>>("module_id")?.unwrap_or(0),
293            package: row.get::<_, Option<String>>("package").unwrap_or(None),
294        })
295    }
296
297    pub(crate) fn load_endpoints(&self, node_id: u64) -> Vec<ExternalEndpoint> {
298        let mut stmt = match self
299            .conn
300            .prepare("SELECT kind, method, path, direction FROM external_endpoints WHERE node_id = ?1")
301        {
302            Ok(s) => s,
303            Err(e) => {
304                eprintln!("[keel] load_endpoints: prepare failed: {e}");
305                return Vec::new();
306            }
307        };
308
309        let result = match stmt.query_map(params![node_id], |row| {
310            Ok(ExternalEndpoint {
311                kind: row.get(0)?,
312                method: row.get(1)?,
313                path: row.get(2)?,
314                direction: row.get(3)?,
315            })
316        }) {
317            Ok(rows) => rows.filter_map(|r| r.ok()).collect(),
318            Err(e) => {
319                eprintln!("[keel] load_endpoints: query failed: {e}");
320                Vec::new()
321            }
322        };
323        result
324    }
325
326    pub(crate) fn load_previous_hashes(&self, node_id: u64) -> Vec<String> {
327        let mut stmt = match self
328            .conn
329            .prepare(
330                "SELECT hash FROM previous_hashes WHERE node_id = ?1 ORDER BY created_at DESC LIMIT 3",
331            )
332        {
333            Ok(s) => s,
334            Err(e) => {
335                eprintln!("[keel] load_previous_hashes: prepare failed: {e}");
336                return Vec::new();
337            }
338        };
339
340        let result = match stmt.query_map(params![node_id], |row| row.get(0)) {
341            Ok(rows) => rows.filter_map(|r| r.ok()).collect(),
342            Err(e) => {
343                eprintln!("[keel] load_previous_hashes: query failed: {e}");
344                Vec::new()
345            }
346        };
347        result
348    }
349
350    pub(crate) fn node_with_relations(&self, mut node: GraphNode) -> GraphNode {
351        node.external_endpoints = self.load_endpoints(node.id);
352        node.previous_hashes = self.load_previous_hashes(node.id);
353        node
354    }
355
356    /// Insert or update module profiles in bulk.
357    /// Uses INSERT ... ON CONFLICT DO UPDATE for upsert semantics.
358    pub fn upsert_module_profiles(
359        &self,
360        profiles: Vec<crate::types::ModuleProfile>,
361    ) -> Result<(), GraphError> {
362        let tx = self.conn.unchecked_transaction()?;
363        {
364            let mut stmt = tx.prepare(
365                "INSERT INTO module_profiles (
366                    module_id, path, function_count, class_count, line_count,
367                    function_name_prefixes, primary_types, import_sources,
368                    export_targets, external_endpoint_count, responsibility_keywords
369                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)
370                ON CONFLICT(module_id) DO UPDATE SET
371                    path = excluded.path,
372                    function_count = excluded.function_count,
373                    class_count = excluded.class_count,
374                    line_count = excluded.line_count,
375                    function_name_prefixes = excluded.function_name_prefixes,
376                    primary_types = excluded.primary_types,
377                    import_sources = excluded.import_sources,
378                    export_targets = excluded.export_targets,
379                    external_endpoint_count = excluded.external_endpoint_count,
380                    responsibility_keywords = excluded.responsibility_keywords",
381            )?;
382            for p in &profiles {
383                let prefixes_json = serde_json::to_string(&p.function_name_prefixes)
384                    .unwrap_or_else(|_| "[]".to_string());
385                let types_json = serde_json::to_string(&p.primary_types)
386                    .unwrap_or_else(|_| "[]".to_string());
387                let imports_json = serde_json::to_string(&p.import_sources)
388                    .unwrap_or_else(|_| "[]".to_string());
389                let exports_json = serde_json::to_string(&p.export_targets)
390                    .unwrap_or_else(|_| "[]".to_string());
391                let keywords_json = serde_json::to_string(&p.responsibility_keywords)
392                    .unwrap_or_else(|_| "[]".to_string());
393                stmt.execute(params![
394                    p.module_id, p.path, p.function_count, p.class_count, p.line_count,
395                    prefixes_json, types_json, imports_json, exports_json,
396                    p.external_endpoint_count, keywords_json,
397                ])?;
398            }
399        }
400        tx.commit()?;
401        Ok(())
402    }
403}
404
405#[cfg(test)]
406#[path = "sqlite_tests.rs"]
407mod tests;
keel_core/sqlite.rs

keel_core/
sqlite.rs
