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    #[arg(long, default_value_t = false, help = "Skip FTS5 index rebuild")]
23    pub skip_fts: bool,
24}
25
26#[derive(Serialize)]
27struct OptimizeResponse {
28    db_path: String,
29    status: String,
30    /// True when the FTS5 index was rebuilt during this optimize run.
31    fts_rebuilt: bool,
32    /// Total execution time in milliseconds from handler start to serialisation.
33    elapsed_ms: u64,
34}
35
36pub fn run(args: OptimizeArgs) -> Result<(), AppError> {
37    let inicio = std::time::Instant::now();
38    let paths = AppPaths::resolve(args.db.as_deref())?;
39
40    crate::storage::connection::ensure_db_ready(&paths)?;
41
42    let conn = open_rw(&paths.db)?;
43    conn.execute_batch("PRAGMA optimize;")?;
44
45    let fts_rebuilt = if !args.skip_fts {
46        conn.execute_batch("INSERT INTO fts_memories(fts_memories) VALUES('rebuild');")
47            .is_ok()
48    } else {
49        false
50    };
51
52    output::emit_json(&OptimizeResponse {
53        db_path: paths.db.display().to_string(),
54        status: "ok".to_string(),
55        fts_rebuilt,
56        elapsed_ms: inicio.elapsed().as_millis() as u64,
57    })?;
58
59    Ok(())
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use serial_test::serial;
66    use tempfile::TempDir;
67
68    #[test]
69    fn optimize_response_serializes_required_fields() {
70        let resp = OptimizeResponse {
71            db_path: "/tmp/graphrag.sqlite".to_string(),
72            status: "ok".to_string(),
73            fts_rebuilt: false,
74            elapsed_ms: 5,
75        };
76        let json = serde_json::to_value(&resp).unwrap();
77        assert_eq!(json["status"], "ok");
78        assert_eq!(json["db_path"], "/tmp/graphrag.sqlite");
79        assert_eq!(json["elapsed_ms"], 5);
80    }
81
82    #[test]
83    #[serial]
84    fn optimize_auto_inits_when_db_missing() {
85        let dir = TempDir::new().unwrap();
86        let db_path = dir.path().join("missing.sqlite");
87        // SAFETY: `#[serial]` guarantees single-threaded execution.
88        unsafe {
89            std::env::set_var("SQLITE_GRAPHRAG_DB_PATH", db_path.to_str().unwrap());
90            std::env::set_var("LOG_LEVEL", "error");
91        }
92
93        let args = OptimizeArgs {
94            json: false,
95            db: Some(db_path.to_string_lossy().to_string()),
96            skip_fts: false,
97        };
98        let result = run(args);
99        assert!(
100            result.is_ok(),
101            "auto-init must succeed and PRAGMA optimize must run on the fresh database, got {result:?}"
102        );
103        assert!(
104            db_path.exists(),
105            "auto-init must create the database file at {}",
106            db_path.display()
107        );
108        // SAFETY: `#[serial]` guarantees single-threaded execution.
109        unsafe {
110            std::env::remove_var("SQLITE_GRAPHRAG_DB_PATH");
111            std::env::remove_var("LOG_LEVEL");
112        }
113    }
114
115    #[test]
116    fn optimize_response_status_ok_fixo() {
117        let resp = OptimizeResponse {
118            db_path: "/qualquer/caminho".to_string(),
119            status: "ok".to_string(),
120            fts_rebuilt: false,
121            elapsed_ms: 0,
122        };
123        let json = serde_json::to_value(&resp).unwrap();
124        assert_eq!(json["status"], "ok", "status deve ser sempre 'ok'");
125    }
126}
sqlite_graphrag/commands/optimize.rs

sqlite_graphrag/commands/
optimize.rs
