sqlite_graphrag/commands/
optimize.rs

1//! Handler for the `optimize` CLI subcommand.
2
3use crate::errors::AppError;
4use crate::output;
5use crate::paths::AppPaths;
6use crate::storage::connection::open_rw;
7use serde::Serialize;
8
9#[derive(clap::Args)]
10#[command(after_long_help = "EXAMPLES:\n  \
11    # Run PRAGMA optimize on the default database\n  \
12    sqlite-graphrag optimize\n\n  \
13    # Optimize a database at a custom path\n  \
14    sqlite-graphrag optimize --db /path/to/graphrag.sqlite\n\n  \
15    # Optimize via SQLITE_GRAPHRAG_DB_PATH env var\n  \
16    SQLITE_GRAPHRAG_DB_PATH=/data/graphrag.sqlite sqlite-graphrag optimize")]
17pub struct OptimizeArgs {
18    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
19    pub json: bool,
20    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
21    pub db: Option<String>,
22}
23
24#[derive(Serialize)]
25struct OptimizeResponse {
26    db_path: String,
27    status: String,
28    /// Total execution time in milliseconds from handler start to serialisation.
29    elapsed_ms: u64,
30}
31
32pub fn run(args: OptimizeArgs) -> Result<(), AppError> {
33    let inicio = std::time::Instant::now();
34    let paths = AppPaths::resolve(args.db.as_deref())?;
35
36    crate::storage::connection::ensure_db_ready(&paths)?;
37
38    let conn = open_rw(&paths.db)?;
39    conn.execute_batch("PRAGMA optimize;")?;
40
41    output::emit_json(&OptimizeResponse {
42        db_path: paths.db.display().to_string(),
43        status: "ok".to_string(),
44        elapsed_ms: inicio.elapsed().as_millis() as u64,
45    })?;
46
47    Ok(())
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53    use serial_test::serial;
54    use tempfile::TempDir;
55
56    #[test]
57    fn optimize_response_serializa_campos_obrigatorios() {
58        let resp = OptimizeResponse {
59            db_path: "/tmp/graphrag.sqlite".to_string(),
60            status: "ok".to_string(),
61            elapsed_ms: 5,
62        };
63        let json = serde_json::to_value(&resp).unwrap();
64        assert_eq!(json["status"], "ok");
65        assert_eq!(json["db_path"], "/tmp/graphrag.sqlite");
66        assert_eq!(json["elapsed_ms"], 5);
67    }
68
69    #[test]
70    #[serial]
71    fn optimize_auto_inits_when_db_missing() {
72        let dir = TempDir::new().unwrap();
73        let db_path = dir.path().join("missing.sqlite");
74        // SAFETY: `#[serial]` guarantees single-threaded execution.
75        unsafe {
76            std::env::set_var("SQLITE_GRAPHRAG_DB_PATH", db_path.to_str().unwrap());
77            std::env::set_var("LOG_LEVEL", "error");
78        }
79
80        let args = OptimizeArgs {
81            json: false,
82            db: Some(db_path.to_string_lossy().to_string()),
83        };
84        let result = run(args);
85        assert!(
86            result.is_ok(),
87            "auto-init must succeed and PRAGMA optimize must run on the fresh database, got {result:?}"
88        );
89        assert!(
90            db_path.exists(),
91            "auto-init must create the database file at {}",
92            db_path.display()
93        );
94        // SAFETY: `#[serial]` guarantees single-threaded execution.
95        unsafe {
96            std::env::remove_var("SQLITE_GRAPHRAG_DB_PATH");
97            std::env::remove_var("LOG_LEVEL");
98        }
99    }
100
101    #[test]
102    fn optimize_response_status_ok_fixo() {
103        let resp = OptimizeResponse {
104            db_path: "/qualquer/caminho".to_string(),
105            status: "ok".to_string(),
106            elapsed_ms: 0,
107        };
108        let json = serde_json::to_value(&resp).unwrap();
109        assert_eq!(json["status"], "ok", "status deve ser sempre 'ok'");
110    }
111}
sqlite_graphrag/commands/optimize.rs

sqlite_graphrag/commands/
optimize.rs
