sqlite_graphrag/storage/
pending_memories.rs

1//! GAP-001 (v1.0.82): DAO for the `pending_memories` table.
2//!
3//! Staged persistence with a resumable checkpoint. Lets `remember` resume
4//! from Stage B (embedding) without re-validating Stage A (parse + validate).
5//!
6//! Status transitions:
7//!   validated → embedding_in_progress → embedding_done → committed
8//!                                                    ↘ abandoned (manual cleanup)
9//!                                                    ↘ failed (max attempts reached)
10
11use rusqlite::{params, Connection};
12
13use crate::errors::AppError;
14
15/// Status enum of a pending entry. Maps 1:1 to the CHECK constraint
16/// of the `pending_memories` table.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum PendingStatus {
20    /// Validated variant.
21    Validated,
22    /// Embedding in progress variant.
23    EmbeddingInProgress,
24    /// Embedding done variant.
25    EmbeddingDone,
26    /// Committed variant.
27    Committed,
28    /// Abandoned variant.
29    Abandoned,
30    /// Failed variant.
31    Failed,
32}
33
34impl PendingStatus {
35    /// Return the canonical string representation.
36    pub fn as_str(&self) -> &'static str {
37        match self {
38            Self::Validated => "validated",
39            Self::EmbeddingInProgress => "embedding_in_progress",
40            Self::EmbeddingDone => "embedding_done",
41            Self::Committed => "committed",
42            Self::Abandoned => "abandoned",
43            Self::Failed => "failed",
44        }
45    }
46}
47
48/// Represents an entry in the `pending_memories` table.
49#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
50pub struct PendingMemory {
51    /// Pending ID.
52    pub pending_id: i64,
53    /// Name of this item.
54    pub name: String,
55    /// Namespace scope.
56    pub namespace: String,
57    /// Memory type classification.
58    pub memory_type: String,
59    /// Human-readable description.
60    pub description: Option<String>,
61    /// Full text body.
62    pub body: Vec<u8>,
63    /// Body hash.
64    pub body_hash: String,
65    /// Entities JSON.
66    pub entities_json: Option<String>,
67    /// Relationships JSON.
68    pub relationships_json: Option<String>,
69    /// Status value.
70    pub status: PendingStatus,
71    /// Embedding vector.
72    pub embedding: Option<Vec<u8>>,
73    /// Embedding dim.
74    pub embedding_dim: Option<i32>,
75    /// Attempt count.
76    pub attempt_count: i32,
77    /// Last error.
78    pub last_error: Option<String>,
79    /// Creation timestamp.
80    pub created_at: i64,
81    /// Last-update timestamp.
82    pub updated_at: i64,
83}
84
85/// Inserts a new entry into `pending_memories` with status `validated`.
86///
87/// Returns the generated `pending_id`.
88#[allow(clippy::too_many_arguments)]
89pub fn insert_validated(
90    conn: &Connection,
91    name: &str,
92    namespace: &str,
93    memory_type: &str,
94    description: Option<&str>,
95    body: &[u8],
96    body_hash: &str,
97    entities_json: Option<&str>,
98    relationships_json: Option<&str>,
99) -> Result<i64, AppError> {
100    conn.execute(
101        "INSERT INTO pending_memories
102            (name, namespace, memory_type, description, body, body_hash,
103             entities_json, relationships_json, status, attempt_count)
104         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, 'validated', 0)",
105        params![
106            name,
107            namespace,
108            memory_type,
109            description,
110            body,
111            body_hash,
112            entities_json,
113            relationships_json,
114        ],
115    )?;
116    Ok(conn.last_insert_rowid())
117}
118
119/// Updates the status to `embedding_in_progress` and increments `attempt_count`.
120pub fn update_to_embedding_in_progress(conn: &Connection, pending_id: i64) -> Result<(), AppError> {
121    conn.execute(
122        "UPDATE pending_memories
123         SET status = 'embedding_in_progress',
124             attempt_count = attempt_count + 1,
125             updated_at = unixepoch()
126         WHERE pending_id = ?1",
127        params![pending_id],
128    )?;
129    Ok(())
130}
131
132/// Updates the status to `embedding_done` and stores the embedding BLOB.
133pub fn update_to_embedding_done(
134    conn: &Connection,
135    pending_id: i64,
136    embedding: &[u8],
137    dim: i32,
138) -> Result<(), AppError> {
139    conn.execute(
140        "UPDATE pending_memories
141         SET status = 'embedding_done',
142             embedding = ?1,
143             embedding_dim = ?2,
144             updated_at = unixepoch()
145         WHERE pending_id = ?3",
146        params![embedding, dim, pending_id],
147    )?;
148    Ok(())
149}
150
151/// Marks as `committed` (called after a successful Stage C).
152pub fn mark_committed(conn: &Connection, pending_id: i64) -> Result<(), AppError> {
153    conn.execute(
154        "UPDATE pending_memories
155         SET status = 'committed',
156             updated_at = unixepoch()
157         WHERE pending_id = ?1",
158        params![pending_id],
159    )?;
160    Ok(())
161}
162
163/// Marks as `failed` with an error message.
164pub fn mark_failed(conn: &Connection, pending_id: i64, error: &str) -> Result<(), AppError> {
165    conn.execute(
166        "UPDATE pending_memories
167         SET status = 'failed',
168             last_error = ?1,
169             updated_at = unixepoch()
170         WHERE pending_id = ?2",
171        params![error, pending_id],
172    )?;
173    Ok(())
174}
175
176/// Lists entries by status, ordered by `updated_at` ascending.
177pub fn list_by_status(
178    conn: &Connection,
179    status: PendingStatus,
180    limit: usize,
181) -> Result<Vec<PendingMemory>, AppError> {
182    let mut stmt = conn.prepare(
183        "SELECT pending_id, name, namespace, memory_type, description, body,
184                body_hash, entities_json, relationships_json, status,
185                embedding, embedding_dim, attempt_count, last_error,
186                created_at, updated_at
187         FROM pending_memories
188         WHERE status = ?1
189         ORDER BY updated_at ASC
190         LIMIT ?2",
191    )?;
192    let rows = stmt.query_map(params![status.as_str(), limit as i64], |row| {
193        Ok(PendingMemory {
194            pending_id: row.get(0)?,
195            name: row.get(1)?,
196            namespace: row.get(2)?,
197            memory_type: row.get(3)?,
198            description: row.get(4)?,
199            body: row.get(5)?,
200            body_hash: row.get(6)?,
201            entities_json: row.get(7)?,
202            relationships_json: row.get(8)?,
203            status: parse_status(&row.get::<_, String>(9)?).map_err(|e| -> rusqlite::Error {
204                rusqlite::Error::FromSqlConversionFailure(
205                    9,
206                    rusqlite::types::Type::Text,
207                    Box::new(std::io::Error::other(e.to_string())),
208                )
209            })?,
210            embedding: row.get(10)?,
211            embedding_dim: row.get(11)?,
212            attempt_count: row.get(12)?,
213            last_error: row.get(13)?,
214            created_at: row.get(14)?,
215            updated_at: row.get(15)?,
216        })
217    })?;
218    let mut pending = Vec::new();
219    for row in rows {
220        pending.push(row?);
221    }
222    Ok(pending)
223}
224
225/// Looks up by `pending_id`.
226pub fn find_by_id(conn: &Connection, pending_id: i64) -> Result<Option<PendingMemory>, AppError> {
227    let mut stmt = conn.prepare(
228        "SELECT pending_id, name, namespace, memory_type, description, body,
229                body_hash, entities_json, relationships_json, status,
230                embedding, embedding_dim, attempt_count, last_error,
231                created_at, updated_at
232         FROM pending_memories
233         WHERE pending_id = ?1",
234    )?;
235    let mut rows = stmt.query(params![pending_id])?;
236    if let Some(row) = rows.next()? {
237        Ok(Some(PendingMemory {
238            pending_id: row.get(0)?,
239            name: row.get(1)?,
240            namespace: row.get(2)?,
241            memory_type: row.get(3)?,
242            description: row.get(4)?,
243            body: row.get(5)?,
244            body_hash: row.get(6)?,
245            entities_json: row.get(7)?,
246            relationships_json: row.get(8)?,
247            status: parse_status(row.get::<_, String>(9)?.as_str())?,
248            embedding: row.get(10)?,
249            embedding_dim: row.get(11)?,
250            attempt_count: row.get(12)?,
251            last_error: row.get(13)?,
252            created_at: row.get(14)?,
253            updated_at: row.get(15)?,
254        }))
255    } else {
256        Ok(None)
257    }
258}
259
260/// Removes `embedding_in_progress` entries older than `older_than_secs`.
261/// Returns the number of entries removed.
262pub fn cleanup_older_than(conn: &Connection, older_than_secs: i64) -> Result<usize, AppError> {
263    let cutoff = chrono::Utc::now().timestamp() - older_than_secs;
264    let count = conn.execute(
265        "DELETE FROM pending_memories
266         WHERE status IN ('embedding_in_progress', 'validated', 'failed')
267           AND updated_at < ?1",
268        params![cutoff],
269    )?;
270    Ok(count)
271}
272
273fn parse_status(s: &str) -> Result<PendingStatus, AppError> {
274    match s {
275        "validated" => Ok(PendingStatus::Validated),
276        "embedding_in_progress" => Ok(PendingStatus::EmbeddingInProgress),
277        "embedding_done" => Ok(PendingStatus::EmbeddingDone),
278        "committed" => Ok(PendingStatus::Committed),
279        "abandoned" => Ok(PendingStatus::Abandoned),
280        "failed" => Ok(PendingStatus::Failed),
281        other => Err(AppError::Validation(crate::i18n::validation::unknown_pending_memories_status(other))),
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use rusqlite::Connection;
289
290    fn fresh_db() -> Connection {
291        let mut conn = Connection::open_in_memory().expect("in-memory db");
292        conn.execute_batch("PRAGMA foreign_keys = ON;")
293            .expect("pragma");
294        crate::migrations::runner()
295            .run(&mut conn)
296            .expect("migrations apply");
297        conn
298    }
299
300    #[test]
301    fn insert_validated_returns_pending_id() {
302        let conn = fresh_db();
303        let id = insert_validated(
304            &conn,
305            "test-pending",
306            "global",
307            "note",
308            Some("desc"),
309            b"body bytes",
310            "blake3-hash-here",
311            None,
312            None,
313        )
314        .expect("insert");
315        assert!(id > 0);
316    }
317
318    #[test]
319    fn status_transition_validated_to_committed() {
320        let conn = fresh_db();
321        let id =
322            insert_validated(&conn, "x", "global", "note", None, b"b", "h", None, None).unwrap();
323        update_to_embedding_in_progress(&conn, id).unwrap();
324        let p = find_by_id(&conn, id).unwrap().unwrap();
325        assert_eq!(p.status, PendingStatus::EmbeddingInProgress);
326        assert_eq!(p.attempt_count, 1);
327
328        // Embedding BLOB is little-endian &[u8] — use raw bytes for the test
329        let fake_emb: Vec<u8> = vec![0u8; 64 * 4]; // 64 * 4 bytes
330        update_to_embedding_done(&conn, id, &fake_emb, 64).unwrap();
331        let p = find_by_id(&conn, id).unwrap().unwrap();
332        assert_eq!(p.status, PendingStatus::EmbeddingDone);
333        assert_eq!(p.embedding_dim, Some(64));
334
335        mark_committed(&conn, id).unwrap();
336        let p = find_by_id(&conn, id).unwrap().unwrap();
337        assert_eq!(p.status, PendingStatus::Committed);
338    }
339
340    #[test]
341    fn list_by_status_filters_correctly() {
342        let conn = fresh_db();
343        let id1 =
344            insert_validated(&conn, "a", "global", "note", None, b"b", "h", None, None).unwrap();
345        let _id2 =
346            insert_validated(&conn, "b", "global", "note", None, b"b", "h", None, None).unwrap();
347        mark_committed(&conn, id1).unwrap();
348        let validated = list_by_status(&conn, PendingStatus::Validated, 10).unwrap();
349        assert_eq!(validated.len(), 1);
350        assert_eq!(validated[0].name, "b");
351    }
352
353    #[test]
354    fn cleanup_older_than_removes_stale() {
355        let conn = fresh_db();
356        let _id = insert_validated(
357            &conn, "stale", "global", "note", None, b"b", "h", None, None,
358        )
359        .unwrap();
360        // Cleanup com cutoff no futuro = remove tudo
361        let removed = cleanup_older_than(&conn, -3600).unwrap();
362        assert_eq!(removed, 1);
363    }
364
365    #[test]
366    fn mark_failed_records_error() {
367        let conn = fresh_db();
368        let id =
369            insert_validated(&conn, "f", "global", "note", None, b"b", "h", None, None).unwrap();
370        mark_failed(&conn, id, "codex exited with OOM").unwrap();
371        let p = find_by_id(&conn, id).unwrap().unwrap();
372        assert_eq!(p.status, PendingStatus::Failed);
373        assert_eq!(p.last_error.as_deref(), Some("codex exited with OOM"));
374    }
375}
sqlite_graphrag/storage/pending_memories.rs

sqlite_graphrag/storage/
pending_memories.rs
