use crate::error::LlmError;
use rusqlite::Connection;
#[derive(Debug, Clone)]
pub struct CodeChunk {
pub id: i64,
pub file_path: String,
pub byte_start: u64,
pub byte_end: u64,
pub content: String,
pub content_hash: String,
pub symbol_name: Option<String>,
pub symbol_kind: Option<String>,
}
pub fn search_chunks_by_symbol_name(
conn: &Connection,
symbol_name: &str,
) -> Result<Vec<CodeChunk>, LlmError> {
let sql = r#"
SELECT id, file_path, byte_start, byte_end, content, content_hash, symbol_name, symbol_kind
FROM code_chunks
WHERE symbol_name = ?
"#;
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map([symbol_name], |row| {
Ok(CodeChunk {
id: row.get(0)?,
file_path: row.get(1)?,
byte_start: row.get(2)?,
byte_end: row.get(3)?,
content: row.get(4)?,
content_hash: row.get(5)?,
symbol_name: row.get(6)?,
symbol_kind: row.get(7)?,
})
})?;
let mut chunks = Vec::new();
for row in rows {
chunks.push(row?);
}
Ok(chunks)
}
pub fn search_chunks_by_span(
conn: &Connection,
file_path: &str,
byte_start: u64,
byte_end: u64,
) -> Result<Option<CodeChunk>, LlmError> {
let sql = r#"
SELECT id, file_path, byte_start, byte_end, content, content_hash, symbol_name, symbol_kind
FROM code_chunks
WHERE file_path = ? AND byte_start = ? AND byte_end = ?
LIMIT 1
"#;
let mut stmt = conn.prepare(sql)?;
let mut rows = stmt.query_map(
rusqlite::params![file_path, byte_start as i64, byte_end as i64],
|row| {
Ok(CodeChunk {
id: row.get(0)?,
file_path: row.get(1)?,
byte_start: row.get(2)?,
byte_end: row.get(3)?,
content: row.get(4)?,
content_hash: row.get(5)?,
symbol_name: row.get(6)?,
symbol_kind: row.get(7)?,
})
},
)?;
match rows.next() {
Some(Ok(chunk)) => Ok(Some(chunk)),
Some(Err(e)) => Err(LlmError::from(e)),
None => Ok(None),
}
}