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/// Embedding model choices exposed through `--model`.
9///
10/// Currently only `multilingual-e5-small` is supported. Additional variants
11/// will be added here as new models are integrated; the `value_enum` derive
12/// ensures the CLI rejects unknown strings at parse time rather than at runtime.
13#[derive(Copy, Clone, Debug, PartialEq, Eq, clap::ValueEnum)]
14pub enum EmbeddingModelChoice {
15    #[value(name = "multilingual-e5-small")]
16    MultilingualE5Small,
17}
18
19#[derive(clap::Args)]
20pub struct InitArgs {
21    /// Path to graphrag.sqlite. Defaults to `./graphrag.sqlite` in the current directory.
22    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
23    pub db: Option<String>,
24    /// Embedding model identifier. Currently only `multilingual-e5-small` is supported.
25    /// Reserved for future multi-model support; safe to omit.
26    #[arg(long, value_enum)]
27    pub model: Option<EmbeddingModelChoice>,
28    /// Force re-initialization, overwriting any existing schema metadata.
29    /// Use only when the schema is corrupted; loses configuration but preserves data.
30    #[arg(long)]
31    pub force: bool,
32    /// Initial namespace to resolve. Aligned with bilingual docs that mention `init --namespace`.
33    /// When provided, overrides `SQLITE_GRAPHRAG_NAMESPACE`; otherwise resolves via env or fallback `global`.
34    #[arg(long)]
35    pub namespace: Option<String>,
36    #[arg(long, help = "No-op; JSON is always emitted on stdout")]
37    pub json: bool,
38}
39
40#[derive(Serialize)]
41struct InitResponse {
42    db_path: String,
43    schema_version: String,
44    model: String,
45    dim: usize,
46    /// Namespace ativo resolvido durante a inicialização, alinhado à doc bilíngue.
47    namespace: String,
48    status: String,
49    /// Tempo total de execução em milissegundos desde início do handler até serialização.
50    elapsed_ms: u64,
51}
52
53pub fn run(args: InitArgs) -> Result<(), AppError> {
54    let inicio = std::time::Instant::now();
55    let paths = AppPaths::resolve(args.db.as_deref())?;
56    paths.ensure_dirs()?;
57
58    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
59
60    let mut conn = open_rw(&paths.db)?;
61
62    apply_init_pragmas(&conn)?;
63
64    crate::migrations::runner()
65        .run(&mut conn)
66        .map_err(|e| AppError::Internal(anyhow::anyhow!("migration failed: {e}")))?;
67
68    conn.execute_batch(&format!(
69        "PRAGMA user_version = {};",
70        crate::constants::SCHEMA_USER_VERSION
71    ))?;
72
73    let schema_version = latest_schema_version(&conn)?;
74
75    conn.execute(
76        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('schema_version', ?1)",
77        rusqlite::params![schema_version],
78    )?;
79    conn.execute(
80        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('model', 'multilingual-e5-small')",
81        [],
82    )?;
83    conn.execute(
84        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('dim', '384')",
85        [],
86    )?;
87    conn.execute(
88        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('created_at', CAST(unixepoch() AS TEXT))",
89        [],
90    )?;
91    conn.execute(
92        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('sqlite-graphrag_version', ?1)",
93        rusqlite::params![crate::constants::SQLITE_GRAPHRAG_VERSION],
94    )?;
95    // Persist the resolved namespace so downstream tools can inspect it without re-resolving.
96    conn.execute(
97        "INSERT OR REPLACE INTO schema_meta (key, value) VALUES ('namespace_initial', ?1)",
98        rusqlite::params![namespace],
99    )?;
100
101    output::emit_progress_i18n(
102        "Initializing embedding model (may download on first run)...",
103        "Inicializando modelo de embedding (pode baixar na primeira execução)...",
104    );
105
106    let test_emb = crate::daemon::embed_passage_or_local(&paths.models, "smoke test")?;
107
108    output::emit_json(&InitResponse {
109        db_path: paths.db.display().to_string(),
110        schema_version,
111        model: "multilingual-e5-small".to_string(),
112        dim: test_emb.len(),
113        namespace,
114        status: "ok".to_string(),
115        elapsed_ms: inicio.elapsed().as_millis() as u64,
116    })?;
117
118    Ok(())
119}
120
121fn latest_schema_version(conn: &rusqlite::Connection) -> Result<String, AppError> {
122    match conn.query_row(
123        "SELECT version FROM refinery_schema_history ORDER BY version DESC LIMIT 1",
124        [],
125        |row| row.get::<_, i64>(0),
126    ) {
127        Ok(version) => Ok(version.to_string()),
128        Err(rusqlite::Error::QueryReturnedNoRows) => Ok("0".to_string()),
129        Err(err) => Err(AppError::Database(err)),
130    }
131}
132
133#[cfg(test)]
134mod testes {
135    use super::*;
136
137    #[test]
138    fn init_response_serializa_todos_campos() {
139        let resp = InitResponse {
140            db_path: "/tmp/test.sqlite".to_string(),
141            schema_version: "6".to_string(),
142            model: "multilingual-e5-small".to_string(),
143            dim: 384,
144            namespace: "global".to_string(),
145            status: "ok".to_string(),
146            elapsed_ms: 100,
147        };
148        let json = serde_json::to_value(&resp).expect("serialização falhou");
149        assert_eq!(json["db_path"], "/tmp/test.sqlite");
150        assert_eq!(json["schema_version"], "6");
151        assert_eq!(json["model"], "multilingual-e5-small");
152        assert_eq!(json["dim"], 384usize);
153        assert_eq!(json["namespace"], "global");
154        assert_eq!(json["status"], "ok");
155        assert!(json["elapsed_ms"].is_number());
156    }
157
158    #[test]
159    fn latest_schema_version_retorna_zero_para_banco_vazio() {
160        let conn = rusqlite::Connection::open_in_memory().expect("falha ao abrir banco em memória");
161        conn.execute_batch("CREATE TABLE refinery_schema_history (version INTEGER NOT NULL);")
162            .expect("falha ao criar tabela");
163
164        let versao = latest_schema_version(&conn).expect("latest_schema_version falhou");
165        assert_eq!(versao, "0", "banco vazio deve retornar schema_version '0'");
166    }
167
168    #[test]
169    fn latest_schema_version_retorna_versao_maxima() {
170        let conn = rusqlite::Connection::open_in_memory().expect("falha ao abrir banco em memória");
171        conn.execute_batch(
172            "CREATE TABLE refinery_schema_history (version INTEGER NOT NULL);
173             INSERT INTO refinery_schema_history VALUES (1);
174             INSERT INTO refinery_schema_history VALUES (3);
175             INSERT INTO refinery_schema_history VALUES (2);",
176        )
177        .expect("falha ao popular tabela");
178
179        let versao = latest_schema_version(&conn).expect("latest_schema_version falhou");
180        assert_eq!(versao, "3", "deve retornar a maior versão presente");
181    }
182
183    #[test]
184    fn init_response_dim_alinhado_com_constante() {
185        assert_eq!(
186            crate::constants::EMBEDDING_DIM,
187            384,
188            "dim deve estar alinhado com EMBEDDING_DIM=384"
189        );
190    }
191
192    #[test]
193    fn init_response_namespace_alinhado_com_schema() {
194        // Verify namespace field survives round-trip serialization with correct value.
195        let resp = InitResponse {
196            db_path: "/tmp/x.sqlite".to_string(),
197            schema_version: "6".to_string(),
198            model: "multilingual-e5-small".to_string(),
199            dim: 384,
200            namespace: "meu-projeto".to_string(),
201            status: "ok".to_string(),
202            elapsed_ms: 0,
203        };
204        let json = serde_json::to_value(&resp).expect("serialização falhou");
205        assert_eq!(json["namespace"], "meu-projeto");
206    }
207}
sqlite_graphrag/commands/init.rs

sqlite_graphrag/commands/
init.rs
