tempo_cli/db/
connection.rs

1use anyhow::Result;
2use rusqlite::{Connection, OpenFlags, OptionalExtension};
3use std::path::Path;
4
5pub struct Database {
6    pub connection: Connection,
7}
8
9impl Database {
10    pub fn new(db_path: &Path) -> Result<Self> {
11        // Create parent directory if it doesn't exist
12        if let Some(parent) = db_path.parent() {
13            std::fs::create_dir_all(parent)?;
14        }
15
16        let connection = Connection::open_with_flags(
17            db_path,
18            OpenFlags::SQLITE_OPEN_READ_WRITE
19                | OpenFlags::SQLITE_OPEN_CREATE
20                | OpenFlags::SQLITE_OPEN_NO_MUTEX,
21        )?;
22
23        // Enable foreign key constraints
24        connection.pragma_update(None, "foreign_keys", "ON")?;
25        
26        // Set WAL mode for better concurrent access
27        connection.pragma_update(None, "journal_mode", "WAL")?;
28        
29        // Set synchronous to NORMAL for better performance
30        connection.pragma_update(None, "synchronous", "NORMAL")?;
31        
32        // Set cache size (negative value means KB)
33        connection.pragma_update(None, "cache_size", "-64000")?;
34
35        let db = Self { connection };
36        
37        // Run migrations automatically
38        crate::db::migrations::run_migrations(&db.connection)?;
39
40        Ok(db)
41    }
42
43    pub fn in_memory() -> Result<Self> {
44        let connection = Connection::open_in_memory()?;
45        connection.execute("PRAGMA foreign_keys = ON", [])?;
46        Ok(Self { connection })
47    }
48
49    pub fn backup_to(&self, backup_path: &Path) -> Result<()> {
50        let mut backup_conn = Connection::open(backup_path)?;
51        let backup = rusqlite::backup::Backup::new(&self.connection, &mut backup_conn)?;
52        backup.run_to_completion(5, std::time::Duration::from_millis(250), None)?;
53        Ok(())
54    }
55
56    pub fn vacuum(&self) -> Result<()> {
57        self.connection.execute("VACUUM", [])?;
58        Ok(())
59    }
60
61    pub fn analyze(&self) -> Result<()> {
62        self.connection.execute("ANALYZE", [])?;
63        Ok(())
64    }
65
66    pub fn get_schema_version(&self) -> Result<Option<i32>> {
67        let mut stmt = self.connection.prepare(
68            "SELECT version FROM schema_version ORDER BY version DESC LIMIT 1"
69        )?;
70        
71        let version = stmt.query_row([], |row| {
72            Ok(row.get::<_, i32>(0)?)
73        }).optional()?;
74        
75        Ok(version)
76    }
77}