sqlite_graphrag/storage/
pending_embeddings.rs

1//! GAP-005 (v1.0.82): DAO for the `pending_embeddings` table.
2//!
3//! Queue of memories persisted with a NULL embedding for later reprocessing
4//! via `embedding retry <PENDING_ID>` or `enrich --operation re-embed`.
5
6use rusqlite::{params, Connection};
7
8use crate::errors::AppError;
9
10/// Pending embedding status.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
12#[serde(rename_all = "snake_case")]
13pub enum PendingEmbeddingStatus {
14    /// Pending variant.
15    Pending,
16    /// In progress variant.
17    InProgress,
18    /// Done variant.
19    Done,
20    /// Abandoned variant.
21    Abandoned,
22}
23
24impl PendingEmbeddingStatus {
25    /// Return the canonical string representation.
26    pub fn as_str(&self) -> &'static str {
27        match self {
28            Self::Pending => "pending",
29            Self::InProgress => "in_progress",
30            Self::Done => "done",
31            Self::Abandoned => "abandoned",
32        }
33    }
34}
35
36/// Pending embedding.
37#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
38pub struct PendingEmbedding {
39    /// Pending ID.
40    pub pending_id: i64,
41    /// Memory identifier.
42    pub memory_id: i64,
43    /// Namespace scope.
44    pub namespace: String,
45    /// Name of this item.
46    pub name: String,
47    /// Backend chain.
48    pub backend_chain: String,
49    /// Last error.
50    pub last_error: Option<String>,
51    /// Last exit code.
52    pub last_exit_code: Option<i32>,
53    /// Last stderr tail.
54    pub last_stderr_tail: Option<String>,
55    /// Attempt count.
56    pub attempt_count: i32,
57    /// Status value.
58    pub status: PendingEmbeddingStatus,
59    /// Creation timestamp.
60    pub created_at: i64,
61    /// Last-update timestamp.
62    pub updated_at: i64,
63}
64
65/// Inserts a new `pending_embeddings` entry with status `pending`.
66// One parameter per column of the `pending_embeddings` row this writes: the
67// arity IS the schema, and a struct here would be that row spelled a second time.
68#[allow(clippy::too_many_arguments)]
69pub fn insert(
70    conn: &Connection,
71    memory_id: i64,
72    namespace: &str,
73    name: &str,
74    backend_chain: &str,
75    last_error: Option<&str>,
76    last_exit_code: Option<i32>,
77    last_stderr_tail: Option<&str>,
78) -> Result<i64, AppError> {
79    conn.execute(
80        "INSERT INTO pending_embeddings
81            (memory_id, namespace, name, backend_chain, last_error,
82             last_exit_code, last_stderr_tail, attempt_count, status)
83         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, 0, 'pending')",
84        params![
85            memory_id,
86            namespace,
87            name,
88            backend_chain,
89            last_error,
90            last_exit_code,
91            last_stderr_tail,
92        ],
93    )?;
94    Ok(conn.last_insert_rowid())
95}
96
97/// Update status.
98pub fn update_status(
99    conn: &Connection,
100    pending_id: i64,
101    status: PendingEmbeddingStatus,
102    last_error: Option<&str>,
103    last_exit_code: Option<i32>,
104    last_stderr_tail: Option<&str>,
105) -> Result<(), AppError> {
106    conn.execute(
107        "UPDATE pending_embeddings
108         SET status = ?1,
109             last_error = COALESCE(?2, last_error),
110             last_exit_code = COALESCE(?3, last_exit_code),
111             last_stderr_tail = COALESCE(?4, last_stderr_tail),
112             attempt_count = attempt_count + 1,
113             updated_at = unixepoch()
114         WHERE pending_id = ?5",
115        params![
116            status.as_str(),
117            last_error,
118            last_exit_code,
119            last_stderr_tail,
120            pending_id
121        ],
122    )?;
123    Ok(())
124}
125
126/// Counts the entries [`list_by_status`] would page through.
127///
128/// GAP-SG-201: `embedding list --limit N` is a page of a countable set, and the
129/// output surface cannot tell a page from a corpus without the size of the
130/// corpus. The `WHERE` clause mirrors [`list_by_status`] exactly.
131///
132/// # Errors
133/// Returns [`AppError`] when the query fails.
134pub fn count_by_status(
135    conn: &Connection,
136    status: PendingEmbeddingStatus,
137) -> Result<usize, AppError> {
138    let count: i64 = conn.query_row(
139        "SELECT COUNT(*) FROM pending_embeddings WHERE status = ?1",
140        params![status.as_str()],
141        |row| row.get(0),
142    )?;
143    Ok(usize::try_from(count).unwrap_or(0))
144}
145
146/// List by status.
147pub fn list_by_status(
148    conn: &Connection,
149    status: PendingEmbeddingStatus,
150    limit: usize,
151) -> Result<Vec<PendingEmbedding>, AppError> {
152    let mut stmt = conn.prepare(
153        "SELECT pending_id, memory_id, namespace, name, backend_chain,
154                last_error, last_exit_code, last_stderr_tail,
155                attempt_count, status, created_at, updated_at
156         FROM pending_embeddings
157         WHERE status = ?1
158         ORDER BY updated_at ASC
159         LIMIT ?2",
160    )?;
161    let rows = stmt.query_map(params![status.as_str(), limit as i64], |row| {
162        Ok(PendingEmbedding {
163            pending_id: row.get(0)?,
164            memory_id: row.get(1)?,
165            namespace: row.get(2)?,
166            name: row.get(3)?,
167            backend_chain: row.get(4)?,
168            last_error: row.get(5)?,
169            last_exit_code: row.get(6)?,
170            last_stderr_tail: row.get(7)?,
171            attempt_count: row.get(8)?,
172            status: parse_status(&row.get::<_, String>(9)?).map_err(|e| -> rusqlite::Error {
173                rusqlite::Error::FromSqlConversionFailure(
174                    9,
175                    rusqlite::types::Type::Text,
176                    Box::new(std::io::Error::other(e.to_string())),
177                )
178            })?,
179            created_at: row.get(10)?,
180            updated_at: row.get(11)?,
181        })
182    })?;
183    let mut out = Vec::new();
184    for row in rows {
185        out.push(row?);
186    }
187    Ok(out)
188}
189
190/// Abandon.
191pub fn abandon(conn: &Connection, pending_id: i64) -> Result<(), AppError> {
192    update_status(
193        conn,
194        pending_id,
195        PendingEmbeddingStatus::Abandoned,
196        None,
197        None,
198        None,
199    )
200}
201
202/// Delete.
203pub fn delete(conn: &Connection, pending_id: i64) -> Result<(), AppError> {
204    conn.execute(
205        "DELETE FROM pending_embeddings WHERE pending_id = ?1",
206        params![pending_id],
207    )?;
208    Ok(())
209}
210
211fn parse_status(s: &str) -> Result<PendingEmbeddingStatus, AppError> {
212    match s {
213        "pending" => Ok(PendingEmbeddingStatus::Pending),
214        "in_progress" => Ok(PendingEmbeddingStatus::InProgress),
215        "done" => Ok(PendingEmbeddingStatus::Done),
216        "abandoned" => Ok(PendingEmbeddingStatus::Abandoned),
217        other => Err(AppError::Validation(
218            crate::i18n::validation::unknown_pending_embeddings_status(other),
219        )),
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use rusqlite::Connection;
227
228    fn fresh_db() -> Connection {
229        let mut conn = Connection::open_in_memory().expect("in-memory db");
230        conn.execute_batch("PRAGMA foreign_keys = ON;")
231            .expect("pragma");
232        crate::migrations::runner()
233            .run(&mut conn)
234            .expect("migrations apply");
235        conn
236    }
237
238    fn insert_test_memory(conn: &Connection, name: &str) -> i64 {
239        conn.execute(
240            "INSERT INTO memories (name, namespace, type, description, body, body_hash, source)
241             VALUES (?1, 'global', 'note', 'desc', 'body', 'h', 'agent')",
242            params![name],
243        )
244        .unwrap();
245        conn.last_insert_rowid()
246    }
247
248    #[test]
249    fn insert_records_pending_with_full_diagnostics() {
250        let conn = fresh_db();
251        let mid = insert_test_memory(&conn, "p");
252        let id = insert(
253            &conn,
254            mid,
255            "global",
256            "p",
257            "openrouter,none",
258            Some("exit 137 SIGKILL"),
259            Some(137),
260            Some("OOM killed by kernel"),
261        )
262        .unwrap();
263        let p = list_by_status(&conn, PendingEmbeddingStatus::Pending, 10)
264            .unwrap()
265            .into_iter()
266            .find(|p| p.pending_id == id)
267            .expect("pending found");
268        assert_eq!(p.backend_chain, "openrouter,none");
269        assert_eq!(p.last_exit_code, Some(137));
270        assert_eq!(p.last_stderr_tail.as_deref(), Some("OOM killed by kernel"));
271    }
272
273    #[test]
274    fn update_status_increments_attempt_count() {
275        let conn = fresh_db();
276        let mid = insert_test_memory(&conn, "p");
277        let id = insert(&conn, mid, "global", "p", "openrouter", None, None, None).unwrap();
278        update_status(
279            &conn,
280            id,
281            PendingEmbeddingStatus::InProgress,
282            None,
283            None,
284            None,
285        )
286        .unwrap();
287        let p = list_by_status(&conn, PendingEmbeddingStatus::InProgress, 10)
288            .unwrap()
289            .into_iter()
290            .find(|p| p.pending_id == id)
291            .expect("found");
292        assert_eq!(p.attempt_count, 1);
293    }
294
295    #[test]
296    fn abandon_sets_status() {
297        let conn = fresh_db();
298        let mid = insert_test_memory(&conn, "p");
299        let id = insert(&conn, mid, "global", "p", "openrouter", None, None, None).unwrap();
300        abandon(&conn, id).unwrap();
301        let abandoned = list_by_status(&conn, PendingEmbeddingStatus::Abandoned, 10).unwrap();
302        assert!(abandoned.iter().any(|p| p.pending_id == id));
303    }
304
305    #[test]
306    fn delete_removes_row() {
307        let conn = fresh_db();
308        let mid = insert_test_memory(&conn, "p");
309        let id = insert(&conn, mid, "global", "p", "openrouter", None, None, None).unwrap();
310        delete(&conn, id).unwrap();
311        let pending = list_by_status(&conn, PendingEmbeddingStatus::Pending, 10).unwrap();
312        assert!(pending.iter().all(|p| p.pending_id != id));
313    }
314}
sqlite_graphrag/storage/pending_embeddings.rs

sqlite_graphrag/storage/
pending_embeddings.rs
