sqlite_graphrag/storage/
urls.rs

1//! Persistence for URLs extracted from memory bodies.
2//!
3//! Manages the `memory_urls` table: insert, deduplicate, and query URLs
4//! linked to a specific memory record.
5
6use crate::errors::AppError;
7use rusqlite::Connection;
8
9/// URL extracted from a memory body.
10pub struct MemoryUrl {
11    /// URL value.
12    pub url: String,
13    /// Pagination offset.
14    pub offset: Option<i64>,
15}
16
17/// Inserts a URL into the `memory_urls` table. Silently ignores duplicates.
18pub fn insert_url(conn: &Connection, memory_id: i64, entry: &MemoryUrl) -> Result<(), AppError> {
19    conn.execute(
20        "INSERT OR IGNORE INTO memory_urls (memory_id, url, url_offset) VALUES (?1, ?2, ?3)",
21        rusqlite::params![memory_id, entry.url, entry.offset],
22    )?;
23    Ok(())
24}
25
26/// Inserts multiple URLs for a memory. Returns the count inserted (duplicates ignored).
27/// Individual errors are logged as warn and not propagated — non-critical path.
28pub fn insert_urls(conn: &Connection, memory_id: i64, urls: &[MemoryUrl]) -> usize {
29    let mut inserted = 0usize;
30    for entry in urls {
31        match insert_url(conn, memory_id, entry) {
32            Ok(()) => {
33                let changed = conn.changes();
34                if changed > 0 {
35                    inserted += 1;
36                }
37            }
38            Err(e) => {
39                tracing::warn!(target: "storage", url = %entry.url, error = %e, "url persistence failed");
40            }
41        }
42    }
43    inserted
44}
45
46/// Lists all URLs associated with a memory.
47pub fn list_by_memory(conn: &Connection, memory_id: i64) -> Result<Vec<MemoryUrl>, AppError> {
48    let mut stmt = conn.prepare_cached(
49        "SELECT url, url_offset FROM memory_urls WHERE memory_id = ?1 ORDER BY id",
50    )?;
51    let rows = stmt.query_map(rusqlite::params![memory_id], |row| {
52        Ok(MemoryUrl {
53            url: row.get(0)?,
54            offset: row.get(1)?,
55        })
56    })?;
57    let mut result = Vec::with_capacity(8);
58    for row in rows {
59        result.push(row?);
60    }
61    Ok(result)
62}
63
64/// Removes all URLs for a memory.
65pub fn delete_by_memory(conn: &Connection, memory_id: i64) -> Result<(), AppError> {
66    conn.execute(
67        "DELETE FROM memory_urls WHERE memory_id = ?1",
68        rusqlite::params![memory_id],
69    )?;
70    Ok(())
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use rusqlite::Connection;
77    use tempfile::TempDir;
78
79    type TestResult = Result<(), Box<dyn std::error::Error>>;
80
81    fn setup_db() -> Result<(TempDir, Connection), Box<dyn std::error::Error>> {
82        crate::storage::connection::register_vec_extension();
83        let tmp = TempDir::new()?;
84        let db_path = tmp.path().join("test.db");
85        let mut conn = Connection::open(&db_path)?;
86        crate::migrations::runner().run(&mut conn)?;
87        Ok((tmp, conn))
88    }
89
90    fn insert_test_memory(conn: &Connection) -> Result<i64, Box<dyn std::error::Error>> {
91        conn.execute(
92            "INSERT INTO memories (name, type, description, body, body_hash) VALUES ('mem', 'user', 'desc', 'body', 'hash')",
93            [],
94        )?;
95        Ok(conn.last_insert_rowid())
96    }
97
98    #[test]
99    fn insert_url_persists_and_list_returns() -> TestResult {
100        let (_tmp, conn) = setup_db()?;
101        let mem_id = insert_test_memory(&conn)?;
102
103        insert_url(
104            &conn,
105            mem_id,
106            &MemoryUrl {
107                url: "https://example.com/page".to_string(),
108                offset: Some(5),
109            },
110        )?;
111
112        let urls = list_by_memory(&conn, mem_id)?;
113        assert_eq!(urls.len(), 1);
114        assert_eq!(urls[0].url, "https://example.com/page");
115        assert_eq!(urls[0].offset, Some(5));
116        Ok(())
117    }
118
119    #[test]
120    fn insert_url_duplicate_ignored() -> TestResult {
121        let (_tmp, conn) = setup_db()?;
122        let mem_id = insert_test_memory(&conn)?;
123
124        let entry = MemoryUrl {
125            url: "https://example.com/dup".to_string(),
126            offset: None,
127        };
128        insert_url(&conn, mem_id, &entry)?;
129        insert_url(&conn, mem_id, &entry)?;
130
131        let urls = list_by_memory(&conn, mem_id)?;
132        assert_eq!(urls.len(), 1, "duplicata deve ser ignorada");
133        Ok(())
134    }
135
136    #[test]
137    fn insert_urls_returns_inserted_count() -> TestResult {
138        let (_tmp, conn) = setup_db()?;
139        let mem_id = insert_test_memory(&conn)?;
140
141        let batch = vec![
142            MemoryUrl {
143                url: "https://alpha.example.com".to_string(),
144                offset: Some(0),
145            },
146            MemoryUrl {
147                url: "https://beta.example.com".to_string(),
148                offset: Some(10),
149            },
150            MemoryUrl {
151                url: "https://alpha.example.com".to_string(),
152                offset: Some(0),
153            },
154        ];
155        let count = insert_urls(&conn, mem_id, &batch);
156        assert_eq!(count, 2, "only 2 unique entries must be inserted");
157        Ok(())
158    }
159
160    #[test]
161    fn delete_by_memory_removes_all_urls() -> TestResult {
162        let (_tmp, conn) = setup_db()?;
163        let mem_id = insert_test_memory(&conn)?;
164
165        insert_url(
166            &conn,
167            mem_id,
168            &MemoryUrl {
169                url: "https://to-delete.example.com".to_string(),
170                offset: None,
171            },
172        )?;
173        assert_eq!(list_by_memory(&conn, mem_id)?.len(), 1);
174
175        delete_by_memory(&conn, mem_id)?;
176        assert_eq!(list_by_memory(&conn, mem_id)?.len(), 0);
177        Ok(())
178    }
179}
sqlite_graphrag/storage/urls.rs

sqlite_graphrag/storage/
urls.rs
