Skip to main content

sqlite_graphrag/commands/
optimize.rs

1//! Handler for the `optimize` CLI subcommand.
2
3use crate::errors::AppError;
4use crate::i18n::errors_msg;
5use crate::output;
6use crate::paths::AppPaths;
7use crate::storage::connection::open_rw;
8use serde::Serialize;
9
10#[derive(clap::Args)]
11pub struct OptimizeArgs {
12    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
13    pub json: bool,
14    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
15    pub db: Option<String>,
16}
17
18#[derive(Serialize)]
19struct OptimizeResponse {
20    db_path: String,
21    status: String,
22    /// Total execution time in milliseconds from handler start to serialisation.
23    elapsed_ms: u64,
24}
25
26pub fn run(args: OptimizeArgs) -> Result<(), AppError> {
27    let inicio = std::time::Instant::now();
28    let paths = AppPaths::resolve(args.db.as_deref())?;
29
30    if !paths.db.exists() {
31        return Err(AppError::NotFound(errors_msg::database_not_found(
32            &paths.db.display().to_string(),
33        )));
34    }
35
36    let conn = open_rw(&paths.db)?;
37    conn.execute_batch("PRAGMA optimize;")?;
38
39    output::emit_json(&OptimizeResponse {
40        db_path: paths.db.display().to_string(),
41        status: "ok".to_string(),
42        elapsed_ms: inicio.elapsed().as_millis() as u64,
43    })?;
44
45    Ok(())
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51    use serial_test::serial;
52    use tempfile::TempDir;
53
54    #[test]
55    fn optimize_response_serializa_campos_obrigatorios() {
56        let resp = OptimizeResponse {
57            db_path: "/tmp/graphrag.sqlite".to_string(),
58            status: "ok".to_string(),
59            elapsed_ms: 5,
60        };
61        let json = serde_json::to_value(&resp).unwrap();
62        assert_eq!(json["status"], "ok");
63        assert_eq!(json["db_path"], "/tmp/graphrag.sqlite");
64        assert_eq!(json["elapsed_ms"], 5);
65    }
66
67    #[test]
68    #[serial]
69    fn optimize_retorna_not_found_quando_db_ausente() {
70        let dir = TempDir::new().unwrap();
71        let db_path = dir.path().join("inexistente.sqlite");
72        std::env::set_var("SQLITE_GRAPHRAG_DB_PATH", db_path.to_str().unwrap());
73        std::env::set_var("LOG_LEVEL", "error");
74
75        let args = OptimizeArgs {
76            json: false,
77            db: Some(db_path.to_string_lossy().to_string()),
78        };
79        let resultado = run(args);
80        assert!(resultado.is_err(), "deve falhar quando db não existe");
81        match resultado.unwrap_err() {
82            AppError::NotFound(_) => {}
83            outro => unreachable!("esperava NotFound, obteve: {outro:?}"),
84        }
85        std::env::remove_var("SQLITE_GRAPHRAG_DB_PATH");
86        std::env::remove_var("LOG_LEVEL");
87    }
88
89    #[test]
90    fn optimize_response_status_ok_fixo() {
91        let resp = OptimizeResponse {
92            db_path: "/qualquer/caminho".to_string(),
93            status: "ok".to_string(),
94            elapsed_ms: 0,
95        };
96        let json = serde_json::to_value(&resp).unwrap();
97        assert_eq!(json["status"], "ok", "status deve ser sempre 'ok'");
98    }
99}