use crate::storage::schema::Storage;
use rusqlite::{OptionalExtension, Result as SqliteResult, params};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeRecord {
pub id: Option<i64>,
pub project_id: String,
pub file_path: String,
pub node_id: String,
pub symbol_name: String,
pub qualified_name: String,
pub language: String,
pub node_type: NodeType,
pub signature: Option<String>,
pub complexity: Option<i32>,
pub content_hash: String,
pub embedding: Option<Vec<u8>>,
pub byte_range_start: Option<i64>,
pub byte_range_end: Option<i64>,
pub embedding_format: Option<i32>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum NodeType {
Function,
Class,
Method,
Variable,
Module,
External,
FileSummary,
}
impl NodeType {
pub fn as_str(&self) -> &'static str {
match self {
NodeType::Function => "function",
NodeType::Class => "class",
NodeType::Method => "method",
NodeType::Variable => "variable",
NodeType::Module => "module",
NodeType::External => "external",
NodeType::FileSummary => "file_summary",
}
}
pub fn from_str_name(s: &str) -> Option<Self> {
match s {
"function" => Some(NodeType::Function),
"class" => Some(NodeType::Class),
"method" => Some(NodeType::Method),
"variable" => Some(NodeType::Variable),
"module" => Some(NodeType::Module),
"external" => Some(NodeType::External),
"file_summary" => Some(NodeType::FileSummary),
_ => None,
}
}
}
pub struct NodeStore<'a> {
storage: &'a mut Storage,
}
impl<'a> NodeStore<'a> {
pub fn new(storage: &'a mut Storage) -> Self {
Self { storage }
}
fn node_metadata_from_row(
row: &rusqlite::Row<'_>,
) -> SqliteResult<(
Option<i64>,
String,
String,
String,
String,
String,
String,
NodeType,
Option<String>,
Option<i32>,
)> {
Ok((
Some(row.get(0)?),
row.get(1)?,
row.get(2)?,
row.get(3)?,
row.get(4)?,
row.get(5)?,
row.get(6)?,
NodeType::from_str_name(&row.get::<_, String>(7)?).unwrap_or(NodeType::Function),
row.get(8)?,
row.get(9)?,
))
}
fn node_storage_from_row(
row: &rusqlite::Row<'_>,
) -> SqliteResult<(
String,
Option<Vec<u8>>,
Option<i64>,
Option<i64>,
Option<i32>,
)> {
Ok((
row.get(10)?,
row.get(11)?,
row.get(12)?,
row.get(13)?,
row.get(14)?,
))
}
fn node_from_row(row: &rusqlite::Row<'_>) -> SqliteResult<NodeRecord> {
let (
id,
project_id,
file_path,
node_id,
symbol_name,
qualified_name,
language,
node_type,
signature,
complexity,
) = Self::node_metadata_from_row(row)?;
let (content_hash, embedding, byte_range_start, byte_range_end, embedding_format) =
Self::node_storage_from_row(row)?;
Ok(NodeRecord {
id,
project_id,
file_path,
node_id,
symbol_name,
qualified_name,
language,
node_type,
signature,
complexity,
content_hash,
embedding,
byte_range_start,
byte_range_end,
embedding_format,
})
}
pub fn insert(&mut self, record: &NodeRecord) -> SqliteResult<i64> {
self.storage.conn().execute(
"INSERT INTO intel_nodes (project_id, file_path, node_id, symbol_name, qualified_name, language, node_type, signature, complexity, content_hash, embedding, byte_range_start, byte_range_end, created_at, updated_at, embedding_format)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
params![
record.project_id,
record.file_path,
record.node_id,
record.symbol_name,
record.qualified_name,
record.language,
record.node_type.as_str(),
record.signature,
record.complexity,
record.content_hash,
record.embedding.as_deref(),
record.byte_range_start,
record.byte_range_end,
chrono::Utc::now().timestamp(),
chrono::Utc::now().timestamp(),
record.embedding_format,
],
)?;
Ok(self.storage.conn().last_insert_rowid())
}
pub fn batch_insert(&mut self, records: &[NodeRecord]) -> SqliteResult<Vec<i64>> {
let tx = self.storage.conn_mut().transaction()?;
let mut ids = Vec::new();
for record in records {
tx.execute(
"INSERT INTO intel_nodes (project_id, file_path, node_id, symbol_name, qualified_name, language, node_type, signature, complexity, content_hash, embedding, byte_range_start, byte_range_end, created_at, updated_at, embedding_format)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
params![
record.project_id,
record.file_path,
record.node_id,
record.symbol_name,
record.qualified_name,
record.language,
record.node_type.as_str(),
record.signature,
record.complexity,
record.content_hash,
record.embedding.as_deref(),
record.byte_range_start,
record.byte_range_end,
chrono::Utc::now().timestamp(),
chrono::Utc::now().timestamp(),
record.embedding_format,
],
)?;
ids.push(tx.last_insert_rowid());
}
tx.commit()?;
Ok(ids)
}
pub fn get(&self, id: i64) -> SqliteResult<Option<NodeRecord>> {
let mut stmt = self.storage.conn().prepare(
"SELECT id, project_id, file_path, node_id, symbol_name, qualified_name, language, node_type, signature, complexity, content_hash, embedding, byte_range_start, byte_range_end, embedding_format
FROM intel_nodes WHERE id = ?1"
)?;
let result = stmt.query_row(params![id], Self::node_from_row);
result.optional()
}
pub fn find_by_hash(&self, hash: &str) -> SqliteResult<Option<NodeRecord>> {
let mut stmt = self.storage.conn().prepare(
"SELECT id, project_id, file_path, node_id, symbol_name, qualified_name, language, node_type, signature, complexity, content_hash, embedding, byte_range_start, byte_range_end, embedding_format
FROM intel_nodes WHERE content_hash = ?1"
)?;
let result = stmt.query_row(params![hash], Self::node_from_row);
result.optional()
}
pub fn get_by_file(&self, file_path: &str) -> SqliteResult<Vec<NodeRecord>> {
let mut stmt = self.storage.conn().prepare(
"SELECT id, project_id, file_path, node_id, symbol_name, qualified_name, language, node_type, signature, complexity, content_hash, embedding, byte_range_start, byte_range_end, embedding_format
FROM intel_nodes WHERE file_path = ?1"
)?;
let nodes = stmt
.query_map(params![file_path], Self::node_from_row)?
.collect::<SqliteResult<Vec<_>>>()?;
Ok(nodes)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::schema::Storage;
use tempfile::NamedTempFile;
#[test]
fn test_node_insert_and_get() {
let temp_file = NamedTempFile::new().unwrap();
let mut storage = Storage::open(temp_file.path()).unwrap();
let mut store = NodeStore::new(&mut storage);
let record = NodeRecord {
id: None,
project_id: "test_project".to_string(),
file_path: "test.py".to_string(),
node_id: "test_project:test_func".to_string(),
symbol_name: "test_func".to_string(),
qualified_name: "test_func".to_string(),
language: "python".to_string(),
node_type: NodeType::Function,
signature: Some("def test_func()".to_string()),
complexity: Some(5),
content_hash: "abc123".to_string(),
embedding: None,
byte_range_start: Some(0),
byte_range_end: Some(100),
embedding_format: None,
};
let id = store.insert(&record).unwrap();
assert!(id > 0);
let retrieved = store.get(id).unwrap().unwrap();
assert_eq!(retrieved.symbol_name, "test_func");
}
#[test]
fn test_find_by_hash() {
let temp_file = NamedTempFile::new().unwrap();
let mut storage = Storage::open(temp_file.path()).unwrap();
let mut store = NodeStore::new(&mut storage);
let record = NodeRecord {
id: None,
project_id: "test_project".to_string(),
file_path: "test.py".to_string(),
node_id: "test_project:test_func".to_string(),
symbol_name: "test_func".to_string(),
qualified_name: "test_func".to_string(),
language: "python".to_string(),
node_type: NodeType::Function,
signature: None,
complexity: None,
content_hash: "hash123".to_string(),
embedding: None,
byte_range_start: Some(0),
byte_range_end: Some(100),
embedding_format: None,
};
store.insert(&record).unwrap();
let found = store.find_by_hash("hash123").unwrap();
assert!(found.is_some());
assert_eq!(found.unwrap().symbol_name, "test_func");
}
}