codeprysm_core/lazy/
partition.rs

1//! SQLite Partition Connection and CRUD Operations
2//!
3//! This module provides a wrapper around rusqlite for partition database operations.
4//! Each partition is a self-contained SQLite database storing nodes and edges
5//! for a subset of the repository.
6
7use crate::graph::{Edge, EdgeType, Node, NodeType};
8use rusqlite::{Connection, OptionalExtension, Result as SqliteResult, params};
9use std::path::Path;
10use thiserror::Error;
11
12use super::schema::{
13    PARTITION_SCHEMA_VERSION, SCHEMA_CREATE_EDGES, SCHEMA_CREATE_INDEXES, SCHEMA_CREATE_METADATA,
14    SCHEMA_CREATE_NODES,
15};
16
17/// Errors that can occur during partition operations
18#[derive(Debug, Error)]
19pub enum PartitionError {
20    #[error("SQLite error: {0}")]
21    Sqlite(#[from] rusqlite::Error),
22
23    #[error("Schema version mismatch: expected {expected}, found {found}")]
24    SchemaVersionMismatch { expected: String, found: String },
25
26    #[error("Node not found: {0}")]
27    NodeNotFound(String),
28
29    #[error("JSON serialization error: {0}")]
30    JsonError(#[from] serde_json::Error),
31
32    #[error("IO error: {0}")]
33    IoError(#[from] std::io::Error),
34}
35
36/// A connection to a partition SQLite database
37pub struct PartitionConnection {
38    conn: Connection,
39    /// The partition ID (usually based on directory path or content hash)
40    partition_id: String,
41}
42
43impl PartitionConnection {
44    /// Open an existing partition database
45    ///
46    /// If the partition is at an older schema version (v1.0), it will be
47    /// automatically migrated to the current version (v1.1).
48    pub fn open(path: &Path, partition_id: &str) -> Result<Self, PartitionError> {
49        let conn = Connection::open(path)?;
50        Self::configure_connection(&conn)?;
51
52        let pc = Self {
53            conn,
54            partition_id: partition_id.to_string(),
55        };
56
57        // Check and migrate schema version if needed
58        if let Some(version) = pc.get_metadata("schema_version")? {
59            match version.as_str() {
60                v if v == PARTITION_SCHEMA_VERSION => {
61                    // Already at current version, nothing to do
62                }
63                "1.0" => {
64                    // Migrate from v1.0 to v1.1
65                    pc.migrate_v1_0_to_v1_1()?;
66                }
67                _ => {
68                    return Err(PartitionError::SchemaVersionMismatch {
69                        expected: PARTITION_SCHEMA_VERSION.to_string(),
70                        found: version,
71                    });
72                }
73            }
74        }
75
76        Ok(pc)
77    }
78
79    /// Migrate partition from v1.0 to v1.1
80    ///
81    /// Adds version_spec and is_dev_dependency columns to edges table.
82    fn migrate_v1_0_to_v1_1(&self) -> Result<(), PartitionError> {
83        // Execute migration SQL (each ALTER TABLE is a separate statement)
84        self.conn
85            .execute("ALTER TABLE edges ADD COLUMN version_spec TEXT", [])?;
86        self.conn
87            .execute("ALTER TABLE edges ADD COLUMN is_dev_dependency INTEGER", [])?;
88
89        // Update schema version
90        self.set_metadata("schema_version", PARTITION_SCHEMA_VERSION)?;
91
92        Ok(())
93    }
94
95    /// Create a new partition database with schema
96    pub fn create(path: &Path, partition_id: &str) -> Result<Self, PartitionError> {
97        // Ensure parent directory exists
98        if let Some(parent) = path.parent() {
99            std::fs::create_dir_all(parent)?;
100        }
101
102        let conn = Connection::open(path)?;
103        Self::configure_connection(&conn)?;
104
105        // Create schema
106        conn.execute(SCHEMA_CREATE_NODES, [])?;
107        conn.execute(SCHEMA_CREATE_EDGES, [])?;
108        conn.execute(SCHEMA_CREATE_METADATA, [])?;
109        conn.execute_batch(SCHEMA_CREATE_INDEXES)?;
110
111        let pc = Self {
112            conn,
113            partition_id: partition_id.to_string(),
114        };
115
116        // Store schema version
117        pc.set_metadata("schema_version", PARTITION_SCHEMA_VERSION)?;
118        pc.set_metadata("partition_id", partition_id)?;
119
120        Ok(pc)
121    }
122
123    /// Create an in-memory partition database (for testing)
124    pub fn in_memory(partition_id: &str) -> Result<Self, PartitionError> {
125        let conn = Connection::open_in_memory()?;
126        Self::configure_connection(&conn)?;
127
128        // Create schema
129        conn.execute(SCHEMA_CREATE_NODES, [])?;
130        conn.execute(SCHEMA_CREATE_EDGES, [])?;
131        conn.execute(SCHEMA_CREATE_METADATA, [])?;
132        conn.execute_batch(SCHEMA_CREATE_INDEXES)?;
133
134        let pc = Self {
135            conn,
136            partition_id: partition_id.to_string(),
137        };
138
139        pc.set_metadata("schema_version", PARTITION_SCHEMA_VERSION)?;
140        pc.set_metadata("partition_id", partition_id)?;
141
142        Ok(pc)
143    }
144
145    /// Configure connection with optimal settings
146    fn configure_connection(conn: &Connection) -> SqliteResult<()> {
147        // Enable WAL mode for better concurrent read performance
148        conn.pragma_update(None, "journal_mode", "WAL")?;
149        // Enable foreign keys (not currently used but good practice)
150        conn.pragma_update(None, "foreign_keys", "ON")?;
151        // Increase cache size (negative value = KB)
152        conn.pragma_update(None, "cache_size", -64000)?; // 64MB cache
153        // Synchronous mode: NORMAL is good balance of safety/speed
154        conn.pragma_update(None, "synchronous", "NORMAL")?;
155        // Temp store in memory for better performance
156        conn.pragma_update(None, "temp_store", "MEMORY")?;
157        // Enable memory-mapped I/O for reads
158        conn.pragma_update(None, "mmap_size", 268435456)?; // 256MB mmap
159        Ok(())
160    }
161
162    /// Get the partition ID
163    pub fn partition_id(&self) -> &str {
164        &self.partition_id
165    }
166
167    // =========================================================================
168    // Metadata Operations
169    // =========================================================================
170
171    /// Get a metadata value
172    pub fn get_metadata(&self, key: &str) -> Result<Option<String>, PartitionError> {
173        let result = self
174            .conn
175            .query_row(
176                "SELECT value FROM partition_metadata WHERE key = ?1",
177                [key],
178                |row| row.get(0),
179            )
180            .optional()?;
181        Ok(result)
182    }
183
184    /// Set a metadata value
185    pub fn set_metadata(&self, key: &str, value: &str) -> Result<(), PartitionError> {
186        self.conn.execute(
187            "INSERT OR REPLACE INTO partition_metadata (key, value) VALUES (?1, ?2)",
188            params![key, value],
189        )?;
190        Ok(())
191    }
192
193    // =========================================================================
194    // Node Operations
195    // =========================================================================
196
197    /// Insert a node into the partition
198    pub fn insert_node(&self, node: &Node) -> Result<(), PartitionError> {
199        let metadata_json = if node.metadata.is_empty() {
200            None
201        } else {
202            Some(serde_json::to_string(&node.metadata)?)
203        };
204
205        self.conn.execute(
206            r#"
207            INSERT OR REPLACE INTO nodes
208                (id, name, node_type, kind, subtype, file, line, end_line, text, hash, metadata_json)
209            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)
210            "#,
211            params![
212                node.id,
213                node.name,
214                node.node_type.as_str(),
215                node.kind,
216                node.subtype,
217                node.file,
218                node.line as i64,
219                node.end_line as i64,
220                node.text,
221                node.hash,
222                metadata_json,
223            ],
224        )?;
225        Ok(())
226    }
227
228    /// Insert multiple nodes in a transaction
229    pub fn insert_nodes(&self, nodes: &[Node]) -> Result<(), PartitionError> {
230        let tx = self.conn.unchecked_transaction()?;
231
232        {
233            let mut stmt = tx.prepare(
234                r#"
235                INSERT OR REPLACE INTO nodes
236                    (id, name, node_type, kind, subtype, file, line, end_line, text, hash, metadata_json)
237                VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)
238                "#,
239            )?;
240
241            for node in nodes {
242                let metadata_json = if node.metadata.is_empty() {
243                    None
244                } else {
245                    Some(serde_json::to_string(&node.metadata)?)
246                };
247
248                stmt.execute(params![
249                    node.id,
250                    node.name,
251                    node.node_type.as_str(),
252                    node.kind,
253                    node.subtype,
254                    node.file,
255                    node.line as i64,
256                    node.end_line as i64,
257                    node.text,
258                    node.hash,
259                    metadata_json,
260                ])?;
261            }
262        }
263
264        tx.commit()?;
265        Ok(())
266    }
267
268    /// Get a node by ID
269    pub fn get_node(&self, id: &str) -> Result<Option<Node>, PartitionError> {
270        let result = self
271            .conn
272            .query_row(
273                r#"
274                SELECT id, name, node_type, kind, subtype, file, line, end_line, text, hash, metadata_json
275                FROM nodes WHERE id = ?1
276                "#,
277                [id],
278                Self::row_to_node,
279            )
280            .optional()?;
281        Ok(result)
282    }
283
284    /// Query nodes by file path
285    pub fn query_nodes_by_file(&self, file: &str) -> Result<Vec<Node>, PartitionError> {
286        let mut stmt = self.conn.prepare(
287            r#"
288            SELECT id, name, node_type, kind, subtype, file, line, end_line, text, hash, metadata_json
289            FROM nodes WHERE file = ?1
290            "#,
291        )?;
292
293        let nodes = stmt
294            .query_map([file], Self::row_to_node)?
295            .collect::<SqliteResult<Vec<_>>>()?;
296
297        Ok(nodes)
298    }
299
300    /// Query all nodes in the partition
301    pub fn query_all_nodes(&self) -> Result<Vec<Node>, PartitionError> {
302        let mut stmt = self.conn.prepare(
303            r#"
304            SELECT id, name, node_type, kind, subtype, file, line, end_line, text, hash, metadata_json
305            FROM nodes
306            "#,
307        )?;
308
309        let nodes = stmt
310            .query_map([], Self::row_to_node)?
311            .collect::<SqliteResult<Vec<_>>>()?;
312
313        Ok(nodes)
314    }
315
316    /// Delete nodes by file path
317    pub fn delete_nodes_by_file(&self, file: &str) -> Result<usize, PartitionError> {
318        let deleted = self
319            .conn
320            .execute("DELETE FROM nodes WHERE file = ?1", [file])?;
321        Ok(deleted)
322    }
323
324    /// Delete a node by ID
325    pub fn delete_node(&self, id: &str) -> Result<bool, PartitionError> {
326        let deleted = self.conn.execute("DELETE FROM nodes WHERE id = ?1", [id])?;
327        Ok(deleted > 0)
328    }
329
330    /// Get node count
331    pub fn node_count(&self) -> Result<usize, PartitionError> {
332        let count: i64 = self
333            .conn
334            .query_row("SELECT COUNT(*) FROM nodes", [], |row| row.get(0))?;
335        Ok(count as usize)
336    }
337
338    /// Convert a database row to a Node
339    fn row_to_node(row: &rusqlite::Row<'_>) -> SqliteResult<Node> {
340        let node_type_str: String = row.get(2)?;
341        let metadata_json: Option<String> = row.get(10)?;
342
343        // Handle legacy FILE type by converting to Container
344        let (node_type, is_legacy_file) = match node_type_str.as_str() {
345            "FILE" => (NodeType::Container, true), // Legacy: convert to Container
346            "Container" => (NodeType::Container, false),
347            "Callable" => (NodeType::Callable, false),
348            "Data" => (NodeType::Data, false),
349            _ => (NodeType::Container, false), // Default fallback
350        };
351
352        // For legacy FILE nodes, ensure kind is "file"
353        let kind: Option<String> = if is_legacy_file {
354            Some("file".to_string())
355        } else {
356            row.get(3)?
357        };
358
359        let metadata = metadata_json
360            .and_then(|json| serde_json::from_str(&json).ok())
361            .unwrap_or_default();
362
363        Ok(Node {
364            id: row.get(0)?,
365            name: row.get(1)?,
366            node_type,
367            kind,
368            subtype: row.get(4)?,
369            file: row.get(5)?,
370            line: row.get::<_, i64>(6)? as usize,
371            end_line: row.get::<_, i64>(7)? as usize,
372            text: row.get(8)?,
373            hash: row.get(9)?,
374            metadata,
375        })
376    }
377
378    // =========================================================================
379    // Edge Operations
380    // =========================================================================
381
382    /// Insert an edge into the partition
383    pub fn insert_edge(&self, edge: &Edge) -> Result<(), PartitionError> {
384        self.conn.execute(
385            r#"
386            INSERT OR IGNORE INTO edges (source, target, edge_type, ref_line, ident, version_spec, is_dev_dependency)
387            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
388            "#,
389            params![
390                edge.source,
391                edge.target,
392                edge.edge_type.as_str(),
393                edge.ref_line.map(|l| l as i64),
394                edge.ident,
395                edge.version_spec,
396                edge.is_dev_dependency,
397            ],
398        )?;
399        Ok(())
400    }
401
402    /// Insert multiple edges in a transaction
403    pub fn insert_edges(&self, edges: &[Edge]) -> Result<(), PartitionError> {
404        let tx = self.conn.unchecked_transaction()?;
405
406        {
407            let mut stmt = tx.prepare(
408                r#"
409                INSERT OR IGNORE INTO edges (source, target, edge_type, ref_line, ident, version_spec, is_dev_dependency)
410                VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
411                "#,
412            )?;
413
414            for edge in edges {
415                stmt.execute(params![
416                    edge.source,
417                    edge.target,
418                    edge.edge_type.as_str(),
419                    edge.ref_line.map(|l| l as i64),
420                    edge.ident,
421                    edge.version_spec,
422                    edge.is_dev_dependency,
423                ])?;
424            }
425        }
426
427        tx.commit()?;
428        Ok(())
429    }
430
431    /// Query edges by source node ID
432    pub fn query_edges_by_source(&self, source: &str) -> Result<Vec<Edge>, PartitionError> {
433        let mut stmt = self.conn.prepare(
434            r#"
435            SELECT source, target, edge_type, ref_line, ident, version_spec, is_dev_dependency
436            FROM edges WHERE source = ?1
437            "#,
438        )?;
439
440        let edges = stmt
441            .query_map([source], Self::row_to_edge)?
442            .collect::<SqliteResult<Vec<_>>>()?;
443
444        Ok(edges)
445    }
446
447    /// Query edges by target node ID
448    pub fn query_edges_by_target(&self, target: &str) -> Result<Vec<Edge>, PartitionError> {
449        let mut stmt = self.conn.prepare(
450            r#"
451            SELECT source, target, edge_type, ref_line, ident, version_spec, is_dev_dependency
452            FROM edges WHERE target = ?1
453            "#,
454        )?;
455
456        let edges = stmt
457            .query_map([target], Self::row_to_edge)?
458            .collect::<SqliteResult<Vec<_>>>()?;
459
460        Ok(edges)
461    }
462
463    /// Query all edges in the partition
464    pub fn query_all_edges(&self) -> Result<Vec<Edge>, PartitionError> {
465        let mut stmt = self.conn.prepare(
466            r#"
467            SELECT source, target, edge_type, ref_line, ident, version_spec, is_dev_dependency
468            FROM edges
469            "#,
470        )?;
471
472        let edges = stmt
473            .query_map([], Self::row_to_edge)?
474            .collect::<SqliteResult<Vec<_>>>()?;
475
476        Ok(edges)
477    }
478
479    /// Delete edges involving a node (as source or target)
480    pub fn delete_edges_involving(&self, node_id: &str) -> Result<usize, PartitionError> {
481        let deleted = self.conn.execute(
482            "DELETE FROM edges WHERE source = ?1 OR target = ?1",
483            [node_id],
484        )?;
485        Ok(deleted)
486    }
487
488    /// Delete edges by source node
489    pub fn delete_edges_by_source(&self, source: &str) -> Result<usize, PartitionError> {
490        let deleted = self
491            .conn
492            .execute("DELETE FROM edges WHERE source = ?1", [source])?;
493        Ok(deleted)
494    }
495
496    /// Get edge count
497    pub fn edge_count(&self) -> Result<usize, PartitionError> {
498        let count: i64 = self
499            .conn
500            .query_row("SELECT COUNT(*) FROM edges", [], |row| row.get(0))?;
501        Ok(count as usize)
502    }
503
504    /// Convert a database row to an Edge
505    fn row_to_edge(row: &rusqlite::Row<'_>) -> SqliteResult<Edge> {
506        let edge_type_str: String = row.get(2)?;
507        let edge_type = match edge_type_str.as_str() {
508            "CONTAINS" => EdgeType::Contains,
509            "USES" => EdgeType::Uses,
510            "DEFINES" => EdgeType::Defines,
511            "DEPENDS_ON" => EdgeType::DependsOn,
512            _ => EdgeType::Uses, // Default fallback
513        };
514
515        Ok(Edge {
516            source: row.get(0)?,
517            target: row.get(1)?,
518            edge_type,
519            ref_line: row.get::<_, Option<i64>>(3)?.map(|l| l as usize),
520            ident: row.get(4)?,
521            version_spec: row.get(5).ok().flatten(),
522            is_dev_dependency: row.get(6).ok().flatten(),
523        })
524    }
525
526    // =========================================================================
527    // Bulk Operations
528    // =========================================================================
529
530    /// Clear all data from the partition (keeps schema)
531    pub fn clear(&self) -> Result<(), PartitionError> {
532        self.conn.execute("DELETE FROM edges", [])?;
533        self.conn.execute("DELETE FROM nodes", [])?;
534        Ok(())
535    }
536
537    /// Get partition statistics
538    pub fn stats(&self) -> Result<PartitionStats, PartitionError> {
539        Ok(PartitionStats {
540            node_count: self.node_count()?,
541            edge_count: self.edge_count()?,
542            partition_id: self.partition_id.clone(),
543        })
544    }
545}
546
547/// Statistics about a partition
548#[derive(Debug, Clone)]
549pub struct PartitionStats {
550    pub node_count: usize,
551    pub edge_count: usize,
552    pub partition_id: String,
553}
554
555#[cfg(test)]
556mod tests {
557    use super::*;
558    use crate::graph::{CallableKind, ContainerKind, NodeMetadata};
559
560    fn create_test_node(id: &str, name: &str, file: &str) -> Node {
561        Node {
562            id: id.to_string(),
563            name: name.to_string(),
564            node_type: NodeType::Callable,
565            kind: Some(CallableKind::Function.as_str().to_string()),
566            subtype: None,
567            file: file.to_string(),
568            line: 1,
569            end_line: 10,
570            text: Some("def test(): pass".to_string()),
571            hash: None,
572            metadata: NodeMetadata::default(),
573        }
574    }
575
576    fn create_test_edge(source: &str, target: &str, edge_type: EdgeType) -> Edge {
577        Edge {
578            source: source.to_string(),
579            target: target.to_string(),
580            edge_type,
581            ref_line: Some(5),
582            ident: Some("test".to_string()),
583            version_spec: None,
584            is_dev_dependency: None,
585        }
586    }
587
588    #[test]
589    fn test_create_in_memory() {
590        let conn = PartitionConnection::in_memory("test_partition").unwrap();
591        assert_eq!(conn.partition_id(), "test_partition");
592        assert_eq!(conn.node_count().unwrap(), 0);
593        assert_eq!(conn.edge_count().unwrap(), 0);
594    }
595
596    #[test]
597    fn test_insert_and_get_node() {
598        let conn = PartitionConnection::in_memory("test").unwrap();
599        let node = create_test_node("test.py:test_func", "test_func", "test.py");
600
601        conn.insert_node(&node).unwrap();
602        assert_eq!(conn.node_count().unwrap(), 1);
603
604        let retrieved = conn.get_node("test.py:test_func").unwrap().unwrap();
605        assert_eq!(retrieved.id, "test.py:test_func");
606        assert_eq!(retrieved.name, "test_func");
607        assert_eq!(retrieved.node_type, NodeType::Callable);
608        assert_eq!(retrieved.kind, Some("function".to_string()));
609        assert_eq!(retrieved.file, "test.py");
610        assert_eq!(retrieved.line, 1);
611        assert_eq!(retrieved.end_line, 10);
612    }
613
614    #[test]
615    fn test_insert_bulk_nodes() {
616        let conn = PartitionConnection::in_memory("test").unwrap();
617        let nodes: Vec<Node> = (0..100)
618            .map(|i| {
619                create_test_node(
620                    &format!("test.py:func_{}", i),
621                    &format!("func_{}", i),
622                    "test.py",
623                )
624            })
625            .collect();
626
627        conn.insert_nodes(&nodes).unwrap();
628        assert_eq!(conn.node_count().unwrap(), 100);
629    }
630
631    #[test]
632    fn test_query_nodes_by_file() {
633        let conn = PartitionConnection::in_memory("test").unwrap();
634
635        conn.insert_node(&create_test_node("a.py:func1", "func1", "a.py"))
636            .unwrap();
637        conn.insert_node(&create_test_node("a.py:func2", "func2", "a.py"))
638            .unwrap();
639        conn.insert_node(&create_test_node("b.py:func3", "func3", "b.py"))
640            .unwrap();
641
642        let a_nodes = conn.query_nodes_by_file("a.py").unwrap();
643        assert_eq!(a_nodes.len(), 2);
644
645        let b_nodes = conn.query_nodes_by_file("b.py").unwrap();
646        assert_eq!(b_nodes.len(), 1);
647    }
648
649    #[test]
650    fn test_delete_node() {
651        let conn = PartitionConnection::in_memory("test").unwrap();
652        let node = create_test_node("test.py:func", "func", "test.py");
653
654        conn.insert_node(&node).unwrap();
655        assert_eq!(conn.node_count().unwrap(), 1);
656
657        let deleted = conn.delete_node("test.py:func").unwrap();
658        assert!(deleted);
659        assert_eq!(conn.node_count().unwrap(), 0);
660
661        // Deleting non-existent node returns false
662        let deleted = conn.delete_node("nonexistent").unwrap();
663        assert!(!deleted);
664    }
665
666    #[test]
667    fn test_insert_and_query_edges() {
668        let conn = PartitionConnection::in_memory("test").unwrap();
669
670        let edge1 = create_test_edge("a", "b", EdgeType::Contains);
671        let edge2 = create_test_edge("a", "c", EdgeType::Uses);
672        let edge3 = create_test_edge("b", "c", EdgeType::Defines);
673
674        conn.insert_edge(&edge1).unwrap();
675        conn.insert_edge(&edge2).unwrap();
676        conn.insert_edge(&edge3).unwrap();
677        assert_eq!(conn.edge_count().unwrap(), 3);
678
679        // Query by source
680        let from_a = conn.query_edges_by_source("a").unwrap();
681        assert_eq!(from_a.len(), 2);
682
683        // Query by target
684        let to_c = conn.query_edges_by_target("c").unwrap();
685        assert_eq!(to_c.len(), 2);
686    }
687
688    #[test]
689    fn test_insert_bulk_edges() {
690        let conn = PartitionConnection::in_memory("test").unwrap();
691        let edges: Vec<Edge> = (0..100)
692            .map(|i| {
693                create_test_edge(
694                    &format!("node_{}", i),
695                    &format!("node_{}", i + 1),
696                    EdgeType::Uses,
697                )
698            })
699            .collect();
700
701        conn.insert_edges(&edges).unwrap();
702        assert_eq!(conn.edge_count().unwrap(), 100);
703    }
704
705    #[test]
706    fn test_delete_edges_involving() {
707        let conn = PartitionConnection::in_memory("test").unwrap();
708
709        conn.insert_edge(&create_test_edge("a", "b", EdgeType::Contains))
710            .unwrap();
711        conn.insert_edge(&create_test_edge("b", "c", EdgeType::Uses))
712            .unwrap();
713        conn.insert_edge(&create_test_edge("c", "a", EdgeType::Defines))
714            .unwrap();
715        assert_eq!(conn.edge_count().unwrap(), 3);
716
717        // Delete edges involving "b" (should delete 2 edges)
718        let deleted = conn.delete_edges_involving("b").unwrap();
719        assert_eq!(deleted, 2);
720        assert_eq!(conn.edge_count().unwrap(), 1);
721    }
722
723    #[test]
724    fn test_node_with_metadata() {
725        let conn = PartitionConnection::in_memory("test").unwrap();
726
727        let mut node = create_test_node("test.py:async_func", "async_func", "test.py");
728        node.metadata.is_async = Some(true);
729        node.metadata.visibility = Some("public".to_string());
730        node.metadata.decorators = Some(vec!["staticmethod".to_string()]);
731
732        conn.insert_node(&node).unwrap();
733
734        let retrieved = conn.get_node("test.py:async_func").unwrap().unwrap();
735        assert_eq!(retrieved.metadata.is_async, Some(true));
736        assert_eq!(retrieved.metadata.visibility, Some("public".to_string()));
737        assert_eq!(
738            retrieved.metadata.decorators,
739            Some(vec!["staticmethod".to_string()])
740        );
741    }
742
743    #[test]
744    fn test_file_node_with_hash() {
745        let conn = PartitionConnection::in_memory("test").unwrap();
746
747        let node = Node {
748            id: "test.py".to_string(),
749            name: "test.py".to_string(),
750            node_type: NodeType::Container,
751            kind: Some(ContainerKind::File.as_str().to_string()),
752            subtype: None,
753            file: "test.py".to_string(),
754            line: 1,
755            end_line: 100,
756            text: None,
757            hash: Some("abc123".to_string()),
758            metadata: NodeMetadata::default(),
759        };
760
761        conn.insert_node(&node).unwrap();
762
763        let retrieved = conn.get_node("test.py").unwrap().unwrap();
764        assert_eq!(retrieved.hash, Some("abc123".to_string()));
765    }
766
767    #[test]
768    fn test_stats() {
769        let conn = PartitionConnection::in_memory("test_partition").unwrap();
770
771        conn.insert_node(&create_test_node("a", "a", "a.py"))
772            .unwrap();
773        conn.insert_node(&create_test_node("b", "b", "b.py"))
774            .unwrap();
775        conn.insert_edge(&create_test_edge("a", "b", EdgeType::Uses))
776            .unwrap();
777
778        let stats = conn.stats().unwrap();
779        assert_eq!(stats.node_count, 2);
780        assert_eq!(stats.edge_count, 1);
781        assert_eq!(stats.partition_id, "test_partition");
782    }
783
784    #[test]
785    fn test_clear() {
786        let conn = PartitionConnection::in_memory("test").unwrap();
787
788        conn.insert_node(&create_test_node("a", "a", "a.py"))
789            .unwrap();
790        conn.insert_edge(&create_test_edge("a", "b", EdgeType::Uses))
791            .unwrap();
792
793        conn.clear().unwrap();
794
795        assert_eq!(conn.node_count().unwrap(), 0);
796        assert_eq!(conn.edge_count().unwrap(), 0);
797    }
798
799    #[test]
800    fn test_depends_on_edge_with_metadata() {
801        let conn = PartitionConnection::in_memory("test").unwrap();
802
803        // Create a DependsOn edge with version_spec and is_dev_dependency
804        let edge = Edge {
805            source: "pkg/foo".to_string(),
806            target: "pkg/bar".to_string(),
807            edge_type: EdgeType::DependsOn,
808            ref_line: None,
809            ident: Some("bar".to_string()),
810            version_spec: Some("^1.0.0".to_string()),
811            is_dev_dependency: Some(true),
812        };
813
814        conn.insert_edge(&edge).unwrap();
815        assert_eq!(conn.edge_count().unwrap(), 1);
816
817        // Query the edge back
818        let edges = conn.query_edges_by_source("pkg/foo").unwrap();
819        assert_eq!(edges.len(), 1);
820
821        let retrieved = &edges[0];
822        assert_eq!(retrieved.source, "pkg/foo");
823        assert_eq!(retrieved.target, "pkg/bar");
824        assert_eq!(retrieved.edge_type, EdgeType::DependsOn);
825        assert_eq!(retrieved.ident, Some("bar".to_string()));
826        assert_eq!(retrieved.version_spec, Some("^1.0.0".to_string()));
827        assert_eq!(retrieved.is_dev_dependency, Some(true));
828    }
829
830    #[test]
831    fn test_migrate_v1_0_to_v1_1() {
832        use rusqlite::Connection;
833
834        // Create a temporary directory for the test
835        let temp_dir = tempfile::tempdir().unwrap();
836        let db_path = temp_dir.path().join("test_partition.db");
837
838        // Create a v1.0 partition manually (old schema without new columns)
839        {
840            let conn = Connection::open(&db_path).unwrap();
841
842            // Create v1.0 schema (without version_spec and is_dev_dependency)
843            conn.execute(
844                r#"
845                CREATE TABLE nodes (
846                    id TEXT PRIMARY KEY NOT NULL,
847                    name TEXT NOT NULL,
848                    node_type TEXT NOT NULL,
849                    kind TEXT,
850                    subtype TEXT,
851                    file TEXT NOT NULL,
852                    line INTEGER NOT NULL,
853                    end_line INTEGER NOT NULL,
854                    text TEXT,
855                    hash TEXT,
856                    metadata_json TEXT
857                )
858                "#,
859                [],
860            )
861            .unwrap();
862
863            conn.execute(
864                r#"
865                CREATE TABLE edges (
866                    id INTEGER PRIMARY KEY AUTOINCREMENT,
867                    source TEXT NOT NULL,
868                    target TEXT NOT NULL,
869                    edge_type TEXT NOT NULL,
870                    ref_line INTEGER,
871                    ident TEXT,
872                    UNIQUE(source, target, edge_type, ref_line)
873                )
874                "#,
875                [],
876            )
877            .unwrap();
878
879            conn.execute(
880                r#"
881                CREATE TABLE partition_metadata (
882                    key TEXT PRIMARY KEY NOT NULL,
883                    value TEXT NOT NULL
884                )
885                "#,
886                [],
887            )
888            .unwrap();
889
890            // Set v1.0 schema version
891            conn.execute(
892                "INSERT INTO partition_metadata (key, value) VALUES ('schema_version', '1.0')",
893                [],
894            )
895            .unwrap();
896
897            // Insert an edge with v1.0 schema
898            conn.execute(
899                "INSERT INTO edges (source, target, edge_type, ref_line, ident) VALUES (?1, ?2, ?3, ?4, ?5)",
900                params!["a", "b", "USES", 10i64, "test_ident"],
901            )
902            .unwrap();
903        }
904
905        // Open the partition with PartitionConnection - this should trigger migration
906        let conn = PartitionConnection::open(&db_path, "test_partition").unwrap();
907
908        // Verify schema version is now 1.1
909        let version = conn.get_metadata("schema_version").unwrap();
910        assert_eq!(version, Some("1.1".to_string()));
911
912        // Verify old edges can still be read (new columns should be None)
913        let edges = conn.query_edges_by_source("a").unwrap();
914        assert_eq!(edges.len(), 1);
915        assert_eq!(edges[0].source, "a");
916        assert_eq!(edges[0].target, "b");
917        assert_eq!(edges[0].edge_type, EdgeType::Uses);
918        assert_eq!(edges[0].ref_line, Some(10));
919        assert_eq!(edges[0].ident, Some("test_ident".to_string()));
920        assert_eq!(edges[0].version_spec, None); // Backward compat
921        assert_eq!(edges[0].is_dev_dependency, None); // Backward compat
922
923        // Insert a new edge with v1.1 schema
924        let new_edge = Edge {
925            source: "c".to_string(),
926            target: "d".to_string(),
927            edge_type: EdgeType::DependsOn,
928            ref_line: None,
929            ident: Some("dep".to_string()),
930            version_spec: Some(">=2.0".to_string()),
931            is_dev_dependency: Some(false),
932        };
933        conn.insert_edge(&new_edge).unwrap();
934
935        // Verify the new edge is stored with all fields
936        let new_edges = conn.query_edges_by_source("c").unwrap();
937        assert_eq!(new_edges.len(), 1);
938        assert_eq!(new_edges[0].version_spec, Some(">=2.0".to_string()));
939        assert_eq!(new_edges[0].is_dev_dependency, Some(false));
940    }
941}
codeprysm_core/lazy/partition.rs

codeprysm_core/lazy/
partition.rs
