nyx/src/database.rs
elipeter 0a62b6f40c Refactor database connection handling with connection pooling and parallel processing
- Introduced `r2d2` connection pooling for SQLite in `database.rs`.
- Updated `Indexer` to use pooled connections for improved concurrency.
- Replaced sequential processing with `rayon` for parallel file scanning.
- Added a `clear` method to `Indexer` for reindexing support.
- Enhanced database initialization with `init` and `from_pool` methods.
- Updated `Cargo.toml` and `Cargo.lock` to include `r2d2`, `r2d2_sqlite`, and new dependencies.
2025-06-17 20:45:33 +02:00

226 lines
7.5 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

pub mod index {
use rusqlite::{params, Connection, OpenFlags, OptionalExtension};
use std::fs;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::commands::scan::Diag;
use crate::patterns::Severity;
use r2d2_sqlite::{SqliteConnectionManager};
use std::ops::Deref;
use std::sync::Arc;
use r2d2::{Pool, PooledConnection};
/// DB schema (foreignkeys enabled).
const SCHEMA: &str = r#"
PRAGMA foreign_keys = ON;
CREATE TABLE IF NOT EXISTS files (
id INTEGER PRIMARY KEY AUTOINCREMENT,
project TEXT NOT NULL,
path TEXT NOT NULL,
hash BLOB NOT NULL,
mtime INTEGER NOT NULL,
scanned_at INTEGER NOT NULL,
UNIQUE(project, path)
);
CREATE TABLE IF NOT EXISTS issues (
file_id INTEGER NOT NULL
REFERENCES files(id)
ON DELETE CASCADE,
rule_id TEXT NOT NULL,
severity TEXT NOT NULL,
line INTEGER NOT NULL,
col INTEGER NOT NULL,
PRIMARY KEY (file_id, rule_id, line, col)
);
"#;
/// A single issue row, ready for insertion.
#[derive(Debug, Clone)]
pub struct IssueRow<'a> {
pub rule_id: &'a str,
pub severity: &'a str,
pub line: i64,
pub col: i64,
}
pub struct Indexer {
conn: PooledConnection<SqliteConnectionManager>,
project: String,
}
impl Indexer {
pub fn init(
database_path: &Path,
) -> Result<std::sync::Arc<Pool<SqliteConnectionManager>>, Box<dyn std::error::Error>> {
let flags = OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_FULL_MUTEX;
let manager = SqliteConnectionManager::file(&database_path).with_flags(flags);
let pool = Arc::new(Pool::new(manager)?);
{
let conn = pool.get()?;
conn.pragma_update(None, "journal_mode", &"WAL")?;
conn.execute_batch(SCHEMA)?;
}
Ok(pool)
}
pub fn from_pool(
project: &str,
pool: &Pool<SqliteConnectionManager>,
) -> Result<Self, Box<dyn std::error::Error>> {
let conn = pool.get()?;
Ok(Self { conn, project: project.to_owned() })
}
// helper so code below can treat PooledConnection like &Connection
fn c(&self) -> &Connection { self.conn.deref() }
/// Open (or create) the DB at `database_path` for the given project name.
// pub fn new(project: &str, database_path: &Path) -> Result<Self, Box<dyn std::error::Error>> {
// let conn = Connection::open(database_path)?;
// conn.pragma_update(None, "journal_mode", &"WAL")?;
// conn.execute_batch(SCHEMA)?;
// Ok(Self { conn, project: project.to_owned() })
// }
/// Return true when the file *content* or *mtime* changed since the last scan.
pub fn should_scan(&self, path: &Path) -> Result<bool, Box<dyn std::error::Error>> {
let meta = fs::metadata(path)?;
let mtime = meta.modified()?.duration_since(UNIX_EPOCH)?.as_secs() as i64;
let digest = Self::digest_file(path)?;
let row: Option<(Vec<u8>, i64)> = self
.conn
.query_row(
"SELECT hash, mtime FROM files WHERE project = ?1 AND path = ?2",
params![self.project, path.to_string_lossy()],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.optional()?;
Ok(match row {
Some((stored_hash, stored_mtime)) => stored_hash != digest || stored_mtime != mtime,
None => true,
})
}
/// Insert or update the `files` row and return its id.
pub fn upsert_file(&self, path: &Path) -> Result<i64, Box<dyn std::error::Error>> {
let meta = fs::metadata(path)?;
let mtime = meta.modified()?.duration_since(UNIX_EPOCH)?.as_secs() as i64;
let scanned_at = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() as i64;
let digest = Self::digest_file(path)?;
self.c().execute(
"INSERT INTO files (project, path, hash, mtime, scanned_at)
VALUES (?1, ?2, ?3, ?4, ?5)
ON CONFLICT(project,path) DO UPDATE
SET hash = excluded.hash,
mtime = excluded.mtime,
scanned_at = excluded.scanned_at",
params![self.project, path.to_string_lossy(), digest, mtime, scanned_at],
)?;
let id: i64 = self.c().query_row(
"SELECT id FROM files WHERE project = ?1 AND path = ?2",
params![self.project, path.to_string_lossy()],
|r| r.get(0),
)?;
Ok(id)
}
/// Replace all issues for `file_id` with the supplied set.
pub fn replace_issues<'a>(&mut self, file_id: i64, issues: impl IntoIterator<Item = IssueRow<'a>>)
-> Result<(), Box<dyn std::error::Error>> {
let tx = self.conn.transaction()?;
tx.execute("DELETE FROM issues WHERE file_id = ?", params![file_id])?;
{
let mut stmt = tx.prepare(
"INSERT INTO issues (file_id, rule_id, severity, line, col)
VALUES (?1, ?2, ?3, ?4, ?5)",
)?;
for iss in issues {
stmt.execute(params![file_id, iss.rule_id, iss.severity, iss.line, iss.col])?;
}
}
tx.commit()?;
Ok(())
}
/// Gets the issues for a specific file so we don't have to rescan
pub fn get_issues_from_file(
&self,
path: &Path,
) -> Result<Vec<Diag>, Box<dyn std::error::Error>> {
let file_id: i64 = self.c().query_row(
"SELECT id FROM files WHERE project = ?1 AND path = ?2",
params![self.project, path.to_string_lossy()],
|r| r.get(0),
)?;
let mut stmt = self.c().prepare(
"SELECT rule_id, severity, line, col
FROM issues
WHERE file_id = ?1",
)?;
let issue_iter = stmt.query_map([file_id], |row| {
let sev_str: String = row.get(1)?;
Ok(Diag {
path: path.to_string_lossy().to_string(),
id: row.get::<_, String>(0)?, // rule_id
line: row.get::<_, i64>(2)? as usize,
col: row.get::<_, i64>(3)? as usize,
severity: Severity::from_str(&sev_str).unwrap(),
})
})?;
Ok(issue_iter.filter_map(Result::ok).collect())
}
/// gets files from the database
pub fn get_files(&self, project: &str) -> Result<Vec<std::path::PathBuf>, Box<dyn std::error::Error>> {
let mut stmt = self.c().prepare(
"SELECT path
FROM files
WHERE project = ?1",
)?;
let file_iter = stmt.query_map([project], |row| row.get::<_, String>(0))?;
Ok(file_iter.map(|p| p.map(PathBuf::from)).collect::<Result<_, _>>()?)
}
/// Clears the tables to prep for a reindex
pub fn clear(&self) -> rusqlite::Result<()> {
self.c().execute_batch(
r#"
PRAGMA foreign_keys = OFF;
DROP TABLE IF EXISTS issues;
DROP TABLE IF EXISTS files;
PRAGMA foreign_keys = ON;
VACUUM;
"#,
)?;
self.c().execute_batch(SCHEMA)?;
Ok(())
}
fn digest_file(path: &Path) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
let mut hasher = blake3::Hasher::new();
let mut file = fs::File::open(path)?;
std::io::copy(&mut file, &mut hasher)?;
Ok(hasher.finalize().as_bytes().to_vec())
}
}
}