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