Skip to main content

sqlite_graphrag/commands/
init.rs

1use crate::errors::AppError;
2use crate::output;
3use crate::paths::AppPaths;
4use crate::pragmas::apply_init_pragmas;
5use crate::storage::connection::open_rw;
6use serde::Serialize;
7
8#[derive(clap::Args)]
9pub struct InitArgs {
10    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
11    pub db: Option<String>,
12    #[arg(long)]
13    pub model: Option<String>,
14    #[arg(long)]
15    pub force: bool,
16    /// Namespace inicial a resolver. Alinhado à documentação bilíngue que prevê `init --namespace`.
17    /// Se fornecido, sobrepõe `SQLITE_GRAPHRAG_NAMESPACE`; caso contrário, resolve via env
18    /// ou fallback `global`.
19    #[arg(long)]
20    pub namespace: Option<String>,
21    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
22    pub json: bool,
23}
24
25#[derive(Serialize)]
26struct InitResponse {
27    db_path: String,
28    schema_version: String,
29    model: String,
30    dim: usize,
31    /// Namespace ativo resolvido durante a inicialização, alinhado à doc bilíngue.
32    namespace: String,
33    status: String,
34    /// Tempo total de execução em milissegundos desde início do handler até serialização.
35    elapsed_ms: u64,
36}
37
38pub fn run(args: InitArgs) -> Result<(), AppError> {
39    let inicio = std::time::Instant::now();
40    let paths = AppPaths::resolve(args.db.as_deref())?;
41    paths.ensure_dirs()?;
42
43    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
44
45    let mut conn = open_rw(&paths.db)?;
46
47    apply_init_pragmas(&conn)?;
48
49    crate::migrations::runner()
50        .run(&mut conn)
51        .map_err(|e| AppError::Internal(anyhow::anyhow!("migration failed: {e}")))?;
52
53    conn.execute_batch(&format!(
54        "PRAGMA user_version = {};",
55        crate::constants::SCHEMA_USER_VERSION
56    ))?;
57
58    let schema_version = latest_schema_version(&conn)?;
59
60    conn.execute(
61        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('schema_version', ?1)",
62        rusqlite::params![schema_version],
63    )?;
64    conn.execute(
65        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('model', 'multilingual-e5-small')",
66        [],
67    )?;
68    conn.execute(
69        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('dim', '384')",
70        [],
71    )?;
72    conn.execute(
73        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('created_at', CAST(unixepoch() AS TEXT))",
74        [],
75    )?;
76    conn.execute(
77        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('sqlite-graphrag_version', ?1)",
78        rusqlite::params![crate::constants::SQLITE_GRAPHRAG_VERSION],
79    )?;
80
81    output::emit_progress_i18n(
82        "Initializing embedding model (may download on first run)...",
83        "Inicializando modelo de embedding (pode baixar na primeira execução)...",
84    );
85
86    let test_emb = crate::daemon::embed_passage_or_local(&paths.models, "smoke test")?;
87
88    output::emit_json(&InitResponse {
89        db_path: paths.db.display().to_string(),
90        schema_version,
91        model: "multilingual-e5-small".to_string(),
92        dim: test_emb.len(),
93        namespace,
94        status: "ok".to_string(),
95        elapsed_ms: inicio.elapsed().as_millis() as u64,
96    })?;
97
98    Ok(())
99}
100
101fn latest_schema_version(conn: &rusqlite::Connection) -> Result<String, AppError> {
102    match conn.query_row(
103        "SELECT version FROM refinery_schema_history ORDER BY version DESC LIMIT 1",
104        [],
105        |row| row.get::<_, i64>(0),
106    ) {
107        Ok(version) => Ok(version.to_string()),
108        Err(rusqlite::Error::QueryReturnedNoRows) => Ok("0".to_string()),
109        Err(err) => Err(AppError::Database(err)),
110    }
111}
112
113#[cfg(test)]
114mod testes {
115    use super::*;
116
117    #[test]
118    fn init_response_serializa_todos_campos() {
119        let resp = InitResponse {
120            db_path: "/tmp/test.sqlite".to_string(),
121            schema_version: "5".to_string(),
122            model: "multilingual-e5-small".to_string(),
123            dim: 384,
124            namespace: "global".to_string(),
125            status: "ok".to_string(),
126            elapsed_ms: 100,
127        };
128        let json = serde_json::to_value(&resp).expect("serialização falhou");
129        assert_eq!(json["db_path"], "/tmp/test.sqlite");
130        assert_eq!(json["schema_version"], "5");
131        assert_eq!(json["model"], "multilingual-e5-small");
132        assert_eq!(json["dim"], 384usize);
133        assert_eq!(json["namespace"], "global");
134        assert_eq!(json["status"], "ok");
135        assert!(json["elapsed_ms"].is_number());
136    }
137
138    #[test]
139    fn latest_schema_version_retorna_zero_para_banco_vazio() {
140        let conn = rusqlite::Connection::open_in_memory().expect("falha ao abrir banco em memória");
141        conn.execute_batch("CREATE TABLE refinery_schema_history (version INTEGER NOT NULL);")
142            .expect("falha ao criar tabela");
143
144        let versao = latest_schema_version(&conn).expect("latest_schema_version falhou");
145        assert_eq!(versao, "0", "banco vazio deve retornar schema_version '0'");
146    }
147
148    #[test]
149    fn latest_schema_version_retorna_versao_maxima() {
150        let conn = rusqlite::Connection::open_in_memory().expect("falha ao abrir banco em memória");
151        conn.execute_batch(
152            "CREATE TABLE refinery_schema_history (version INTEGER NOT NULL);
153             INSERT INTO refinery_schema_history VALUES (1);
154             INSERT INTO refinery_schema_history VALUES (3);
155             INSERT INTO refinery_schema_history VALUES (2);",
156        )
157        .expect("falha ao popular tabela");
158
159        let versao = latest_schema_version(&conn).expect("latest_schema_version falhou");
160        assert_eq!(versao, "3", "deve retornar a maior versão presente");
161    }
162
163    #[test]
164    fn init_response_dim_alinhado_com_constante() {
165        assert_eq!(
166            crate::constants::EMBEDDING_DIM,
167            384,
168            "dim deve estar alinhado com EMBEDDING_DIM=384"
169        );
170    }
171}