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