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 embedder = crate::embedder::get_embedder(&paths.models)?;
87    let test_emb = crate::embedder::embed_passage(embedder, "smoke test")?;
88
89    output::emit_json(&InitResponse {
90        db_path: paths.db.display().to_string(),
91        schema_version,
92        model: "multilingual-e5-small".to_string(),
93        dim: test_emb.len(),
94        namespace,
95        status: "ok".to_string(),
96        elapsed_ms: inicio.elapsed().as_millis() as u64,
97    })?;
98
99    Ok(())
100}
101
102fn latest_schema_version(conn: &rusqlite::Connection) -> Result<String, AppError> {
103    match conn.query_row(
104        "SELECT version FROM refinery_schema_history ORDER BY version DESC LIMIT 1",
105        [],
106        |row| row.get::<_, i64>(0),
107    ) {
108        Ok(version) => Ok(version.to_string()),
109        Err(rusqlite::Error::QueryReturnedNoRows) => Ok("0".to_string()),
110        Err(err) => Err(AppError::Database(err)),
111    }
112}
113
114#[cfg(test)]
115mod testes {
116    use super::*;
117
118    #[test]
119    fn init_response_serializa_todos_campos() {
120        let resp = InitResponse {
121            db_path: "/tmp/test.sqlite".to_string(),
122            schema_version: "5".to_string(),
123            model: "multilingual-e5-small".to_string(),
124            dim: 384,
125            namespace: "global".to_string(),
126            status: "ok".to_string(),
127            elapsed_ms: 100,
128        };
129        let json = serde_json::to_value(&resp).expect("serialização falhou");
130        assert_eq!(json["db_path"], "/tmp/test.sqlite");
131        assert_eq!(json["schema_version"], "5");
132        assert_eq!(json["model"], "multilingual-e5-small");
133        assert_eq!(json["dim"], 384usize);
134        assert_eq!(json["namespace"], "global");
135        assert_eq!(json["status"], "ok");
136        assert!(json["elapsed_ms"].is_number());
137    }
138
139    #[test]
140    fn latest_schema_version_retorna_zero_para_banco_vazio() {
141        let conn = rusqlite::Connection::open_in_memory().expect("falha ao abrir banco em memória");
142        conn.execute_batch("CREATE TABLE refinery_schema_history (version INTEGER NOT NULL);")
143            .expect("falha ao criar tabela");
144
145        let versao = latest_schema_version(&conn).expect("latest_schema_version falhou");
146        assert_eq!(versao, "0", "banco vazio deve retornar schema_version '0'");
147    }
148
149    #[test]
150    fn latest_schema_version_retorna_versao_maxima() {
151        let conn = rusqlite::Connection::open_in_memory().expect("falha ao abrir banco em memória");
152        conn.execute_batch(
153            "CREATE TABLE refinery_schema_history (version INTEGER NOT NULL);
154             INSERT INTO refinery_schema_history VALUES (1);
155             INSERT INTO refinery_schema_history VALUES (3);
156             INSERT INTO refinery_schema_history VALUES (2);",
157        )
158        .expect("falha ao popular tabela");
159
160        let versao = latest_schema_version(&conn).expect("latest_schema_version falhou");
161        assert_eq!(versao, "3", "deve retornar a maior versão presente");
162    }
163
164    #[test]
165    fn init_response_dim_alinhado_com_constante() {
166        assert_eq!(
167            crate::constants::EMBEDDING_DIM,
168            384,
169            "dim deve estar alinhado com EMBEDDING_DIM=384"
170        );
171    }
172}