Skip to main content

sqlite_graphrag/storage/
entities.rs

1//! Persistence layer for entities, relationships and their junction tables.
2//!
3//! The entity graph mirrors the conceptual content of memories: `entities`
4//! holds nodes, `relationships` holds typed edges and `memory_entities` and
5//! `memory_relationships` connect each memory to the graph slice it emitted.
6
7use crate::embedder::f32_to_bytes;
8use crate::errors::AppError;
9use rusqlite::{params, Connection};
10use serde::{Deserialize, Serialize};
11
12/// Input payload used to upsert a single entity.
13///
14/// `name` is normalized to kebab-case by the caller. `description` is
15/// optional and preserved across upserts when the new value is `None`.
16#[derive(Debug, Serialize, Deserialize, Clone)]
17#[serde(deny_unknown_fields)]
18pub struct NewEntity {
19    pub name: String,
20    #[serde(alias = "type")]
21    pub entity_type: String,
22    pub description: Option<String>,
23}
24
25/// Input payload used to upsert a typed relationship between entities.
26///
27/// `strength` must lie within `[0.0, 1.0]` and is mapped to the `weight`
28/// column of the `relationships` table.
29#[derive(Debug, Serialize, Deserialize, Clone)]
30#[serde(deny_unknown_fields)]
31pub struct NewRelationship {
32    pub source: String,
33    pub target: String,
34    pub relation: String,
35    pub strength: f64,
36    pub description: Option<String>,
37}
38
39/// Upserts an entity and returns its primary key.
40///
41/// Uses `ON CONFLICT(namespace, name)` to keep one row per entity within a
42/// namespace, refreshing `type` and `description` opportunistically.
43///
44/// # Errors
45///
46/// Returns `Err(AppError::Database)` on any `rusqlite` failure.
47pub fn upsert_entity(conn: &Connection, namespace: &str, e: &NewEntity) -> Result<i64, AppError> {
48    conn.execute(
49        "INSERT INTO entities (namespace, name, type, description)
50         VALUES (?1, ?2, ?3, ?4)
51         ON CONFLICT(namespace, name) DO UPDATE SET
52           type        = excluded.type,
53           description = COALESCE(excluded.description, entities.description),
54           updated_at  = unixepoch()",
55        params![namespace, e.name, e.entity_type, e.description],
56    )?;
57    let id: i64 = conn.query_row(
58        "SELECT id FROM entities WHERE namespace = ?1 AND name = ?2",
59        params![namespace, e.name],
60        |r| r.get(0),
61    )?;
62    Ok(id)
63}
64
65/// Replaces the vector row for an entity in `vec_entities`.
66///
67/// vec0 virtual tables do not honour `INSERT OR REPLACE` when the primary key
68/// already exists — they raise a UNIQUE constraint error instead of silently
69/// replacing the row. The workaround is an explicit DELETE before INSERT so
70/// that the insert never conflicts. `embedding` must have length
71/// [`crate::constants::EMBEDDING_DIM`].
72///
73/// # Errors
74///
75/// Returns `Err(AppError::Database)` on any `rusqlite` failure.
76pub fn upsert_entity_vec(
77    conn: &Connection,
78    entity_id: i64,
79    namespace: &str,
80    entity_type: &str,
81    embedding: &[f32],
82    name: &str,
83) -> Result<(), AppError> {
84    conn.execute(
85        "DELETE FROM vec_entities WHERE entity_id = ?1",
86        params![entity_id],
87    )?;
88    conn.execute(
89        "INSERT INTO vec_entities(entity_id, namespace, type, embedding, name)
90         VALUES (?1, ?2, ?3, ?4, ?5)",
91        params![
92            entity_id,
93            namespace,
94            entity_type,
95            f32_to_bytes(embedding),
96            name
97        ],
98    )?;
99    Ok(())
100}
101
102/// Upserts a typed relationship between two entity ids.
103///
104/// Conflicts on `(source_id, target_id, relation)` refresh `weight` and
105/// preserve a non-null `description`. Returns the `rowid` of the stored row.
106///
107/// # Errors
108///
109/// Returns `Err(AppError::Database)` on any `rusqlite` failure.
110pub fn upsert_relationship(
111    conn: &Connection,
112    namespace: &str,
113    source_id: i64,
114    target_id: i64,
115    rel: &NewRelationship,
116) -> Result<i64, AppError> {
117    conn.execute(
118        "INSERT INTO relationships (namespace, source_id, target_id, relation, weight, description)
119         VALUES (?1, ?2, ?3, ?4, ?5, ?6)
120         ON CONFLICT(source_id, target_id, relation) DO UPDATE SET
121           weight = excluded.weight,
122           description = COALESCE(excluded.description, relationships.description)",
123        params![
124            namespace,
125            source_id,
126            target_id,
127            rel.relation,
128            rel.strength,
129            rel.description
130        ],
131    )?;
132    let id: i64 = conn.query_row(
133        "SELECT id FROM relationships WHERE source_id=?1 AND target_id=?2 AND relation=?3",
134        params![source_id, target_id, rel.relation],
135        |r| r.get(0),
136    )?;
137    Ok(id)
138}
139
140pub fn link_memory_entity(
141    conn: &Connection,
142    memory_id: i64,
143    entity_id: i64,
144) -> Result<(), AppError> {
145    conn.execute(
146        "INSERT OR IGNORE INTO memory_entities (memory_id, entity_id) VALUES (?1, ?2)",
147        params![memory_id, entity_id],
148    )?;
149    Ok(())
150}
151
152pub fn link_memory_relationship(
153    conn: &Connection,
154    memory_id: i64,
155    rel_id: i64,
156) -> Result<(), AppError> {
157    conn.execute(
158        "INSERT OR IGNORE INTO memory_relationships (memory_id, relationship_id) VALUES (?1, ?2)",
159        params![memory_id, rel_id],
160    )?;
161    Ok(())
162}
163
164pub fn increment_degree(conn: &Connection, entity_id: i64) -> Result<(), AppError> {
165    conn.execute(
166        "UPDATE entities SET degree = degree + 1 WHERE id = ?1",
167        params![entity_id],
168    )?;
169    Ok(())
170}
171
172/// Busca a entidade por nome e namespace. Retorna o id se existir.
173pub fn find_entity_id(
174    conn: &Connection,
175    namespace: &str,
176    name: &str,
177) -> Result<Option<i64>, AppError> {
178    let mut stmt =
179        conn.prepare_cached("SELECT id FROM entities WHERE namespace = ?1 AND name = ?2")?;
180    match stmt.query_row(params![namespace, name], |r| r.get::<_, i64>(0)) {
181        Ok(id) => Ok(Some(id)),
182        Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
183        Err(e) => Err(AppError::Database(e)),
184    }
185}
186
187/// Estrutura representando uma relação existente.
188#[derive(Debug, Serialize)]
189pub struct RelationshipRow {
190    pub id: i64,
191    pub namespace: String,
192    pub source_id: i64,
193    pub target_id: i64,
194    pub relation: String,
195    pub weight: f64,
196    pub description: Option<String>,
197}
198
199/// Busca uma relação específica por (source_id, target_id, relation).
200pub fn find_relationship(
201    conn: &Connection,
202    source_id: i64,
203    target_id: i64,
204    relation: &str,
205) -> Result<Option<RelationshipRow>, AppError> {
206    let mut stmt = conn.prepare_cached(
207        "SELECT id, namespace, source_id, target_id, relation, weight, description
208         FROM relationships
209         WHERE source_id = ?1 AND target_id = ?2 AND relation = ?3",
210    )?;
211    match stmt.query_row(params![source_id, target_id, relation], |r| {
212        Ok(RelationshipRow {
213            id: r.get(0)?,
214            namespace: r.get(1)?,
215            source_id: r.get(2)?,
216            target_id: r.get(3)?,
217            relation: r.get(4)?,
218            weight: r.get(5)?,
219            description: r.get(6)?,
220        })
221    }) {
222        Ok(row) => Ok(Some(row)),
223        Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
224        Err(e) => Err(AppError::Database(e)),
225    }
226}
227
228/// Cria uma relação se não existir (retorna action="created")
229/// ou retorna a relação existente (action="already_exists") com peso atualizado.
230pub fn create_or_fetch_relationship(
231    conn: &Connection,
232    namespace: &str,
233    source_id: i64,
234    target_id: i64,
235    relation: &str,
236    weight: f64,
237    description: Option<&str>,
238) -> Result<(i64, bool), AppError> {
239    // Check if it exists first.
240    let existing = find_relationship(conn, source_id, target_id, relation)?;
241    if let Some(row) = existing {
242        return Ok((row.id, false));
243    }
244    conn.execute(
245        "INSERT INTO relationships (namespace, source_id, target_id, relation, weight, description)
246         VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
247        params![
248            namespace,
249            source_id,
250            target_id,
251            relation,
252            weight,
253            description
254        ],
255    )?;
256    let id: i64 = conn.query_row(
257        "SELECT id FROM relationships WHERE source_id = ?1 AND target_id = ?2 AND relation = ?3",
258        params![source_id, target_id, relation],
259        |r| r.get(0),
260    )?;
261    Ok((id, true))
262}
263
264/// Remove uma relação pelo id e limpa memory_relationships.
265pub fn delete_relationship_by_id(conn: &Connection, relationship_id: i64) -> Result<(), AppError> {
266    conn.execute(
267        "DELETE FROM memory_relationships WHERE relationship_id = ?1",
268        params![relationship_id],
269    )?;
270    conn.execute(
271        "DELETE FROM relationships WHERE id = ?1",
272        params![relationship_id],
273    )?;
274    Ok(())
275}
276
277/// Recalcula o campo `degree` de uma entidade.
278pub fn recalculate_degree(conn: &Connection, entity_id: i64) -> Result<(), AppError> {
279    conn.execute(
280        "UPDATE entities
281         SET degree = (SELECT COUNT(*) FROM relationships
282                       WHERE source_id = entities.id OR target_id = entities.id)
283         WHERE id = ?1",
284        params![entity_id],
285    )?;
286    Ok(())
287}
288
289/// Linha de entidade com dados suficientes para exportação/consulta de grafo.
290#[derive(Debug, Serialize, Clone)]
291pub struct EntityNode {
292    pub id: i64,
293    pub name: String,
294    pub namespace: String,
295    pub kind: String,
296}
297
298/// Lista entidades, filtrando por namespace se fornecido.
299pub fn list_entities(
300    conn: &Connection,
301    namespace: Option<&str>,
302) -> Result<Vec<EntityNode>, AppError> {
303    if let Some(ns) = namespace {
304        let mut stmt = conn.prepare(
305            "SELECT id, name, namespace, type FROM entities WHERE namespace = ?1 ORDER BY id",
306        )?;
307        let rows = stmt
308            .query_map(params![ns], |r| {
309                Ok(EntityNode {
310                    id: r.get(0)?,
311                    name: r.get(1)?,
312                    namespace: r.get(2)?,
313                    kind: r.get(3)?,
314                })
315            })?
316            .collect::<Result<Vec<_>, _>>()?;
317        Ok(rows)
318    } else {
319        let mut stmt =
320            conn.prepare("SELECT id, name, namespace, type FROM entities ORDER BY namespace, id")?;
321        let rows = stmt
322            .query_map([], |r| {
323                Ok(EntityNode {
324                    id: r.get(0)?,
325                    name: r.get(1)?,
326                    namespace: r.get(2)?,
327                    kind: r.get(3)?,
328                })
329            })?
330            .collect::<Result<Vec<_>, _>>()?;
331        Ok(rows)
332    }
333}
334
335/// Lista relações filtradas por namespace (das entidades de origem/destino).
336pub fn list_relationships_by_namespace(
337    conn: &Connection,
338    namespace: Option<&str>,
339) -> Result<Vec<RelationshipRow>, AppError> {
340    if let Some(ns) = namespace {
341        let mut stmt = conn.prepare(
342            "SELECT r.id, r.namespace, r.source_id, r.target_id, r.relation, r.weight, r.description
343             FROM relationships r
344             JOIN entities se ON se.id = r.source_id AND se.namespace = ?1
345             JOIN entities te ON te.id = r.target_id AND te.namespace = ?1
346             ORDER BY r.id",
347        )?;
348        let rows = stmt
349            .query_map(params![ns], |r| {
350                Ok(RelationshipRow {
351                    id: r.get(0)?,
352                    namespace: r.get(1)?,
353                    source_id: r.get(2)?,
354                    target_id: r.get(3)?,
355                    relation: r.get(4)?,
356                    weight: r.get(5)?,
357                    description: r.get(6)?,
358                })
359            })?
360            .collect::<Result<Vec<_>, _>>()?;
361        Ok(rows)
362    } else {
363        let mut stmt = conn.prepare(
364            "SELECT id, namespace, source_id, target_id, relation, weight, description
365             FROM relationships ORDER BY id",
366        )?;
367        let rows = stmt
368            .query_map([], |r| {
369                Ok(RelationshipRow {
370                    id: r.get(0)?,
371                    namespace: r.get(1)?,
372                    source_id: r.get(2)?,
373                    target_id: r.get(3)?,
374                    relation: r.get(4)?,
375                    weight: r.get(5)?,
376                    description: r.get(6)?,
377                })
378            })?
379            .collect::<Result<Vec<_>, _>>()?;
380        Ok(rows)
381    }
382}
383
384/// Localiza entidades órfãs: sem vínculo em memory_entities e sem relações.
385pub fn find_orphan_entity_ids(
386    conn: &Connection,
387    namespace: Option<&str>,
388) -> Result<Vec<i64>, AppError> {
389    if let Some(ns) = namespace {
390        let mut stmt = conn.prepare(
391            "SELECT e.id FROM entities e
392             WHERE e.namespace = ?1
393               AND NOT EXISTS (SELECT 1 FROM memory_entities me WHERE me.entity_id = e.id)
394               AND NOT EXISTS (
395                   SELECT 1 FROM relationships r
396                   WHERE r.source_id = e.id OR r.target_id = e.id
397               )",
398        )?;
399        let ids = stmt
400            .query_map(params![ns], |r| r.get::<_, i64>(0))?
401            .collect::<Result<Vec<_>, _>>()?;
402        Ok(ids)
403    } else {
404        let mut stmt = conn.prepare(
405            "SELECT e.id FROM entities e
406             WHERE NOT EXISTS (SELECT 1 FROM memory_entities me WHERE me.entity_id = e.id)
407               AND NOT EXISTS (
408                   SELECT 1 FROM relationships r
409                   WHERE r.source_id = e.id OR r.target_id = e.id
410               )",
411        )?;
412        let ids = stmt
413            .query_map([], |r| r.get::<_, i64>(0))?
414            .collect::<Result<Vec<_>, _>>()?;
415        Ok(ids)
416    }
417}
418
419/// Deleta entidades e seus vetores associados. Retorna o número de entidades removidas.
420pub fn delete_entities_by_ids(conn: &Connection, entity_ids: &[i64]) -> Result<usize, AppError> {
421    if entity_ids.is_empty() {
422        return Ok(0);
423    }
424    let mut removed = 0usize;
425    for id in entity_ids {
426        // vec0 lacks FK CASCADE — clean vec_entities explicitly.
427        let _ = conn.execute("DELETE FROM vec_entities WHERE entity_id = ?1", params![id]);
428        let affected = conn.execute("DELETE FROM entities WHERE id = ?1", params![id])?;
429        removed += affected;
430    }
431    Ok(removed)
432}
433
434pub fn knn_search(
435    conn: &Connection,
436    embedding: &[f32],
437    namespace: &str,
438    k: usize,
439) -> Result<Vec<(i64, f32)>, AppError> {
440    let bytes = f32_to_bytes(embedding);
441    let mut stmt = conn.prepare(
442        "SELECT entity_id, distance FROM vec_entities
443         WHERE embedding MATCH ?1 AND namespace = ?2
444         ORDER BY distance LIMIT ?3",
445    )?;
446    let rows = stmt
447        .query_map(params![bytes, namespace, k as i64], |r| {
448            Ok((r.get::<_, i64>(0)?, r.get::<_, f32>(1)?))
449        })?
450        .collect::<Result<Vec<_>, _>>()?;
451    Ok(rows)
452}
453
454#[cfg(test)]
455mod tests {
456    use super::*;
457    use crate::constants::EMBEDDING_DIM;
458    use crate::storage::connection::register_vec_extension;
459    use rusqlite::Connection;
460    use tempfile::TempDir;
461
462    type Resultado = Result<(), Box<dyn std::error::Error>>;
463
464    fn setup_db() -> Result<(TempDir, Connection), Box<dyn std::error::Error>> {
465        register_vec_extension();
466        let tmp = TempDir::new()?;
467        let db_path = tmp.path().join("test.db");
468        let mut conn = Connection::open(&db_path)?;
469        crate::migrations::runner().run(&mut conn)?;
470        Ok((tmp, conn))
471    }
472
473    fn insert_memory(conn: &Connection) -> Result<i64, Box<dyn std::error::Error>> {
474        conn.execute(
475            "INSERT INTO memories (namespace, name, type, description, body, body_hash)
476             VALUES ('global', 'test-mem', 'user', 'desc', 'body', 'hash1')",
477            [],
478        )?;
479        Ok(conn.last_insert_rowid())
480    }
481
482    fn nova_entidade(name: &str) -> NewEntity {
483        NewEntity {
484            name: name.to_string(),
485            entity_type: "project".to_string(),
486            description: None,
487        }
488    }
489
490    fn embedding_zero() -> Vec<f32> {
491        vec![0.0f32; EMBEDDING_DIM]
492    }
493
494    // ------------------------------------------------------------------ //
495    // upsert_entity
496    // ------------------------------------------------------------------ //
497
498    #[test]
499    fn test_upsert_entity_cria_nova() -> Resultado {
500        let (_tmp, conn) = setup_db()?;
501        let e = nova_entidade("projeto-alpha");
502        let id = upsert_entity(&conn, "global", &e)?;
503        assert!(id > 0);
504        Ok(())
505    }
506
507    #[test]
508    fn test_upsert_entity_idempotente_retorna_mesmo_id() -> Resultado {
509        let (_tmp, conn) = setup_db()?;
510        let e = nova_entidade("projeto-beta");
511        let id1 = upsert_entity(&conn, "global", &e)?;
512        let id2 = upsert_entity(&conn, "global", &e)?;
513        assert_eq!(id1, id2);
514        Ok(())
515    }
516
517    #[test]
518    fn test_upsert_entity_atualiza_descricao() -> Resultado {
519        let (_tmp, conn) = setup_db()?;
520        let e1 = nova_entidade("projeto-gamma");
521        let id1 = upsert_entity(&conn, "global", &e1)?;
522
523        let e2 = NewEntity {
524            name: "projeto-gamma".to_string(),
525            entity_type: "tool".to_string(),
526            description: Some("nova desc".to_string()),
527        };
528        let id2 = upsert_entity(&conn, "global", &e2)?;
529        assert_eq!(id1, id2);
530
531        let desc: Option<String> = conn.query_row(
532            "SELECT description FROM entities WHERE id = ?1",
533            params![id1],
534            |r| r.get(0),
535        )?;
536        assert_eq!(desc.as_deref(), Some("nova desc"));
537        Ok(())
538    }
539
540    #[test]
541    fn test_upsert_entity_namespaces_diferentes_criam_registros_distintos() -> Resultado {
542        let (_tmp, conn) = setup_db()?;
543        let e = nova_entidade("compartilhada");
544        let id1 = upsert_entity(&conn, "ns1", &e)?;
545        let id2 = upsert_entity(&conn, "ns2", &e)?;
546        assert_ne!(id1, id2);
547        Ok(())
548    }
549
550    // ------------------------------------------------------------------ //
551    // upsert_entity_vec — cobre DELETE+INSERT (branch novo após fix OOM)
552    // ------------------------------------------------------------------ //
553
554    #[test]
555    fn test_upsert_entity_vec_primeira_vez_sem_conflito() -> Resultado {
556        let (_tmp, conn) = setup_db()?;
557        let e = nova_entidade("vec-nova");
558        let entity_id = upsert_entity(&conn, "global", &e)?;
559        let emb = embedding_zero();
560
561        let resultado = upsert_entity_vec(&conn, entity_id, "global", "project", &emb, "vec-nova");
562        assert!(resultado.is_ok(), "primeira inserção deve ter sucesso");
563
564        let count: i64 = conn.query_row(
565            "SELECT COUNT(*) FROM vec_entities WHERE entity_id = ?1",
566            params![entity_id],
567            |r| r.get(0),
568        )?;
569        assert_eq!(count, 1, "deve existir exatamente uma linha após inserção");
570        Ok(())
571    }
572
573    #[test]
574    fn test_upsert_entity_vec_segunda_vez_substitui_sem_erro() -> Resultado {
575        // Cobre o branch onde DELETE remove a linha existente antes do INSERT.
576        let (_tmp, conn) = setup_db()?;
577        let e = nova_entidade("vec-existente");
578        let entity_id = upsert_entity(&conn, "global", &e)?;
579        let emb = embedding_zero();
580
581        upsert_entity_vec(&conn, entity_id, "global", "project", &emb, "vec-existente")?;
582
583        // Segunda chamada: DELETE retorna 1 linha removida, INSERT deve ter sucesso.
584        let resultado =
585            upsert_entity_vec(&conn, entity_id, "global", "tool", &emb, "vec-existente");
586        assert!(
587            resultado.is_ok(),
588            "segunda inserção (replace) deve ter sucesso: {resultado:?}"
589        );
590
591        let count: i64 = conn.query_row(
592            "SELECT COUNT(*) FROM vec_entities WHERE entity_id = ?1",
593            params![entity_id],
594            |r| r.get(0),
595        )?;
596        assert_eq!(
597            count, 1,
598            "deve existir exatamente uma linha após substituição"
599        );
600        Ok(())
601    }
602
603    #[test]
604    fn test_upsert_entity_vec_multiplas_entidades_independentes() -> Resultado {
605        let (_tmp, conn) = setup_db()?;
606        let emb = embedding_zero();
607
608        for i in 0..3i64 {
609            let nome = format!("ent-{i}");
610            let e = nova_entidade(&nome);
611            let entity_id = upsert_entity(&conn, "global", &e)?;
612            upsert_entity_vec(&conn, entity_id, "global", "project", &emb, &nome)?;
613        }
614
615        let count: i64 = conn.query_row("SELECT COUNT(*) FROM vec_entities", [], |r| r.get(0))?;
616        assert_eq!(count, 3, "deve haver três linhas distintas em vec_entities");
617        Ok(())
618    }
619
620    // ------------------------------------------------------------------ //
621    // find_entity_id
622    // ------------------------------------------------------------------ //
623
624    #[test]
625    fn test_find_entity_id_existente_retorna_some() -> Resultado {
626        let (_tmp, conn) = setup_db()?;
627        let e = nova_entidade("entidade-busca");
628        let id_inserido = upsert_entity(&conn, "global", &e)?;
629        let id_encontrado = find_entity_id(&conn, "global", "entidade-busca")?;
630        assert_eq!(id_encontrado, Some(id_inserido));
631        Ok(())
632    }
633
634    #[test]
635    fn test_find_entity_id_inexistente_retorna_none() -> Resultado {
636        let (_tmp, conn) = setup_db()?;
637        let id = find_entity_id(&conn, "global", "nao-existe")?;
638        assert_eq!(id, None);
639        Ok(())
640    }
641
642    // ------------------------------------------------------------------ //
643    // delete_entities_by_ids
644    // ------------------------------------------------------------------ //
645
646    #[test]
647    fn test_delete_entities_by_ids_lista_vazia_retorna_zero() -> Resultado {
648        let (_tmp, conn) = setup_db()?;
649        let removidos = delete_entities_by_ids(&conn, &[])?;
650        assert_eq!(removidos, 0);
651        Ok(())
652    }
653
654    #[test]
655    fn test_delete_entities_by_ids_remove_entidade_valida() -> Resultado {
656        let (_tmp, conn) = setup_db()?;
657        let e = nova_entidade("para-deletar");
658        let entity_id = upsert_entity(&conn, "global", &e)?;
659
660        let removidos = delete_entities_by_ids(&conn, &[entity_id])?;
661        assert_eq!(removidos, 1);
662
663        let id = find_entity_id(&conn, "global", "para-deletar")?;
664        assert_eq!(id, None, "entidade deve ter sido removida");
665        Ok(())
666    }
667
668    #[test]
669    fn test_delete_entities_by_ids_id_inexistente_retorna_zero() -> Resultado {
670        let (_tmp, conn) = setup_db()?;
671        let removidos = delete_entities_by_ids(&conn, &[9999])?;
672        assert_eq!(removidos, 0);
673        Ok(())
674    }
675
676    #[test]
677    fn test_delete_entities_by_ids_remove_multiplas() -> Resultado {
678        let (_tmp, conn) = setup_db()?;
679        let id1 = upsert_entity(&conn, "global", &nova_entidade("del-a"))?;
680        let id2 = upsert_entity(&conn, "global", &nova_entidade("del-b"))?;
681        let id3 = upsert_entity(&conn, "global", &nova_entidade("del-c"))?;
682
683        let removidos = delete_entities_by_ids(&conn, &[id1, id2])?;
684        assert_eq!(removidos, 2);
685
686        assert!(find_entity_id(&conn, "global", "del-a")?.is_none());
687        assert!(find_entity_id(&conn, "global", "del-b")?.is_none());
688        assert!(find_entity_id(&conn, "global", "del-c")?.is_some());
689        let _ = id3;
690        Ok(())
691    }
692
693    #[test]
694    fn test_delete_entities_by_ids_tambem_remove_vec() -> Resultado {
695        let (_tmp, conn) = setup_db()?;
696        let e = nova_entidade("del-com-vec");
697        let entity_id = upsert_entity(&conn, "global", &e)?;
698        let emb = embedding_zero();
699        upsert_entity_vec(&conn, entity_id, "global", "project", &emb, "del-com-vec")?;
700
701        let count_antes: i64 = conn.query_row(
702            "SELECT COUNT(*) FROM vec_entities WHERE entity_id = ?1",
703            params![entity_id],
704            |r| r.get(0),
705        )?;
706        assert_eq!(count_antes, 1);
707
708        delete_entities_by_ids(&conn, &[entity_id])?;
709
710        let count_depois: i64 = conn.query_row(
711            "SELECT COUNT(*) FROM vec_entities WHERE entity_id = ?1",
712            params![entity_id],
713            |r| r.get(0),
714        )?;
715        assert_eq!(
716            count_depois, 0,
717            "vec_entities deve ser limpo junto com entities"
718        );
719        Ok(())
720    }
721
722    // ------------------------------------------------------------------ //
723    // upsert_relationship / find_relationship
724    // ------------------------------------------------------------------ //
725
726    #[test]
727    fn test_upsert_relationship_cria_nova() -> Resultado {
728        let (_tmp, conn) = setup_db()?;
729        let id_a = upsert_entity(&conn, "global", &nova_entidade("rel-a"))?;
730        let id_b = upsert_entity(&conn, "global", &nova_entidade("rel-b"))?;
731
732        let rel = NewRelationship {
733            source: "rel-a".to_string(),
734            target: "rel-b".to_string(),
735            relation: "uses".to_string(),
736            strength: 0.8,
737            description: None,
738        };
739        let rel_id = upsert_relationship(&conn, "global", id_a, id_b, &rel)?;
740        assert!(rel_id > 0);
741        Ok(())
742    }
743
744    #[test]
745    fn test_upsert_relationship_idempotente() -> Resultado {
746        let (_tmp, conn) = setup_db()?;
747        let id_a = upsert_entity(&conn, "global", &nova_entidade("idem-a"))?;
748        let id_b = upsert_entity(&conn, "global", &nova_entidade("idem-b"))?;
749
750        let rel = NewRelationship {
751            source: "idem-a".to_string(),
752            target: "idem-b".to_string(),
753            relation: "uses".to_string(),
754            strength: 0.5,
755            description: None,
756        };
757        let id1 = upsert_relationship(&conn, "global", id_a, id_b, &rel)?;
758        let id2 = upsert_relationship(&conn, "global", id_a, id_b, &rel)?;
759        assert_eq!(id1, id2);
760        Ok(())
761    }
762
763    #[test]
764    fn test_find_relationship_existente() -> Resultado {
765        let (_tmp, conn) = setup_db()?;
766        let id_a = upsert_entity(&conn, "global", &nova_entidade("fr-a"))?;
767        let id_b = upsert_entity(&conn, "global", &nova_entidade("fr-b"))?;
768
769        let rel = NewRelationship {
770            source: "fr-a".to_string(),
771            target: "fr-b".to_string(),
772            relation: "depends_on".to_string(),
773            strength: 0.7,
774            description: None,
775        };
776        upsert_relationship(&conn, "global", id_a, id_b, &rel)?;
777
778        let encontrada = find_relationship(&conn, id_a, id_b, "depends_on")?;
779        let row = encontrada.ok_or("relação deveria existir")?;
780        assert_eq!(row.source_id, id_a);
781        assert_eq!(row.target_id, id_b);
782        assert!((row.weight - 0.7).abs() < 1e-9);
783        Ok(())
784    }
785
786    #[test]
787    fn test_find_relationship_inexistente_retorna_none() -> Resultado {
788        let (_tmp, conn) = setup_db()?;
789        let resultado = find_relationship(&conn, 9999, 8888, "uses")?;
790        assert!(resultado.is_none());
791        Ok(())
792    }
793
794    // ------------------------------------------------------------------ //
795    // link_memory_entity / link_memory_relationship
796    // ------------------------------------------------------------------ //
797
798    #[test]
799    fn test_link_memory_entity_idempotente() -> Resultado {
800        let (_tmp, conn) = setup_db()?;
801        let memory_id = insert_memory(&conn)?;
802        let entity_id = upsert_entity(&conn, "global", &nova_entidade("me-ent"))?;
803
804        link_memory_entity(&conn, memory_id, entity_id)?;
805        let resultado = link_memory_entity(&conn, memory_id, entity_id);
806        assert!(
807            resultado.is_ok(),
808            "INSERT OR IGNORE não deve falhar em duplicata"
809        );
810        Ok(())
811    }
812
813    #[test]
814    fn test_link_memory_relationship_idempotente() -> Resultado {
815        let (_tmp, conn) = setup_db()?;
816        let memory_id = insert_memory(&conn)?;
817        let id_a = upsert_entity(&conn, "global", &nova_entidade("mr-a"))?;
818        let id_b = upsert_entity(&conn, "global", &nova_entidade("mr-b"))?;
819
820        let rel = NewRelationship {
821            source: "mr-a".to_string(),
822            target: "mr-b".to_string(),
823            relation: "uses".to_string(),
824            strength: 0.5,
825            description: None,
826        };
827        let rel_id = upsert_relationship(&conn, "global", id_a, id_b, &rel)?;
828
829        link_memory_relationship(&conn, memory_id, rel_id)?;
830        let resultado = link_memory_relationship(&conn, memory_id, rel_id);
831        assert!(
832            resultado.is_ok(),
833            "INSERT OR IGNORE não deve falhar em duplicata"
834        );
835        Ok(())
836    }
837
838    // ------------------------------------------------------------------ //
839    // increment_degree / recalculate_degree
840    // ------------------------------------------------------------------ //
841
842    #[test]
843    fn test_increment_degree_aumenta_contador() -> Resultado {
844        let (_tmp, conn) = setup_db()?;
845        let entity_id = upsert_entity(&conn, "global", &nova_entidade("grau-ent"))?;
846
847        increment_degree(&conn, entity_id)?;
848        increment_degree(&conn, entity_id)?;
849
850        let degree: i64 = conn.query_row(
851            "SELECT degree FROM entities WHERE id = ?1",
852            params![entity_id],
853            |r| r.get(0),
854        )?;
855        assert_eq!(degree, 2);
856        Ok(())
857    }
858
859    #[test]
860    fn test_recalculate_degree_reflete_relacoes_reais() -> Resultado {
861        let (_tmp, conn) = setup_db()?;
862        let id_a = upsert_entity(&conn, "global", &nova_entidade("rc-a"))?;
863        let id_b = upsert_entity(&conn, "global", &nova_entidade("rc-b"))?;
864        let id_c = upsert_entity(&conn, "global", &nova_entidade("rc-c"))?;
865
866        let rel1 = NewRelationship {
867            source: "rc-a".to_string(),
868            target: "rc-b".to_string(),
869            relation: "uses".to_string(),
870            strength: 0.5,
871            description: None,
872        };
873        let rel2 = NewRelationship {
874            source: "rc-c".to_string(),
875            target: "rc-a".to_string(),
876            relation: "depends_on".to_string(),
877            strength: 0.5,
878            description: None,
879        };
880        upsert_relationship(&conn, "global", id_a, id_b, &rel1)?;
881        upsert_relationship(&conn, "global", id_c, id_a, &rel2)?;
882
883        recalculate_degree(&conn, id_a)?;
884
885        let degree: i64 = conn.query_row(
886            "SELECT degree FROM entities WHERE id = ?1",
887            params![id_a],
888            |r| r.get(0),
889        )?;
890        assert_eq!(degree, 2, "rc-a aparece em duas relações (source+target)");
891        Ok(())
892    }
893
894    // ------------------------------------------------------------------ //
895    // find_orphan_entity_ids
896    // ------------------------------------------------------------------ //
897
898    #[test]
899    fn test_find_orphan_entity_ids_sem_orfaos() -> Resultado {
900        let (_tmp, conn) = setup_db()?;
901        let memory_id = insert_memory(&conn)?;
902        let entity_id = upsert_entity(&conn, "global", &nova_entidade("nao-orfa"))?;
903        link_memory_entity(&conn, memory_id, entity_id)?;
904
905        let orfas = find_orphan_entity_ids(&conn, Some("global"))?;
906        assert!(!orfas.contains(&entity_id));
907        Ok(())
908    }
909
910    #[test]
911    fn test_find_orphan_entity_ids_detecta_orfas() -> Resultado {
912        let (_tmp, conn) = setup_db()?;
913        let entity_id = upsert_entity(&conn, "global", &nova_entidade("sim-orfa"))?;
914
915        let orfas = find_orphan_entity_ids(&conn, Some("global"))?;
916        assert!(orfas.contains(&entity_id));
917        Ok(())
918    }
919
920    #[test]
921    fn test_find_orphan_entity_ids_sem_namespace_retorna_todas() -> Resultado {
922        let (_tmp, conn) = setup_db()?;
923        let id1 = upsert_entity(&conn, "ns-a", &nova_entidade("orfa-a"))?;
924        let id2 = upsert_entity(&conn, "ns-b", &nova_entidade("orfa-b"))?;
925
926        let orfas = find_orphan_entity_ids(&conn, None)?;
927        assert!(orfas.contains(&id1));
928        assert!(orfas.contains(&id2));
929        Ok(())
930    }
931
932    // ------------------------------------------------------------------ //
933    // list_entities / list_relationships_by_namespace
934    // ------------------------------------------------------------------ //
935
936    #[test]
937    fn test_list_entities_com_namespace() -> Resultado {
938        let (_tmp, conn) = setup_db()?;
939        upsert_entity(&conn, "le-ns", &nova_entidade("le-ent-1"))?;
940        upsert_entity(&conn, "le-ns", &nova_entidade("le-ent-2"))?;
941        upsert_entity(&conn, "outro-ns", &nova_entidade("le-ent-3"))?;
942
943        let lista = list_entities(&conn, Some("le-ns"))?;
944        assert_eq!(lista.len(), 2);
945        assert!(lista.iter().all(|e| e.namespace == "le-ns"));
946        Ok(())
947    }
948
949    #[test]
950    fn test_list_entities_sem_namespace_retorna_todas() -> Resultado {
951        let (_tmp, conn) = setup_db()?;
952        upsert_entity(&conn, "ns1", &nova_entidade("all-ent-1"))?;
953        upsert_entity(&conn, "ns2", &nova_entidade("all-ent-2"))?;
954
955        let lista = list_entities(&conn, None)?;
956        assert!(lista.len() >= 2);
957        Ok(())
958    }
959
960    #[test]
961    fn test_list_relationships_by_namespace_filtra_corretamente() -> Resultado {
962        let (_tmp, conn) = setup_db()?;
963        let id_a = upsert_entity(&conn, "rel-ns", &nova_entidade("lr-a"))?;
964        let id_b = upsert_entity(&conn, "rel-ns", &nova_entidade("lr-b"))?;
965
966        let rel = NewRelationship {
967            source: "lr-a".to_string(),
968            target: "lr-b".to_string(),
969            relation: "uses".to_string(),
970            strength: 0.5,
971            description: None,
972        };
973        upsert_relationship(&conn, "rel-ns", id_a, id_b, &rel)?;
974
975        let lista = list_relationships_by_namespace(&conn, Some("rel-ns"))?;
976        assert!(!lista.is_empty());
977        assert!(lista.iter().all(|r| r.namespace == "rel-ns"));
978        Ok(())
979    }
980
981    // ------------------------------------------------------------------ //
982    // delete_relationship_by_id / create_or_fetch_relationship
983    // ------------------------------------------------------------------ //
984
985    #[test]
986    fn test_delete_relationship_by_id_remove_relacao() -> Resultado {
987        let (_tmp, conn) = setup_db()?;
988        let id_a = upsert_entity(&conn, "global", &nova_entidade("dr-a"))?;
989        let id_b = upsert_entity(&conn, "global", &nova_entidade("dr-b"))?;
990
991        let rel = NewRelationship {
992            source: "dr-a".to_string(),
993            target: "dr-b".to_string(),
994            relation: "uses".to_string(),
995            strength: 0.5,
996            description: None,
997        };
998        let rel_id = upsert_relationship(&conn, "global", id_a, id_b, &rel)?;
999
1000        delete_relationship_by_id(&conn, rel_id)?;
1001
1002        let encontrada = find_relationship(&conn, id_a, id_b, "uses")?;
1003        assert!(encontrada.is_none(), "relação deve ter sido removida");
1004        Ok(())
1005    }
1006
1007    #[test]
1008    fn test_create_or_fetch_relationship_cria_nova() -> Resultado {
1009        let (_tmp, conn) = setup_db()?;
1010        let id_a = upsert_entity(&conn, "global", &nova_entidade("cf-a"))?;
1011        let id_b = upsert_entity(&conn, "global", &nova_entidade("cf-b"))?;
1012
1013        let (rel_id, criada) =
1014            create_or_fetch_relationship(&conn, "global", id_a, id_b, "uses", 0.5, None)?;
1015        assert!(rel_id > 0);
1016        assert!(criada);
1017        Ok(())
1018    }
1019
1020    #[test]
1021    fn test_create_or_fetch_relationship_retorna_existente() -> Resultado {
1022        let (_tmp, conn) = setup_db()?;
1023        let id_a = upsert_entity(&conn, "global", &nova_entidade("cf2-a"))?;
1024        let id_b = upsert_entity(&conn, "global", &nova_entidade("cf2-b"))?;
1025
1026        create_or_fetch_relationship(&conn, "global", id_a, id_b, "uses", 0.5, None)?;
1027        let (_, criada) =
1028            create_or_fetch_relationship(&conn, "global", id_a, id_b, "uses", 0.5, None)?;
1029        assert!(!criada, "segunda chamada deve retornar a relação existente");
1030        Ok(())
1031    }
1032
1033    // ------------------------------------------------------------------ //
1034    // serde alias: campo "type" aceito como sinônimo de "entity_type"
1035    // ------------------------------------------------------------------ //
1036
1037    #[test]
1038    fn aceita_campo_type_como_alias() -> Resultado {
1039        let json = r#"{"name": "X", "type": "concept"}"#;
1040        let ent: NewEntity = serde_json::from_str(json)?;
1041        assert_eq!(ent.entity_type, "concept");
1042        Ok(())
1043    }
1044
1045    #[test]
1046    fn aceita_campo_entity_type_canonico() -> Resultado {
1047        let json = r#"{"name": "X", "entity_type": "concept"}"#;
1048        let ent: NewEntity = serde_json::from_str(json)?;
1049        assert_eq!(ent.entity_type, "concept");
1050        Ok(())
1051    }
1052
1053    #[test]
1054    fn ambos_campos_presentes_gera_erro_de_duplicata() {
1055        // serde trata alias como nome alternativo do mesmo campo;
1056        // ter entity_type e type no mesmo JSON é duplicata e deve falhar
1057        let json = r#"{"name": "X", "entity_type": "A", "type": "B"}"#;
1058        let resultado: Result<NewEntity, _> = serde_json::from_str(json);
1059        assert!(
1060            resultado.is_err(),
1061            "ambos os campos no mesmo JSON é duplicata"
1062        );
1063    }
1064}