1use rustc_hash::FxHashMap;
2use std::collections::{BTreeMap, HashSet, VecDeque};
3use std::num::NonZeroUsize;
4use std::path::Path;
5use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};
6use std::time::Duration;
7
8use parking_lot::{Mutex, MutexGuard};
9use rusqlite::{Connection, OpenFlags, params};
10
11use crate::errors::{MCSError, Result};
12use crate::mutation::{
13 MutationContext, MutationRequest, MutationResult, MutationService, ObservationUpdate,
14};
15use crate::storage::{Durability, SqliteTuning};
16use crate::types::{
17 AttributeDelete, AttributeSet, Degree, Entity, EntityDescription, EntityInput, Observation,
18 ObservationInput, Relation, RelationDetail, RelationInput, RelationObservationUpdate,
19};
20
21const OBSERVATION_JSON: &str = "json_object('body',o.body,'createdAtUs',o.created_us,'occurredAtUs',o.occurred_us,'originEntityName',o.origin_entity_name)";
23
24const MAX_TRAVERSAL_ENTITIES: usize = 500_000;
27const MAX_TRAVERSAL_RELS: usize = 2_000_000;
28
29fn sqlite_err(e: rusqlite::Error) -> MCSError {
30 MCSError::IoError(std::io::Error::other(e))
31}
32
33const fn is_not_found(e: &rusqlite::Error) -> bool {
34 matches!(e, rusqlite::Error::QueryReturnedNoRows)
35}
36
37#[inline(always)]
38pub fn name_hash(name: &str) -> i64 {
39 let mut h: u64 = 0xcbf29ce484222325;
40 for b in name.bytes() {
41 h ^= u64::from(b);
42 h = h.wrapping_mul(0x100000001b3);
43 }
44 h as i64
45}
46
47fn entity_name_lookup(conn: &Connection, name: &str) -> Result<Option<i64>> {
48 let h = name_hash(name);
49 let mut stmt = conn
50 .prepare_cached("SELECT id FROM entity WHERE name_hash = ?1 AND name = ?2 AND flags = 0")
51 .map_err(sqlite_err)?;
52 match stmt.query_row(params![h, name], |row| row.get::<_, i64>(0)) {
53 Ok(id) => Ok(Some(id)),
54 Err(e) if is_not_found(&e) => Ok(None),
55 Err(e) => Err(sqlite_err(e)),
56 }
57}
58
59fn lookup_type_id(conn: &Connection, type_name: &str, kind: i64) -> Option<i64> {
63 conn.prepare_cached("SELECT id FROM type_dict WHERE kind = ?1 AND name = ?2")
64 .ok()?
65 .query_row(params![kind, type_name], |row| row.get::<_, i64>(0))
66 .ok()
67}
68
69fn attributes_for(
73 conn: &Connection,
74 owner_kind: &str,
75 owner_id: i64,
76) -> Result<BTreeMap<String, String>> {
77 let mut stmt = conn
78 .prepare_cached(
79 "SELECT key, value FROM attribute
80 WHERE owner_kind = ?1 AND owner_id = ?2
81 ORDER BY key",
82 )
83 .map_err(sqlite_err)?;
84 let rows = stmt
85 .query_map(params![owner_kind, owner_id], |row| {
86 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
87 })
88 .map_err(sqlite_err)?;
89 rows.collect::<rusqlite::Result<BTreeMap<String, String>>>()
90 .map_err(sqlite_err)
91}
92
93fn some_nonempty(map: BTreeMap<String, String>) -> Option<BTreeMap<String, String>> {
96 if map.is_empty() { None } else { Some(map) }
97}
98
99fn relation_details(
104 conn: &Connection,
105 triples: impl Iterator<Item = Relation>,
106) -> Result<Vec<RelationDetail>> {
107 use rusqlite::OptionalExtension;
108 let mut out = Vec::new();
109 for relation in triples {
110 let Some(mirror_id) = conn
111 .query_row(
112 "SELECT m.id FROM taxonomy_relation m
113 JOIN entity f ON f.id = m.from_id AND f.name = ?1 AND f.flags = 0
114 JOIN entity t ON t.id = m.to_id AND t.name = ?2 AND t.flags = 0
115 JOIN type_dict d ON d.id = m.type_id AND d.kind = 1 AND d.name = ?3
116 WHERE m.deleted = 0",
117 params![relation.from, relation.to, relation.relation_type],
118 |row| row.get(0),
119 )
120 .optional()
121 .map_err(sqlite_err)?
122 else {
123 continue;
124 };
125 let observations: Vec<Observation> = {
126 let mut stmt = conn
127 .prepare_cached(
128 "SELECT body, created_us, occurred_us
129 FROM relation_observation WHERE relation_id = ?1
130 ORDER BY idx, id",
131 )
132 .map_err(sqlite_err)?;
133 stmt.query_map([mirror_id], |row| {
134 Ok(Observation {
135 body: row.get(0)?,
136 created_at_us: Some(row.get(1)?),
137 occurred_at_us: row.get(2)?,
138 origin_entity_name: None,
139 })
140 })
141 .map_err(sqlite_err)?
142 .collect::<rusqlite::Result<Vec<_>>>()
143 .map_err(sqlite_err)?
144 };
145 out.push(RelationDetail {
146 from: relation.from,
147 to: relation.to,
148 relation_type: relation.relation_type,
149 observations,
150 attributes: attributes_for(conn, "relation", mirror_id)?,
151 });
152 }
153 Ok(out)
154}
155
156fn read_graph_stat(conn: &Connection, key: &str) -> Result<i64> {
157 conn.query_row(
158 "SELECT value FROM graph_stat WHERE key = ?1",
159 params![key],
160 |row| row.get(0),
161 )
162 .map_err(sqlite_err)
163}
164
165fn select_all_types(conn: &Connection, kind: i64) -> Result<Vec<(String, usize)>> {
166 let mut stmt = conn
167 .prepare_cached(
168 "SELECT name, count FROM type_dict WHERE kind = ?1 AND count > 0 ORDER BY count DESC",
169 )
170 .map_err(sqlite_err)?;
171 let rows = stmt
172 .query_map(params![kind], |row| {
173 Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)? as usize))
174 })
175 .map_err(sqlite_err)?
176 .filter_map(|r| r.ok())
177 .collect();
178 Ok(rows)
179}
180
181fn select_type_catalog(
186 conn: &Connection,
187 kind: i64,
188) -> Result<Vec<(String, usize, Option<String>)>> {
189 let mut stmt = conn
190 .prepare_cached(
191 "SELECT name, count, desc FROM type_dict
192 WHERE kind = ?1 AND (count > 0 OR desc IS NOT NULL)
193 ORDER BY count DESC",
194 )
195 .map_err(sqlite_err)?;
196 let rows = stmt
197 .query_map(params![kind], |row| {
198 Ok((
199 row.get::<_, String>(0)?,
200 row.get::<_, i64>(1)? as usize,
201 row.get::<_, Option<String>>(2)?,
202 ))
203 })
204 .map_err(sqlite_err)?
205 .filter_map(|r| r.ok())
206 .collect();
207 Ok(rows)
208}
209
210fn int_csv(ids: &[i64]) -> String {
218 use std::fmt::Write as _;
219 let mut s = String::with_capacity(ids.len() * 8);
220 for (i, id) in ids.iter().enumerate() {
221 if i > 0 {
222 s.push(',');
223 }
224 let _ = write!(s, "{id}");
225 }
226 s
227}
228
229fn rel_values_literal(rels: &HashSet<(i64, i64, i64)>) -> String {
235 use std::fmt::Write as _;
236 let mut s = String::with_capacity(rels.len() * 16);
237 for (i, (f, t, tp)) in rels.iter().enumerate() {
238 if i > 0 {
239 s.push(',');
240 }
241 let _ = write!(s, "({f},{t},{tp})");
242 }
243 s
244}
245
246fn batch_entities_by_ids(conn: &Connection, ids: &[i64]) -> FxHashMap<i64, Entity> {
250 let mut map = FxHashMap::default();
251 if ids.is_empty() {
252 return map;
253 }
254 let sql = format!(
255 "SELECT e.id, e.name, t.name,
256 COALESCE((SELECT json_group_array({OBSERVATION_JSON} ORDER BY o.idx, o.id)
257 FROM observation o WHERE o.entity_id = e.id), '[]')
258 FROM entity e JOIN type_dict t ON t.id = e.type_id
259 WHERE e.id IN ({}) AND e.flags = 0",
260 int_csv(ids)
261 );
262 if let Ok(mut stmt) = conn.prepare(&sql)
263 && let Ok(rows) = stmt.query_map([], |row| {
264 Ok((
265 row.get::<_, i64>(0)?,
266 row.get::<_, String>(1)?,
267 row.get::<_, String>(2)?,
268 row.get::<_, String>(3)?,
269 ))
270 })
271 {
272 for (id, name, etype, obs_json) in rows.flatten() {
273 let observations: Vec<Observation> =
274 serde_json::from_str(&obs_json).unwrap_or_default();
275 map.insert(
276 id,
277 Entity {
278 name,
279 entity_type: etype,
280 observations,
281 attributes: None,
282 },
283 );
284 }
285 }
286 map
287}
288
289fn batch_entity_lite_by_ids(
295 conn: &Connection,
296 ids: &[i64],
297) -> FxHashMap<i64, (String, String, i64)> {
298 let mut map = FxHashMap::default();
299 if ids.is_empty() {
300 return map;
301 }
302 let sql = format!(
303 "SELECT e.id, e.name, t.name, e.obs_count
304 FROM entity e JOIN type_dict t ON t.id = e.type_id
305 WHERE e.id IN ({}) AND e.flags = 0",
306 int_csv(ids)
307 );
308 if let Ok(mut stmt) = conn.prepare(&sql)
309 && let Ok(rows) = stmt.query_map([], |row| {
310 Ok((
311 row.get::<_, i64>(0)?,
312 row.get::<_, String>(1)?,
313 row.get::<_, String>(2)?,
314 row.get::<_, i64>(3)?,
315 ))
316 })
317 {
318 for (id, name, etype, oc) in rows.flatten() {
319 map.insert(id, (name, etype, oc));
320 }
321 }
322 map
323}
324
325fn fts_candidate_ids(conn: &Connection, query: &str, cap: usize) -> Vec<i64> {
329 let mut ids: Vec<i64> = Vec::new();
330 let mut seen: HashSet<i64> = HashSet::new();
331 let cap_i64 = cap as i64;
332
333 if let Ok(mut stmt) =
334 conn.prepare("SELECT rowid FROM name_fts WHERE name_fts MATCH ?1 ORDER BY rank LIMIT ?2")
335 && let Ok(rows) = stmt.query_map(params![query, cap_i64], |row| row.get::<_, i64>(0))
336 {
337 for id in rows.flatten() {
338 if seen.insert(id) {
339 ids.push(id);
340 }
341 }
342 }
343
344 if let Ok(mut stmt) = conn.prepare(
345 "SELECT entity_id FROM obs_fts JOIN observation ON obs_fts.rowid = observation.id
346 WHERE obs_fts MATCH ?1
347 GROUP BY entity_id
348 LIMIT ?2",
349 ) && let Ok(rows) = stmt.query_map(params![query, cap_i64], |row| row.get::<_, i64>(0))
350 {
351 for id in rows.flatten() {
352 if seen.insert(id) {
353 ids.push(id);
354 }
355 }
356 }
357
358 ids
359}
360
361#[derive(Debug, Clone, Copy, PartialEq, Eq)]
363pub enum Direction {
364 Outgoing,
365 Incoming,
366 Both,
367}
368
369impl Direction {
370 pub fn parse(s: Option<&str>) -> Self {
371 match s {
372 Some("OUTGOING") => Direction::Outgoing,
373 Some("INCOMING") => Direction::Incoming,
374 _ => Direction::Both,
375 }
376 }
377}
378
379pub fn push_json_str(buf: &mut String, raw: &str) {
382 buf.push('"');
383 let mut start = 0;
384 let bytes = raw.as_bytes();
385 for (i, &b) in bytes.iter().enumerate() {
386 let esc: u8 = match b {
387 b'"' => b'"',
388 b'\\' => b'\\',
389 b'\n' => b'n',
390 b'\r' => b'r',
391 b'\t' => b't',
392 0x08 => b'b',
393 0x0C => b'f',
394 0x00..=0x07 | 0x0B | 0x0E..=0x1F => continue, _ => continue,
396 };
397 buf.push_str(&raw[start..i]);
398 buf.push('\\');
399 buf.push(esc as char);
400 start = i + 1;
401 }
402 for (i, &b) in bytes.iter().enumerate().skip(start) {
404 if b <= 0x07 || b == 0x0B || (0x0E..=0x1F).contains(&b) {
405 buf.push_str(&raw[start..i]);
406 write_escape_unicode(buf, b);
407 start = i + 1;
408 }
409 }
410 buf.push_str(&raw[start..]);
411 buf.push('"');
412}
413
414#[inline(never)]
415fn write_escape_unicode(buf: &mut String, b: u8) {
416 use std::fmt::Write;
417 write!(buf, "\\u{:04x}", b).unwrap();
418}
419
420pub(crate) struct TxGuard<'a> {
423 conn: &'a Connection,
424 done: bool,
425}
426
427impl<'a> TxGuard<'a> {
428 pub(crate) fn begin(conn: &'a Connection) -> Result<Self> {
429 conn.execute_batch("BEGIN IMMEDIATE").map_err(sqlite_err)?;
434 Ok(Self { conn, done: false })
435 }
436
437 pub(crate) fn commit(mut self) -> Result<()> {
438 self.conn.execute_batch("COMMIT").map_err(sqlite_err)?;
439 self.done = true;
440 Ok(())
441 }
442}
443
444impl Drop for TxGuard<'_> {
445 fn drop(&mut self) {
446 if !self.done {
447 let _ = self.conn.execute_batch("ROLLBACK");
448 }
449 }
450}
451
452struct ReaderPool {
459 conns: Vec<Mutex<Connection>>,
460 next: AtomicUsize,
461}
462
463impl ReaderPool {
464 fn get(&self) -> MutexGuard<'_, Connection> {
468 for c in &self.conns {
469 if let Some(g) = c.try_lock() {
470 return g;
471 }
472 }
473 let i = self.next.fetch_add(1, Ordering::Relaxed) % self.conns.len();
474 self.conns[i].lock()
475 }
476}
477
478pub struct GraphHandle {
481 pub(crate) writer: Mutex<Connection>,
484 readers: ReaderPool,
486 seq_entity: AtomicI64,
487 seq_obs: AtomicI64,
488 seq_rel_obs: AtomicI64,
489}
490
491fn open_reader(path: &Path, tuning: &SqliteTuning) -> Result<Connection> {
498 let conn = Connection::open_with_flags(
499 path,
500 OpenFlags::SQLITE_OPEN_READ_WRITE
501 | OpenFlags::SQLITE_OPEN_NO_MUTEX
502 | OpenFlags::SQLITE_OPEN_URI,
503 )
504 .map_err(sqlite_err)?;
505 conn.busy_timeout(Duration::from_millis(tuning.busy_timeout_ms))
506 .map_err(sqlite_err)?;
507 conn.execute_batch(&format!(
508 "PRAGMA query_only = ON;
509 PRAGMA cache_size = -{};
510 PRAGMA temp_store = MEMORY;
511 PRAGMA mmap_size = {};",
512 tuning.cache_size_kb, tuning.mmap_size
513 ))
514 .map_err(sqlite_err)?;
515 Ok(conn)
516}
517
518impl GraphHandle {
519 pub fn new(
520 path: &Path,
521 durability: Durability,
522 tuning: SqliteTuning,
523 _lru_cache_size: NonZeroUsize,
524 read_pool_size: usize,
525 ) -> Result<Self> {
526 let conn = Connection::open(path).map_err(sqlite_err)?;
527 conn.busy_timeout(Duration::from_millis(tuning.busy_timeout_ms))
530 .map_err(sqlite_err)?;
531
532 conn.execute_batch(&format!(
538 "PRAGMA page_size = {};
539 PRAGMA auto_vacuum = INCREMENTAL;",
540 tuning.page_size
541 ))
542 .map_err(sqlite_err)?;
543
544 conn.execute_batch(&format!(
545 "PRAGMA journal_mode = WAL;
546 PRAGMA foreign_keys = OFF;
547 PRAGMA cache_size = -{};
548 PRAGMA temp_store = MEMORY;
549 PRAGMA busy_timeout = {};
550 PRAGMA synchronous = NORMAL;
551 PRAGMA journal_size_limit = {};",
552 tuning.cache_size_kb, tuning.busy_timeout_ms, tuning.journal_size_limit
553 ))
554 .map_err(sqlite_err)?;
555
556 crate::schema::initialize_database(&conn)?;
557
558 conn.execute_batch(&format!("PRAGMA mmap_size = {};", tuning.mmap_size))
559 .map_err(sqlite_err)?;
560
561 let sync_pragma = match durability {
562 Durability::Sync => "PRAGMA synchronous = FULL",
563 Durability::Async => "PRAGMA synchronous = NORMAL",
564 };
565 conn.execute_batch(sync_pragma).map_err(sqlite_err)?;
566
567 conn.execute_batch("PRAGMA analysis_limit = 400;")
570 .map_err(sqlite_err)?;
571
572 conn.execute_batch("PRAGMA optimize;").map_err(sqlite_err)?;
573
574 let seq_entity = read_graph_stat(&conn, "entity_seq").unwrap_or(0);
575 let seq_obs = read_graph_stat(&conn, "obs_seq").unwrap_or(0);
576 let seq_rel_obs = read_graph_stat(&conn, "rel_obs_seq").unwrap_or(0);
577
578 let pool_size = read_pool_size.max(1);
581 let mut conns = Vec::with_capacity(pool_size);
582 for _ in 0..pool_size {
583 conns.push(Mutex::new(open_reader(path, &tuning)?));
584 }
585 let readers = ReaderPool {
586 conns,
587 next: AtomicUsize::new(0),
588 };
589
590 Ok(Self {
591 writer: Mutex::new(conn),
592 readers,
593 seq_entity: AtomicI64::new(seq_entity),
594 seq_obs: AtomicI64::new(seq_obs),
595 seq_rel_obs: AtomicI64::new(seq_rel_obs),
596 })
597 }
598
599 pub(crate) fn next_entity_id(&self) -> i64 {
600 self.seq_entity.fetch_add(1, Ordering::Relaxed) + 1
601 }
602
603 pub(crate) fn refresh_seqs(&self, conn: &Connection) -> Result<()> {
606 self.seq_entity
607 .fetch_max(read_graph_stat(conn, "entity_seq")?, Ordering::Relaxed);
608 self.seq_obs
609 .fetch_max(read_graph_stat(conn, "obs_seq")?, Ordering::Relaxed);
610 self.seq_rel_obs
611 .fetch_max(read_graph_stat(conn, "rel_obs_seq")?, Ordering::Relaxed);
612 Ok(())
613 }
614
615 pub(crate) fn next_obs_id(&self) -> i64 {
616 self.seq_obs.fetch_add(1, Ordering::Relaxed) + 1
617 }
618
619 pub(crate) fn next_rel_obs_id(&self) -> i64 {
622 self.seq_rel_obs.fetch_add(1, Ordering::Relaxed) + 1
623 }
624
625 fn get_entity_id(&self, conn: &Connection, name: &str) -> Result<Option<(i64, i64, i64, i64)>> {
626 use rusqlite::OptionalExtension;
627 conn.query_row(
628 "SELECT id, type_id, out_deg, in_deg FROM entity WHERE name_hash = ?1 AND name = ?2 AND flags = 0",
629 params![name_hash(name), name],
630 |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
631 ).optional().map_err(sqlite_err)
632 }
633
634 pub(crate) fn sync_seqs(&self, conn: &Connection) -> Result<()> {
635 let seq_e = self.seq_entity.load(Ordering::Relaxed);
636 let seq_o = self.seq_obs.load(Ordering::Relaxed);
637 let seq_r = self.seq_rel_obs.load(Ordering::Relaxed);
638 conn.execute(
639 "UPDATE graph_stat SET value = CASE key
640 WHEN 'entity_seq' THEN ?1
641 WHEN 'obs_seq' THEN ?2
642 WHEN 'rel_obs_seq' THEN ?3
643 ELSE value END
644 WHERE key IN ('entity_seq', 'obs_seq', 'rel_obs_seq')",
645 params![seq_e, seq_o, seq_r],
646 )
647 .map_err(sqlite_err)?;
648 Ok(())
649 }
650
651 pub fn get_entity(&self, name: &str) -> Result<Option<Entity>> {
654 let conn = self.readers.get();
655 let tx = conn.unchecked_transaction().map_err(sqlite_err)?;
658 let entity = match crate::mutation::read_entity(&tx, name)? {
659 Some(snapshot) => {
660 let mut entity = snapshot.entity();
661 entity.attributes =
662 some_nonempty(attributes_for(&tx, "entity", snapshot.entity_id)?);
663 Some(entity)
664 }
665 None => None,
666 };
667 tx.commit().map_err(sqlite_err)?;
668 Ok(entity)
669 }
670
671 fn mutate(&self, request: MutationRequest) -> Result<MutationResult> {
672 MutationService::new(self)
673 .apply_with_result(request, MutationContext::local())
674 .map(|(_, result)| result)
675 }
676
677 pub fn create_entities(&self, entities: &[EntityInput]) -> Result<Vec<Entity>> {
678 match self.mutate(MutationRequest::CreateEntities {
679 entities: entities.to_vec(),
680 })? {
681 MutationResult::Entities(result) => Ok(result),
682 _ => unreachable!("create_entities always returns entities"),
683 }
684 }
685
686 pub fn upsert_entities(&self, entities: &[EntityInput]) -> Result<Vec<Entity>> {
687 match self.mutate(MutationRequest::UpsertEntities {
688 entities: entities.to_vec(),
689 })? {
690 MutationResult::Entities(result) => Ok(result),
691 _ => unreachable!("upsert_entities always returns entities"),
692 }
693 }
694
695 pub fn delete_entities(&self, names: &[String]) -> Result<()> {
696 self.mutate(MutationRequest::DeleteEntities {
697 names: names.to_vec(),
698 })
699 .map(|_| ())
700 }
701
702 pub fn create_relations(&self, relations: &[RelationInput]) -> Result<Vec<Relation>> {
703 match self.mutate(MutationRequest::CreateRelations {
704 relations: relations.to_vec(),
705 })? {
706 MutationResult::Relations(result) => Ok(result),
707 _ => unreachable!("create_relations always returns relations"),
708 }
709 }
710
711 pub fn delete_relations(&self, relations: &[Relation]) -> Result<()> {
712 self.mutate(MutationRequest::DeleteRelations {
713 relations: relations.to_vec(),
714 })
715 .map(|_| ())
716 }
717
718 pub fn add_relation_observations(
719 &self,
720 from: &str,
721 to: &str,
722 relation_type: &str,
723 contents: &[ObservationInput],
724 ) -> Result<Vec<Observation>> {
725 match self.mutate(MutationRequest::AddRelationObservations {
726 relations: vec![RelationObservationUpdate {
727 relation: Relation {
728 from: from.into(),
729 to: to.into(),
730 relation_type: relation_type.into(),
731 },
732 contents: contents.to_vec(),
733 }],
734 })? {
735 MutationResult::RelationObservations(mut result) => {
736 Ok(result.remove(0).added_observations)
737 }
738 _ => unreachable!("add_relation_observations always returns observations"),
739 }
740 }
741
742 pub fn delete_relation_observations(
743 &self,
744 from: &str,
745 to: &str,
746 relation_type: &str,
747 observations: &[ObservationInput],
748 ) -> Result<()> {
749 self.mutate(MutationRequest::DeleteRelationObservations {
750 relations: vec![RelationObservationUpdate {
751 relation: Relation {
752 from: from.into(),
753 to: to.into(),
754 relation_type: relation_type.into(),
755 },
756 contents: observations.to_vec(),
757 }],
758 })
759 .map(|_| ())
760 }
761
762 pub fn set_attributes(&self, targets: &[AttributeSet]) -> Result<()> {
763 self.mutate(MutationRequest::SetAttributes {
764 targets: targets.to_vec(),
765 })
766 .map(|_| ())
767 }
768
769 pub fn delete_attributes(&self, targets: &[AttributeDelete]) -> Result<()> {
770 self.mutate(MutationRequest::DeleteAttributes {
771 targets: targets.to_vec(),
772 })
773 .map(|_| ())
774 }
775
776 pub fn add_observations(
777 &self,
778 entity_name: &str,
779 contents: &[ObservationInput],
780 ) -> Result<Vec<Observation>> {
781 match self.mutate(MutationRequest::AddObservations {
782 observations: vec![ObservationUpdate {
783 entity_name: entity_name.into(),
784 contents: contents.to_vec(),
785 }],
786 })? {
787 MutationResult::Observations(mut result) => Ok(result.remove(0).added_observations),
788 _ => unreachable!("add_observations always returns observations"),
789 }
790 }
791
792 pub fn delete_observations(
793 &self,
794 entity_name: &str,
795 observations: &[ObservationInput],
796 ) -> Result<()> {
797 self.mutate(MutationRequest::DeleteObservations {
798 observations: vec![ObservationUpdate {
799 entity_name: entity_name.into(),
800 contents: observations.to_vec(),
801 }],
802 })
803 .map(|_| ())
804 }
805
806 pub fn merge_entities(&self, source: &str, target: &str) -> Result<Entity> {
807 match self.mutate(MutationRequest::MergeEntities {
808 source: source.into(),
809 target: target.into(),
810 })? {
811 MutationResult::Entity(result) => Ok(result),
812 _ => unreachable!("merge_entities always returns an entity"),
813 }
814 }
815
816 pub fn rename_entity(&self, old_name: &str, new_name: &str) -> Result<Entity> {
817 match self.mutate(MutationRequest::RenameEntity {
818 old_name: old_name.into(),
819 new_name: new_name.into(),
820 })? {
821 MutationResult::Entity(result) => Ok(result),
822 _ => unreachable!("rename_entity always returns an entity"),
823 }
824 }
825
826 pub fn code_purge_file(&self, rel_path: &str) -> Result<usize> {
828 match self.mutate(MutationRequest::PurgeDefinedEntities {
829 name: rel_path.into(),
830 })? {
831 MutationResult::Count(count) => Ok(count),
832 _ => unreachable!("purge always returns a count"),
833 }
834 }
835
836 pub fn search_nodes_filtered(
837 &self,
838 query: &str,
839 filter_type: Option<&str>,
840 offset: usize,
841 limit: usize,
842 ) -> Vec<Entity> {
843 if query.is_empty() {
844 return Vec::new();
845 }
846 let conn = self.readers.get();
847
848 let cap = offset.saturating_add(limit);
853 let candidates = fts_candidate_ids(&conn, query, cap);
854 let mut by_id = batch_entities_by_ids(&conn, &candidates);
855
856 let mut results = Vec::new();
857 let mut count: usize = 0;
858 for eid in candidates {
859 let Some(entity) = by_id.remove(&eid) else {
860 continue;
861 };
862 if let Some(ft) = filter_type
863 && !ft.is_empty()
864 && entity.entity_type != ft
865 {
866 continue;
867 }
868 if count < offset {
869 count += 1;
870 continue;
871 }
872 if results.len() >= limit {
873 break;
874 }
875 results.push(entity);
876 count += 1;
877 }
878
879 results
880 }
881
882 pub fn search_nodes_lite_json(
889 &self,
890 query: &str,
891 filter_type: Option<&str>,
892 offset: usize,
893 limit: usize,
894 ) -> (String, usize, bool) {
895 use std::fmt::Write as _;
896 if query.is_empty() {
897 return ("[]".to_string(), 0, false);
898 }
899 let conn = self.readers.get();
900 let cap = offset.saturating_add(limit).saturating_add(1);
902 let candidates = fts_candidate_ids(&conn, query, cap);
903 let by_id = batch_entity_lite_by_ids(&conn, &candidates);
904 let ft = filter_type.filter(|s| !s.is_empty());
905
906 let mut arr = String::from("[");
907 let mut count: usize = 0; let mut returned: usize = 0;
909 let mut has_more = false;
910 for eid in candidates {
911 let Some((name, etype, oc)) = by_id.get(&eid) else {
912 continue;
913 };
914 if let Some(f) = ft
915 && etype != f
916 {
917 continue;
918 }
919 if count < offset {
920 count += 1;
921 continue;
922 }
923 if returned >= limit {
924 has_more = true;
925 break;
926 }
927 if returned > 0 {
928 arr.push(',');
929 }
930 arr.push_str("{\"name\":");
931 push_json_str(&mut arr, name);
932 arr.push_str(",\"entityType\":");
933 push_json_str(&mut arr, etype);
934 let _ = write!(arr, ",\"obsCount\":{oc}}}");
935 returned += 1;
936 count += 1;
937 }
938 arr.push(']');
939 (arr, returned, has_more)
940 }
941
942 pub fn read_graph_filtered(
943 &self,
944 filter_type: Option<&str>,
945 offset: usize,
946 limit: usize,
947 ) -> Result<String> {
948 self.read_graph_page(filter_type, offset, limit, true)
949 .map(|(json, _)| json)
950 }
951
952 pub fn read_graph_filtered_lite(
958 &self,
959 filter_type: Option<&str>,
960 offset: usize,
961 limit: usize,
962 ) -> Result<(String, usize)> {
963 self.read_graph_page(filter_type, offset, limit, false)
964 }
965
966 fn read_graph_page(
967 &self,
968 filter_type: Option<&str>,
969 offset: usize,
970 limit: usize,
971 include_obs: bool,
972 ) -> Result<(String, usize)> {
973 let conn = self.readers.get();
974
975 let limit_sql: i64 = if limit == usize::MAX {
976 -1
977 } else {
978 limit.min(i64::MAX as usize) as i64
979 };
980 let offset_sql: i64 = offset as i64;
981
982 let filter = filter_type.filter(|ft| !ft.is_empty());
988 let ids: Vec<i64> = if let Some(ft) = filter {
989 let mut stmt = conn
990 .prepare_cached(
991 "SELECT e.id FROM entity e
992 WHERE e.type_id = (SELECT id FROM type_dict WHERE kind = 0 AND name = ?1)
993 AND e.flags = 0
994 ORDER BY e.id LIMIT ?2 OFFSET ?3",
995 )
996 .map_err(sqlite_err)?;
997 stmt.query_map(params![ft, limit_sql, offset_sql], |r| r.get::<_, i64>(0))
998 .map_err(sqlite_err)?
999 .filter_map(|r| r.ok())
1000 .collect()
1001 } else {
1002 let mut stmt = conn
1003 .prepare_cached(
1004 "SELECT e.id FROM entity e WHERE e.flags = 0
1005 ORDER BY e.id LIMIT ?1 OFFSET ?2",
1006 )
1007 .map_err(sqlite_err)?;
1008 stmt.query_map(params![limit_sql, offset_sql], |r| r.get::<_, i64>(0))
1009 .map_err(sqlite_err)?
1010 .filter_map(|r| r.ok())
1011 .collect()
1012 };
1013
1014 if ids.is_empty() {
1015 return Ok((r#"{"entities":[],"relations":[]}"#.to_string(), 0));
1016 }
1017
1018 let idlist = int_csv(&ids);
1023 let returned = ids.len();
1024
1025 let obs_field = if include_obs {
1030 format!("'observations', COALESCE((SELECT json_group_array({OBSERVATION_JSON} ORDER BY o.idx, o.id)
1031 FROM observation o WHERE o.entity_id = e.id), json('[]'))")
1032 } else {
1033 "'obsCount', e.obs_count".to_owned()
1034 };
1035
1036 let entities_json: String = {
1037 let sql = format!(
1038 "SELECT COALESCE(json_group_array(json_object(
1039 'name', e.name,
1040 'entityType', t.name,
1041 {obs_field}
1042 ) ORDER BY e.id), json('[]'))
1043 FROM entity e
1044 JOIN type_dict t ON t.id = e.type_id
1045 WHERE e.id IN ({idlist}) AND e.flags = 0"
1046 );
1047 conn.query_row(&sql, [], |row| row.get::<_, String>(0))
1048 .map_err(sqlite_err)?
1049 };
1050
1051 let relations_json: String = {
1052 let sql = format!(
1053 "SELECT COALESCE(json_group_array(json_object(
1054 'from', e1.name,
1055 'to', e2.name,
1056 'relationType', t.name
1057 )), json('[]'))
1058 FROM relation r
1059 JOIN entity e1 ON e1.id = r.from_id
1060 JOIN entity e2 ON e2.id = r.to_id
1061 JOIN type_dict t ON t.id = r.type_id
1062 WHERE r.from_id IN ({idlist}) AND r.to_id IN ({idlist})
1063 AND e1.flags = 0 AND e2.flags = 0"
1064 );
1065 conn.query_row(&sql, [], |row| row.get::<_, String>(0))
1066 .map_err(sqlite_err)?
1067 };
1068
1069 let mut out = String::with_capacity(32 + entities_json.len() + relations_json.len());
1070 out.push_str("{\"entities\":");
1071 out.push_str(&entities_json);
1072 out.push_str(",\"relations\":");
1073 out.push_str(&relations_json);
1074 out.push('}');
1075 Ok((out, returned))
1076 }
1077
1078 pub fn open_nodes(&self, names: &[String]) -> String {
1079 let conn = self.readers.get();
1080 let mut entity_ids: Vec<i64> = Vec::new();
1081
1082 for name in names {
1083 let h = name_hash(name);
1084 if let Ok(Some(id)) = conn
1085 .query_row(
1086 "SELECT id FROM entity WHERE name_hash = ?1 AND name = ?2 AND flags = 0",
1087 params![h, name],
1088 |row| row.get::<_, i64>(0),
1089 )
1090 .map(Some)
1091 .or_else(|e| {
1092 if is_not_found(&e) {
1093 Ok(None)
1094 } else {
1095 Err(sqlite_err(e))
1096 }
1097 })
1098 {
1099 entity_ids.push(id);
1100 }
1101 }
1102
1103 if entity_ids.is_empty() {
1104 return r#"{"entities":[],"relations":[]}"#.to_string();
1105 }
1106
1107 let placeholders: Vec<String> = entity_ids.iter().map(|_| "?".to_string()).collect();
1108 let ids_str = placeholders.join(",");
1109
1110 let entities_json: String = {
1111 let sql = format!(
1112 "SELECT COALESCE(json_group_array(json_object(
1113 'name', e.name,
1114 'entityType', t.name,
1115 'observations', COALESCE((
1116 SELECT json_group_array({OBSERVATION_JSON} ORDER BY o.idx, o.id)
1117 FROM observation o WHERE o.entity_id = e.id
1118 ), json('[]'))
1119 ) ORDER BY e.id), json('[]'))
1120 FROM entity e
1121 JOIN type_dict t ON t.id = e.type_id
1122 WHERE e.id IN ({ids_str}) AND e.flags = 0"
1123 );
1124 conn.query_row(&sql, rusqlite::params_from_iter(&entity_ids), |row| {
1125 row.get::<_, String>(0)
1126 })
1127 .unwrap_or_else(|_| "[]".to_string())
1128 };
1129
1130 let relations_json: String = {
1131 let sql = format!(
1132 "SELECT COALESCE(json_group_array(json_object(
1133 'from', e1.name,
1134 'to', e2.name,
1135 'relationType', t.name
1136 )), json('[]'))
1137 FROM relation r
1138 JOIN entity e1 ON e1.id = r.from_id
1139 JOIN entity e2 ON e2.id = r.to_id
1140 JOIN type_dict t ON t.id = r.type_id
1141 WHERE (r.from_id IN ({ids_str}) OR r.to_id IN ({ids_str}))
1142 AND e1.flags = 0 AND e2.flags = 0"
1143 );
1144 let all_params: Vec<&dyn rusqlite::types::ToSql> = entity_ids
1145 .iter()
1146 .map(|id| id as &dyn rusqlite::types::ToSql)
1147 .chain(
1148 entity_ids
1149 .iter()
1150 .map(|id| id as &dyn rusqlite::types::ToSql),
1151 )
1152 .collect();
1153 let mut stmt = conn.prepare(&sql).unwrap();
1154 stmt.query_row(all_params.as_slice(), |row| row.get::<_, String>(0))
1155 .unwrap_or_else(|_| "[]".to_string())
1156 };
1157
1158 let mut out = String::with_capacity(32 + entities_json.len() + relations_json.len());
1159 out.push_str("{\"entities\":");
1160 out.push_str(&entities_json);
1161 out.push_str(",\"relations\":");
1162 out.push_str(&relations_json);
1163 out.push('}');
1164 out
1165 }
1166
1167 pub fn entities_exist(&self, names: &[String]) -> Result<Vec<bool>> {
1168 let conn = self.readers.get();
1169 let mut results = Vec::with_capacity(names.len());
1170 for name in names {
1171 let h = name_hash(name);
1172 let exists: bool = conn
1173 .query_row(
1174 "SELECT 1 FROM entity WHERE name_hash = ?1 AND name = ?2 AND flags = 0",
1175 params![h, name],
1176 |_| Ok(()),
1177 )
1178 .is_ok();
1179 results.push(exists);
1180 }
1181 Ok(results)
1182 }
1183
1184 pub fn degree(&self, name: &str, direction: Direction) -> Result<usize> {
1185 let conn = self.readers.get();
1186 let (_, _, out_d, in_d) = match self.get_entity_id(&conn, name)? {
1187 Some(v) => v,
1188 None => {
1189 return Err(MCSError::InvalidParams(format!(
1190 "Entity '{name}' not found"
1191 )));
1192 }
1193 };
1194 Ok(match direction {
1195 Direction::Outgoing => out_d as usize,
1196 Direction::Incoming => in_d as usize,
1197 Direction::Both => (out_d + in_d) as usize,
1198 })
1199 }
1200
1201 pub fn get_entity_count(&self) -> Result<usize> {
1202 let conn = self.readers.get();
1203 read_graph_stat(&conn, "entities")
1204 .map(|v| v as usize)
1205 .map_err(|_| MCSError::MemoryError("Failed to read entity count".into()))
1206 }
1207
1208 pub fn get_relation_count(&self) -> Result<usize> {
1209 let conn = self.readers.get();
1210 read_graph_stat(&conn, "relations")
1211 .map(|v| v as usize)
1212 .map_err(|_| MCSError::MemoryError("Failed to read relation count".into()))
1213 }
1214
1215 pub fn search_relations(
1216 &self,
1217 from: Option<&str>,
1218 to: Option<&str>,
1219 rtype: Option<&str>,
1220 query: Option<&str>,
1221 limit: Option<usize>,
1222 ) -> Result<Vec<RelationDetail>> {
1223 let conn = self.readers.get();
1224
1225 let from_id = from
1231 .filter(|f| !f.is_empty())
1232 .map(|f| entity_name_lookup(&conn, f).ok().flatten().unwrap_or(-1));
1233 let to_id = to
1234 .filter(|t| !t.is_empty())
1235 .map(|t| entity_name_lookup(&conn, t).ok().flatten().unwrap_or(-1));
1236 let type_id = rtype
1237 .filter(|rt| !rt.is_empty())
1238 .map(|rt| lookup_type_id(&conn, rt, 1).unwrap_or(-1));
1239
1240 if let Some(query) = query.filter(|q| !q.trim().is_empty()) {
1245 let mut sql = String::from(
1246 "SELECT f.name, t.name, d.name
1247 FROM rel_obs_fts ft
1248 JOIN relation_observation ro ON ro.id = ft.rowid
1249 JOIN taxonomy_relation m ON m.id = ro.relation_id
1250 JOIN entity f ON f.id = m.from_id
1251 JOIN entity t ON t.id = m.to_id
1252 JOIN type_dict d ON d.id = m.type_id
1253 WHERE rel_obs_fts MATCH ?1 AND m.deleted = 0
1254 AND f.flags = 0 AND t.flags = 0",
1255 );
1256 if let Some(fid) = from_id {
1261 sql.push_str(&format!(" AND m.from_id = {fid}"));
1262 }
1263 if let Some(tid) = to_id {
1264 sql.push_str(&format!(" AND m.to_id = {tid}"));
1265 }
1266 if let Some(tpid) = type_id {
1267 sql.push_str(&format!(" AND m.type_id = {tpid}"));
1268 }
1269 sql.push_str(" ORDER BY rank");
1270 if let Some(lim) = limit
1271 && lim > 0
1272 {
1273 sql.push_str(&format!(" LIMIT {lim}"));
1274 }
1275 let mut triples: Vec<Relation> = Vec::new();
1276 let mut seen: HashSet<(String, String, String)> = HashSet::new();
1277 let mut stmt = conn.prepare(&sql).map_err(sqlite_err)?;
1278 let rows = stmt
1279 .query_map(params![query], |row| {
1280 Ok((
1281 row.get::<_, String>(0)?,
1282 row.get::<_, String>(1)?,
1283 row.get::<_, String>(2)?,
1284 ))
1285 })
1286 .map_err(sqlite_err)?;
1287 for row in rows {
1288 let (from, to, relation_type) = row.map_err(sqlite_err)?;
1289 if seen.insert((from.clone(), to.clone(), relation_type.clone())) {
1290 triples.push(Relation {
1291 from,
1292 to,
1293 relation_type,
1294 });
1295 }
1296 }
1297 return relation_details(&conn, triples.into_iter());
1298 }
1299
1300 let mut triples: Vec<Relation> = Vec::new();
1304 match (from_id, to_id, type_id) {
1305 (Some(fid), Some(tid), Some(tpid)) => {
1306 if let Ok(mut stmt) = conn.prepare_cached(
1307 "SELECT e1.name, e2.name, t.name
1308 FROM relation r
1309 JOIN entity e1 ON e1.id = r.from_id
1310 JOIN entity e2 ON e2.id = r.to_id
1311 JOIN type_dict t ON t.id = r.type_id
1312 WHERE r.from_id = ?1 AND r.to_id = ?2 AND r.type_id = ?3
1313 AND e1.flags = 0 AND e2.flags = 0
1314 ORDER BY r.from_id, r.to_id",
1315 ) && let Ok(rows) = stmt.query_map(params![fid, tid, tpid], |row| {
1316 Ok(Relation {
1317 from: row.get(0)?,
1318 to: row.get(1)?,
1319 relation_type: row.get(2)?,
1320 })
1321 }) {
1322 for row in rows.flatten() {
1323 triples.push(row);
1324 }
1325 }
1326 }
1327 (Some(fid), Some(tid), None) => {
1328 if let Ok(mut stmt) = conn.prepare_cached(
1329 "SELECT e1.name, e2.name, t.name
1330 FROM relation r
1331 JOIN entity e1 ON e1.id = r.from_id
1332 JOIN entity e2 ON e2.id = r.to_id
1333 JOIN type_dict t ON t.id = r.type_id
1334 WHERE r.from_id = ?1 AND r.to_id = ?2
1335 AND e1.flags = 0 AND e2.flags = 0
1336 ORDER BY r.from_id, r.to_id",
1337 ) && let Ok(rows) = stmt.query_map(params![fid, tid], |row| {
1338 Ok(Relation {
1339 from: row.get(0)?,
1340 to: row.get(1)?,
1341 relation_type: row.get(2)?,
1342 })
1343 }) {
1344 for row in rows.flatten() {
1345 triples.push(row);
1346 }
1347 }
1348 }
1349 (Some(fid), None, Some(tpid)) => {
1350 if let Ok(mut stmt) = conn.prepare_cached(
1351 "SELECT e1.name, e2.name, t.name
1352 FROM relation r
1353 JOIN entity e1 ON e1.id = r.from_id
1354 JOIN entity e2 ON e2.id = r.to_id
1355 JOIN type_dict t ON t.id = r.type_id
1356 WHERE r.from_id = ?1 AND r.type_id = ?2
1357 AND e1.flags = 0 AND e2.flags = 0
1358 ORDER BY r.from_id, r.to_id",
1359 ) && let Ok(rows) = stmt.query_map(params![fid, tpid], |row| {
1360 Ok(Relation {
1361 from: row.get(0)?,
1362 to: row.get(1)?,
1363 relation_type: row.get(2)?,
1364 })
1365 }) {
1366 for row in rows.flatten() {
1367 triples.push(row);
1368 }
1369 }
1370 }
1371 (None, Some(tid), Some(tpid)) => {
1372 if let Ok(mut stmt) = conn.prepare_cached(
1373 "SELECT e1.name, e2.name, t.name
1374 FROM relation r
1375 JOIN entity e1 ON e1.id = r.from_id
1376 JOIN entity e2 ON e2.id = r.to_id
1377 JOIN type_dict t ON t.id = r.type_id
1378 WHERE r.to_id = ?1 AND r.type_id = ?2
1379 AND e1.flags = 0 AND e2.flags = 0
1380 ORDER BY r.from_id, r.to_id",
1381 ) && let Ok(rows) = stmt.query_map(params![tid, tpid], |row| {
1382 Ok(Relation {
1383 from: row.get(0)?,
1384 to: row.get(1)?,
1385 relation_type: row.get(2)?,
1386 })
1387 }) {
1388 for row in rows.flatten() {
1389 triples.push(row);
1390 }
1391 }
1392 }
1393 (Some(fid), None, None) => {
1394 if let Ok(mut stmt) = conn.prepare_cached(
1395 "SELECT e1.name, e2.name, t.name
1396 FROM relation r
1397 JOIN entity e1 ON e1.id = r.from_id
1398 JOIN entity e2 ON e2.id = r.to_id
1399 JOIN type_dict t ON t.id = r.type_id
1400 WHERE r.from_id = ?1
1401 AND e1.flags = 0 AND e2.flags = 0
1402 ORDER BY r.from_id, r.to_id",
1403 ) && let Ok(rows) = stmt.query_map(params![fid], |row| {
1404 Ok(Relation {
1405 from: row.get(0)?,
1406 to: row.get(1)?,
1407 relation_type: row.get(2)?,
1408 })
1409 }) {
1410 for row in rows.flatten() {
1411 triples.push(row);
1412 }
1413 }
1414 }
1415 (None, Some(tid), None) => {
1416 if let Ok(mut stmt) = conn.prepare_cached(
1417 "SELECT e1.name, e2.name, t.name
1418 FROM relation r
1419 JOIN entity e1 ON e1.id = r.from_id
1420 JOIN entity e2 ON e2.id = r.to_id
1421 JOIN type_dict t ON t.id = r.type_id
1422 WHERE r.to_id = ?1
1423 AND e1.flags = 0 AND e2.flags = 0
1424 ORDER BY r.from_id, r.to_id",
1425 ) && let Ok(rows) = stmt.query_map(params![tid], |row| {
1426 Ok(Relation {
1427 from: row.get(0)?,
1428 to: row.get(1)?,
1429 relation_type: row.get(2)?,
1430 })
1431 }) {
1432 for row in rows.flatten() {
1433 triples.push(row);
1434 }
1435 }
1436 }
1437 (None, None, Some(tpid)) => {
1438 if let Ok(mut stmt) = conn.prepare_cached(
1439 "SELECT e1.name, e2.name, t.name
1440 FROM relation r
1441 JOIN entity e1 ON e1.id = r.from_id
1442 JOIN entity e2 ON e2.id = r.to_id
1443 JOIN type_dict t ON t.id = r.type_id
1444 WHERE r.type_id = ?1
1445 AND e1.flags = 0 AND e2.flags = 0
1446 ORDER BY r.from_id, r.to_id",
1447 ) && let Ok(rows) = stmt.query_map(params![tpid], |row| {
1448 Ok(Relation {
1449 from: row.get(0)?,
1450 to: row.get(1)?,
1451 relation_type: row.get(2)?,
1452 })
1453 }) {
1454 for row in rows.flatten() {
1455 triples.push(row);
1456 }
1457 }
1458 }
1459 (None, None, None) => {
1460 if let Ok(mut stmt) = conn.prepare_cached(
1461 "SELECT e1.name, e2.name, t.name
1462 FROM relation r
1463 JOIN entity e1 ON e1.id = r.from_id
1464 JOIN entity e2 ON e2.id = r.to_id
1465 JOIN type_dict t ON t.id = r.type_id
1466 WHERE e1.flags = 0 AND e2.flags = 0
1467 ORDER BY r.from_id, r.to_id",
1468 ) && let Ok(rows) = stmt.query_map([], |row| {
1469 Ok(Relation {
1470 from: row.get(0)?,
1471 to: row.get(1)?,
1472 relation_type: row.get(2)?,
1473 })
1474 }) {
1475 for row in rows.flatten() {
1476 triples.push(row);
1477 }
1478 }
1479 }
1480 }
1481 if let Some(lim) = limit {
1482 triples.truncate(lim);
1483 }
1484 relation_details(&conn, triples.into_iter())
1485 }
1486
1487 pub fn find_path(&self, from: &str, to: &str) -> Result<Option<Vec<String>>> {
1488 let conn = self.readers.get();
1489 let (from_id, _, _, _) = match self.get_entity_id(&conn, from)? {
1490 Some(v) => v,
1491 None => {
1492 return Err(MCSError::InvalidParams(format!(
1493 "Source entity '{from}' not found"
1494 )));
1495 }
1496 };
1497 let (to_id, _, _, _) = match self.get_entity_id(&conn, to)? {
1498 Some(v) => v,
1499 None => {
1500 return Err(MCSError::InvalidParams(format!(
1501 "Target entity '{to}' not found"
1502 )));
1503 }
1504 };
1505
1506 if from_id == to_id {
1507 return Ok(Some(vec![from.to_string()]));
1508 }
1509
1510 let mut visited = HashSet::new();
1512 let mut parent: FxHashMap<i64, i64> = FxHashMap::default();
1513 let mut queue = VecDeque::new();
1514 visited.insert(from_id);
1515 queue.push_back(from_id);
1516
1517 while let Some(cur) = queue.pop_front() {
1518 if cur == to_id {
1519 break;
1520 }
1521 if let Ok(mut stmt) =
1523 conn.prepare_cached("SELECT to_id FROM relation WHERE from_id = ?1")
1524 && let Ok(rows) = stmt.query_map(params![cur], |row| row.get::<_, i64>(0))
1525 {
1526 for row in rows.flatten() {
1527 if visited.insert(row) {
1528 parent.insert(row, cur);
1529 queue.push_back(row);
1530 }
1531 }
1532 }
1533 if let Ok(mut stmt) =
1535 conn.prepare_cached("SELECT from_id FROM relation WHERE to_id = ?1")
1536 && let Ok(rows) = stmt.query_map(params![cur], |row| row.get::<_, i64>(0))
1537 {
1538 for row in rows.flatten() {
1539 if visited.insert(row) {
1540 parent.insert(row, cur);
1541 queue.push_back(row);
1542 }
1543 }
1544 }
1545 }
1546
1547 if !parent.contains_key(&to_id) && to_id != from_id {
1548 return Ok(None);
1549 }
1550
1551 let mut path = Vec::new();
1552 let mut cur = to_id;
1553 path.push(cur);
1554 while let Some(&p) = parent.get(&cur) {
1555 path.push(p);
1556 cur = p;
1557 if cur == from_id {
1558 break;
1559 }
1560 }
1561 path.reverse();
1562
1563 let placeholders: Vec<String> = path.iter().map(|_| "?".to_string()).collect();
1564 let sql = format!(
1565 "SELECT id, name FROM entity WHERE id IN ({})",
1566 placeholders.join(",")
1567 );
1568 let name_map: FxHashMap<i64, String> = if let Ok(mut stmt) = conn.prepare(&sql)
1569 && let Ok(rows) = stmt.query_map(rusqlite::params_from_iter(&path), |row| {
1570 Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?))
1571 }) {
1572 rows.flatten().collect()
1573 } else {
1574 FxHashMap::default()
1575 };
1576
1577 let name_path: Vec<String> = path
1578 .iter()
1579 .filter_map(|id| name_map.get(id).cloned())
1580 .collect();
1581
1582 Ok(Some(name_path))
1583 }
1584
1585 pub fn compact(&self) -> Result<()> {
1586 self.mutate(MutationRequest::Compact).map(|_| ())
1587 }
1588
1589 pub fn neighbors(
1590 &self,
1591 name: &str,
1592 direction: Direction,
1593 rtype: Option<&str>,
1594 depth: u32,
1595 ) -> Result<String> {
1596 self._traverse(name, direction, rtype, depth, true)
1597 }
1598
1599 pub fn extract_subgraph(&self, names: &[String], depth: u32) -> Result<String> {
1600 if names.is_empty() {
1601 return Ok(r#"{"entities":[],"relations":[]}"#.to_string());
1602 }
1603
1604 let conn = self.readers.get();
1605 let mut all_entity_ids: HashSet<i64> = HashSet::new();
1606 let mut frontier: HashSet<i64> = HashSet::new();
1607 let mut all_rel_pairs: HashSet<(i64, i64, i64)> = HashSet::new();
1608
1609 for name in names {
1611 let h = name_hash(name);
1612 if let Ok(Some(id)) = conn
1613 .query_row(
1614 "SELECT id FROM entity WHERE name_hash = ?1 AND name = ?2 AND flags = 0",
1615 params![h, name],
1616 |row| row.get::<_, i64>(0),
1617 )
1618 .map(Some)
1619 .or_else(|e| {
1620 if is_not_found(&e) {
1621 Ok(None)
1622 } else {
1623 Err(sqlite_err(e))
1624 }
1625 })
1626 {
1627 all_entity_ids.insert(id);
1628 frontier.insert(id);
1629 }
1630 }
1631
1632 let mut current_depth = 0u32;
1633 while current_depth < depth && !frontier.is_empty() {
1634 let mut next_frontier: HashSet<i64> = HashSet::new();
1635
1636 const CHUNK: usize = 500;
1639 let frontier_ids: Vec<i64> = frontier.iter().copied().collect();
1640 for chunk in frontier_ids.chunks(CHUNK) {
1641 let placeholders: Vec<String> = chunk.iter().map(|_| "?".to_string()).collect();
1642 let in_clause = placeholders.join(",");
1643
1644 if let Ok(mut stmt) = conn.prepare(&format!(
1646 "SELECT from_id, to_id, type_id FROM relation WHERE from_id IN ({in_clause})",
1647 )) && let Ok(rows) = stmt.query_map(rusqlite::params_from_iter(chunk), |row| {
1648 Ok((
1649 row.get::<_, i64>(0)?,
1650 row.get::<_, i64>(1)?,
1651 row.get::<_, i64>(2)?,
1652 ))
1653 }) {
1654 for row in rows.flatten() {
1655 let (from_id, to_id, type_id) = row;
1656 all_rel_pairs.insert((from_id, to_id, type_id));
1657 if all_entity_ids.insert(to_id) {
1658 next_frontier.insert(to_id);
1659 }
1660 }
1661 }
1662
1663 if let Ok(mut stmt) = conn.prepare(&format!(
1665 "SELECT from_id, to_id, type_id FROM relation WHERE to_id IN ({in_clause})",
1666 )) && let Ok(rows) = stmt.query_map(rusqlite::params_from_iter(chunk), |row| {
1667 Ok((
1668 row.get::<_, i64>(0)?,
1669 row.get::<_, i64>(1)?,
1670 row.get::<_, i64>(2)?,
1671 ))
1672 }) {
1673 for row in rows.flatten() {
1674 let (from_id, to_id, type_id) = row;
1675 all_rel_pairs.insert((from_id, to_id, type_id));
1676 if all_entity_ids.insert(from_id) {
1677 next_frontier.insert(from_id);
1678 }
1679 }
1680 }
1681 }
1682 if all_entity_ids.len() > MAX_TRAVERSAL_ENTITIES
1683 || all_rel_pairs.len() > MAX_TRAVERSAL_RELS
1684 {
1685 break;
1686 }
1687 frontier = next_frontier;
1688 current_depth += 1;
1689 }
1690
1691 let entities_json: String = if all_entity_ids.is_empty() {
1692 "[]".to_string()
1693 } else {
1694 let ids: Vec<i64> = all_entity_ids.iter().copied().collect();
1695 let sql = format!(
1696 "SELECT COALESCE(json_group_array(json_object(
1697 'name', e.name,
1698 'entityType', t.name,
1699 'observations', COALESCE((
1700 SELECT json_group_array({OBSERVATION_JSON} ORDER BY o.idx, o.id)
1701 FROM observation o WHERE o.entity_id = e.id
1702 ), json('[]'))
1703 ) ORDER BY e.id), json('[]'))
1704 FROM entity e
1705 JOIN type_dict t ON t.id = e.type_id
1706 WHERE e.id IN ({}) AND e.flags = 0",
1707 int_csv(&ids)
1708 );
1709 conn.query_row(&sql, [], |row| row.get::<_, String>(0))
1710 .map_err(sqlite_err)?
1711 };
1712
1713 let relations_json: String = if all_rel_pairs.is_empty() {
1714 "[]".to_string()
1715 } else {
1716 let sql = format!(
1717 "WITH r(from_id, to_id, type_id) AS (VALUES {})
1718 SELECT COALESCE(json_group_array(json_object(
1719 'from', e1.name,
1720 'to', e2.name,
1721 'relationType', t.name
1722 )), json('[]'))
1723 FROM r
1724 JOIN entity e1 ON e1.id = r.from_id
1725 JOIN entity e2 ON e2.id = r.to_id
1726 JOIN type_dict t ON t.id = r.type_id
1727 WHERE e1.flags = 0 AND e2.flags = 0",
1728 rel_values_literal(&all_rel_pairs)
1729 );
1730 conn.query_row(&sql, [], |row| row.get::<_, String>(0))
1731 .map_err(sqlite_err)?
1732 };
1733
1734 let mut out = String::with_capacity(32 + entities_json.len() + relations_json.len());
1735 out.push_str("{\"entities\":");
1736 out.push_str(&entities_json);
1737 out.push_str(",\"relations\":");
1738 out.push_str(&relations_json);
1739 out.push('}');
1740 Ok(out)
1741 }
1742
1743 pub fn describe_entity(&self, name: &str) -> Result<EntityDescription> {
1744 let conn = self.readers.get();
1745 let tx = conn.unchecked_transaction().map_err(sqlite_err)?;
1749 let entity = crate::mutation::read_entity(&tx, name)?
1750 .ok_or_else(|| MCSError::InvalidParams(format!("Entity '{name}' not found")))?;
1751 let relations = crate::mutation::relations_for(&tx, name)?;
1752 let mut neighbors: Vec<String> = relations
1753 .iter()
1754 .map(|relation| {
1755 if relation.from == name {
1756 relation.to.clone()
1757 } else {
1758 relation.from.clone()
1759 }
1760 })
1761 .collect();
1762 neighbors.sort();
1763 neighbors.dedup();
1764 let incoming = relations
1768 .iter()
1769 .filter(|relation| relation.to == name)
1770 .count() as i64;
1771 let outgoing = relations
1772 .iter()
1773 .filter(|relation| relation.from == name)
1774 .count() as i64;
1775 let attributes = some_nonempty(attributes_for(&tx, "entity", entity.entity_id)?);
1776 tx.commit().map_err(sqlite_err)?;
1777
1778 Ok(EntityDescription {
1779 name: entity.name,
1780 entity_type: entity.entity_type,
1781 observations: entity.observations,
1782 relations,
1783 neighbors,
1784 degree: Degree { incoming, outgoing },
1785 attributes,
1786 })
1787 }
1788
1789 pub fn entity_type_counts(&self) -> Vec<(String, usize)> {
1790 let conn = self.readers.get();
1791 select_all_types(&conn, 0).unwrap_or_default()
1792 }
1793
1794 pub fn entity_type_catalog(&self) -> Vec<(String, usize, Option<String>)> {
1797 let conn = self.readers.get();
1798 select_type_catalog(&conn, 0).unwrap_or_default()
1799 }
1800
1801 pub fn ui_meta(&self) -> (Vec<(String, usize)>, usize, usize) {
1809 let conn = self.readers.get();
1810 let types = select_all_types(&conn, 0).unwrap_or_default();
1811 let entities = read_graph_stat(&conn, "entities").unwrap_or(0).max(0) as usize;
1812 let relations = read_graph_stat(&conn, "relations").unwrap_or(0).max(0) as usize;
1813 (types, entities, relations)
1814 }
1815
1816 pub fn relation_type_counts(&self) -> Vec<(String, usize)> {
1817 let conn = self.readers.get();
1818 select_all_types(&conn, 1).unwrap_or_default()
1819 }
1820
1821 pub fn relation_type_catalog(&self) -> Vec<(String, usize, Option<String>)> {
1824 let conn = self.readers.get();
1825 select_type_catalog(&conn, 1).unwrap_or_default()
1826 }
1827
1828 pub fn set_type_description(
1835 &self,
1836 kind: i64,
1837 name: &str,
1838 description: Option<&str>,
1839 ) -> Result<()> {
1840 let conn = self.writer.lock();
1841 let existing = lookup_type_id(&conn, name, kind);
1844 let description: Option<String> = description.map(|d| d.into());
1845 match existing {
1846 Some(id) => {
1847 conn.execute(
1848 "UPDATE type_dict SET desc = ?1 WHERE id = ?2",
1849 params![description, id],
1850 )
1851 .map_err(sqlite_err)?;
1852 }
1853 None => {
1854 conn.execute(
1855 "INSERT INTO type_dict(kind, name, count, desc) VALUES(?1, ?2, 0, ?3)",
1856 params![kind, name, description],
1857 )
1858 .map_err(sqlite_err)?;
1859 }
1860 }
1861 Ok(())
1862 }
1863
1864 pub fn entity_type_exists(&self, name: &str) -> bool {
1867 let conn = self.readers.get();
1868 lookup_type_id(&conn, name, 0).is_some()
1869 }
1870
1871 pub fn relation_type_exists(&self, name: &str) -> bool {
1874 let conn = self.readers.get();
1875 lookup_type_id(&conn, name, 1).is_some()
1876 }
1877
1878 pub fn batch_get_entities(&self, names: &[String]) -> Vec<Option<Entity>> {
1879 let conn = self.readers.get();
1880 let tx = match conn.unchecked_transaction().map_err(sqlite_err) {
1887 Ok(tx) => tx,
1888 Err(_) => return names.iter().map(|_| None).collect(),
1889 };
1890 let (_, by_name): (Vec<Option<i64>>, FxHashMap<String, Entity>) = {
1891 let ids: Vec<Option<i64>> = names
1892 .iter()
1893 .map(|n| entity_name_lookup(&tx, n).ok().flatten())
1894 .collect();
1895 let resolved: Vec<i64> = ids.iter().flatten().copied().collect();
1896 let mut by_id = batch_entities_by_ids(&tx, &resolved);
1897 let mut attrs: FxHashMap<i64, BTreeMap<String, String>> = FxHashMap::default();
1898 if !resolved.is_empty() {
1899 let sql = format!(
1900 "SELECT owner_id, key, value FROM attribute
1901 WHERE owner_kind = 'entity' AND owner_id IN ({})
1902 ORDER BY owner_id, key",
1903 int_csv(&resolved)
1904 );
1905 if let Ok(mut stmt) = tx.prepare(&sql)
1906 && let Ok(rows) = stmt.query_map([], |row| {
1907 Ok((
1908 row.get::<_, i64>(0)?,
1909 row.get::<_, String>(1)?,
1910 row.get::<_, String>(2)?,
1911 ))
1912 })
1913 {
1914 for row in rows.flatten() {
1915 attrs.entry(row.0).or_default().insert(row.1, row.2);
1916 }
1917 }
1918 }
1919 let mut by_name: FxHashMap<String, Entity> = FxHashMap::default();
1923 for id in &resolved {
1924 if let Some(mut entity) = by_id.remove(id) {
1925 if let Some(map) = attrs.remove(id) {
1926 entity.attributes = some_nonempty(map);
1927 }
1928 by_name.insert(entity.name.clone(), entity);
1929 }
1930 }
1931 (ids, by_name)
1932 };
1933 let _ = tx.commit().map_err(sqlite_err);
1934 names.iter().map(|n| by_name.get(n).cloned()).collect()
1935 }
1936
1937 pub fn find_all_paths(
1938 &self,
1939 from: &str,
1940 to: &str,
1941 max_depth: usize,
1942 max_paths: usize,
1943 ) -> Result<Vec<Vec<String>>> {
1944 let conn = self.readers.get();
1945 let (from_id, _, _, _) = match self.get_entity_id(&conn, from)? {
1946 Some(v) => v,
1947 None => {
1948 return Err(MCSError::InvalidParams(format!(
1949 "Source entity '{from}' not found"
1950 )));
1951 }
1952 };
1953 let (to_id, _, _, _) = match self.get_entity_id(&conn, to)? {
1954 Some(v) => v,
1955 None => {
1956 return Err(MCSError::InvalidParams(format!(
1957 "Target entity '{to}' not found"
1958 )));
1959 }
1960 };
1961
1962 if from_id == to_id {
1963 return Ok(vec![vec![from.to_string()]]);
1964 }
1965
1966 let mut all_paths: Vec<Vec<i64>> = Vec::new();
1968 let mut queue: VecDeque<(i64, Vec<i64>)> = VecDeque::new();
1969 queue.push_back((from_id, vec![from_id]));
1970
1971 const MAX_QUEUE_SIZE: usize = 10_000_000;
1972
1973 while let Some((cur, path)) = queue.pop_front() {
1974 if all_paths.len() >= max_paths {
1975 break;
1976 }
1977 if path.len() > max_depth {
1978 continue;
1979 }
1980
1981 if let Ok(mut stmt) =
1983 conn.prepare_cached("SELECT to_id FROM relation WHERE from_id = ?1")
1984 && let Ok(rows) = stmt.query_map(params![cur], |row| row.get::<_, i64>(0))
1985 {
1986 for next_id in rows.flatten() {
1987 if next_id == to_id {
1988 let mut full_path = path.clone();
1989 full_path.push(next_id);
1990 all_paths.push(full_path);
1991 if all_paths.len() >= max_paths {
1992 break;
1993 }
1994 } else if !path.contains(&next_id) && path.len() < max_depth {
1995 if queue.len() >= MAX_QUEUE_SIZE {
1996 return Err(MCSError::InvalidParams(
1997 "Path exploration queue exceeded limit (too many paths on highly connected graph)".to_string()
1998 ));
1999 }
2000 let mut new_path = path.clone();
2001 new_path.push(next_id);
2002 queue.push_back((next_id, new_path));
2003 }
2004 }
2005 }
2006
2007 if let Ok(mut stmt) =
2009 conn.prepare_cached("SELECT from_id FROM relation WHERE to_id = ?1")
2010 && let Ok(rows) = stmt.query_map(params![cur], |row| row.get::<_, i64>(0))
2011 {
2012 for next_id in rows.flatten() {
2013 if next_id == to_id {
2014 let mut full_path = path.clone();
2015 full_path.push(next_id);
2016 all_paths.push(full_path);
2017 if all_paths.len() >= max_paths {
2018 break;
2019 }
2020 } else if !path.contains(&next_id) && path.len() < max_depth {
2021 if queue.len() >= MAX_QUEUE_SIZE {
2022 return Err(MCSError::InvalidParams(
2023 "Path exploration queue exceeded limit (too many paths on highly connected graph)".to_string()
2024 ));
2025 }
2026 let mut new_path = path.clone();
2027 new_path.push(next_id);
2028 queue.push_back((next_id, new_path));
2029 }
2030 }
2031 }
2032 }
2033
2034 let all_ids: HashSet<i64> = all_paths.iter().flat_map(|p| p.iter()).copied().collect();
2036 let id_list: Vec<i64> = all_ids.into_iter().collect();
2037 let name_map: FxHashMap<i64, String> = if id_list.is_empty() {
2038 FxHashMap::default()
2039 } else {
2040 let placeholders: Vec<String> = id_list.iter().map(|_| "?".to_string()).collect();
2041 let sql = format!(
2042 "SELECT id, name FROM entity WHERE id IN ({})",
2043 placeholders.join(",")
2044 );
2045 if let Ok(mut stmt) = conn.prepare(&sql)
2046 && let Ok(rows) = stmt.query_map(rusqlite::params_from_iter(&id_list), |row| {
2047 Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?))
2048 })
2049 {
2050 rows.flatten().collect()
2051 } else {
2052 FxHashMap::default()
2053 }
2054 };
2055
2056 let mut named_paths: Vec<Vec<String>> = Vec::with_capacity(all_paths.len());
2057 for path_ids in all_paths {
2058 let named: Vec<String> = path_ids
2059 .iter()
2060 .filter_map(|id| name_map.get(id).cloned())
2061 .collect();
2062 named_paths.push(named);
2063 }
2064
2065 Ok(named_paths)
2066 }
2067
2068 pub fn export(&self, _format: &str, max_rows: i64) -> Result<String> {
2073 let conn = self.readers.get();
2074 conn.query_row(
2077 &format!(
2078 "SELECT json_object(
2079 'entities', COALESCE((
2080 SELECT json_group_array(json_object(
2081 'name', e.name,
2082 'entityType', t.name,
2083 'observations', COALESCE((
2084 SELECT json_group_array({OBSERVATION_JSON} ORDER BY o.idx, o.id)
2085 FROM observation o WHERE o.entity_id = e.id
2086 ), json('[]')),
2087 'attributes', COALESCE((
2088 SELECT json_group_object(key, value) FROM attribute
2089 WHERE owner_kind = 'entity' AND owner_id = e.id
2090 ), json('{{}}'))
2091 ) ORDER BY e.id)
2092 FROM (
2093 SELECT id, name, type_id FROM entity
2094 WHERE flags = 0 ORDER BY id LIMIT ?1
2095 ) e
2096 JOIN type_dict t ON t.id = e.type_id
2097 ), json('[]')),
2098 'relations', COALESCE((
2099 SELECT json_group_array(json_object(
2100 'from', e1.name,
2101 'to', e2.name,
2102 'relationType', t.name,
2103 'observations', COALESCE((
2104 SELECT json_group_array(json_object(
2105 'body', ro.body,
2106 'createdAtUs', ro.created_us,
2107 'occurredAtUs', ro.occurred_us,
2108 'originEntityName', NULL))
2109 FROM relation_observation ro
2110 JOIN taxonomy_relation m ON m.id = ro.relation_id AND m.deleted = 0
2111 WHERE m.from_id = r.from_id AND m.to_id = r.to_id AND m.type_id = r.type_id
2112 ORDER BY ro.idx
2113 ), json('[]')),
2114 'attributes', COALESCE((
2115 SELECT json_group_object(key, value) FROM attribute
2116 WHERE owner_kind = 'relation' AND owner_id = (
2117 SELECT id FROM taxonomy_relation
2118 WHERE from_id = r.from_id AND to_id = r.to_id
2119 AND type_id = r.type_id AND deleted = 0)
2120 ), json('{{}}'))
2121 ))
2122 FROM (
2123 SELECT from_id, to_id, type_id FROM relation LIMIT ?1
2124 ) r
2125 JOIN entity e1 ON e1.id = r.from_id
2126 JOIN entity e2 ON e2.id = r.to_id
2127 JOIN type_dict t ON t.id = r.type_id
2128 WHERE e1.flags = 0 AND e2.flags = 0
2129 ), json('[]'))
2130 )"
2131 ),
2132 params![max_rows],
2133 |row| row.get::<_, String>(0),
2134 )
2135 .map_err(sqlite_err)
2136 }
2137
2138 pub fn wipe(&self) -> Result<()> {
2139 self.mutate(MutationRequest::Wipe).map(|_| ())
2140 }
2141
2142 pub fn run_maintenance(&self) -> Result<()> {
2145 let conn = self.writer.lock();
2146
2147 conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")
2148 .map_err(sqlite_err)?;
2149
2150 conn.execute_batch("PRAGMA optimize(0x10000);")
2151 .map_err(sqlite_err)?;
2152
2153 let tx = TxGuard::begin(&conn)?;
2154 conn.execute_batch(
2155 "INSERT INTO name_fts(name_fts) VALUES('optimize');
2156 INSERT INTO obs_fts(obs_fts) VALUES('optimize');
2157 INSERT INTO rel_obs_fts(rel_obs_fts) VALUES('optimize');",
2158 )
2159 .map_err(sqlite_err)?;
2160 tx.commit()?;
2161
2162 Ok(())
2163 }
2164
2165 pub fn checkpoint_passive(&self) -> Result<()> {
2169 let conn = self.writer.lock();
2170 conn.execute_batch("PRAGMA wal_checkpoint(PASSIVE);")
2171 .map_err(sqlite_err)?;
2172 Ok(())
2173 }
2174
2175 fn _traverse(
2176 &self,
2177 name: &str,
2178 direction: Direction,
2179 rtype: Option<&str>,
2180 depth: u32,
2181 _include_relations: bool,
2183 ) -> Result<String> {
2184 let conn = self.readers.get();
2185 let (start_id, _, _, _) = match self.get_entity_id(&conn, name)? {
2186 Some(v) => v,
2187 None => {
2188 return Err(MCSError::InvalidParams(format!(
2189 "Entity '{name}' not found"
2190 )));
2191 }
2192 };
2193
2194 let mut all_ids: HashSet<i64> = HashSet::new();
2195 let mut all_rels: HashSet<(i64, i64, i64)> = HashSet::new();
2196 let mut frontier: HashSet<i64> = HashSet::new();
2197 all_ids.insert(start_id);
2198 frontier.insert(start_id);
2199
2200 let type_filter: Option<i64> = rtype
2206 .filter(|rt| !rt.is_empty())
2207 .map(|rt| lookup_type_id(&conn, rt, 1).unwrap_or(-1));
2208
2209 let mut q_out_t = conn.prepare_cached(
2211 "SELECT to_id, type_id FROM relation WHERE from_id = ?1 AND type_id = ?2",
2212 );
2213 let mut q_out =
2214 conn.prepare_cached("SELECT to_id, type_id FROM relation WHERE from_id = ?1");
2215 let mut q_in_t = conn.prepare_cached(
2216 "SELECT from_id, type_id FROM relation WHERE to_id = ?1 AND type_id = ?2",
2217 );
2218 let mut q_in =
2219 conn.prepare_cached("SELECT from_id, type_id FROM relation WHERE to_id = ?1");
2220
2221 let mut cur_depth = 0u32;
2222 while cur_depth < depth && !frontier.is_empty() {
2223 let mut next_frontier: HashSet<i64> = HashSet::new();
2224
2225 for &fid in &frontier {
2226 if direction == Direction::Outgoing || direction == Direction::Both {
2227 if let Some(tid) = type_filter {
2228 if let Ok(ref mut stmt) = q_out_t
2229 && let Ok(rows) = stmt.query_map(params![fid, tid], |row| {
2230 Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
2231 })
2232 {
2233 for row in rows.flatten() {
2234 let (to_id, t_id) = row;
2235 all_rels.insert((fid, to_id, t_id));
2236 if all_ids.insert(to_id) {
2237 next_frontier.insert(to_id);
2238 }
2239 }
2240 }
2241 } else if let Ok(ref mut stmt) = q_out
2242 && let Ok(rows) = stmt.query_map(params![fid], |row| {
2243 Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
2244 })
2245 {
2246 for row in rows.flatten() {
2247 let (to_id, t_id) = row;
2248 all_rels.insert((fid, to_id, t_id));
2249 if all_ids.insert(to_id) {
2250 next_frontier.insert(to_id);
2251 }
2252 }
2253 }
2254 }
2255
2256 if direction == Direction::Incoming || direction == Direction::Both {
2257 if let Some(tid) = type_filter {
2258 if let Ok(ref mut stmt) = q_in_t
2259 && let Ok(rows) = stmt.query_map(params![fid, tid], |row| {
2260 Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
2261 })
2262 {
2263 for row in rows.flatten() {
2264 let (from_id, t_id) = row;
2265 all_rels.insert((from_id, fid, t_id));
2266 if all_ids.insert(from_id) {
2267 next_frontier.insert(from_id);
2268 }
2269 }
2270 }
2271 } else if let Ok(ref mut stmt) = q_in
2272 && let Ok(rows) = stmt.query_map(params![fid], |row| {
2273 Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
2274 })
2275 {
2276 for row in rows.flatten() {
2277 let (from_id, t_id) = row;
2278 all_rels.insert((from_id, fid, t_id));
2279 if all_ids.insert(from_id) {
2280 next_frontier.insert(from_id);
2281 }
2282 }
2283 }
2284 }
2285 }
2286
2287 if all_ids.len() > MAX_TRAVERSAL_ENTITIES || all_rels.len() > MAX_TRAVERSAL_RELS {
2291 break;
2292 }
2293
2294 frontier = next_frontier;
2295 cur_depth += 1;
2296 }
2297
2298 let entities_json: String = if all_ids.is_empty() {
2299 "[]".to_string()
2300 } else {
2301 let ids: Vec<i64> = all_ids.iter().copied().collect();
2302 let sql = format!(
2303 "SELECT COALESCE(json_group_array(json_object(
2304 'name', e.name,
2305 'entityType', t.name,
2306 'observations', COALESCE((
2307 SELECT json_group_array({OBSERVATION_JSON} ORDER BY o.idx, o.id)
2308 FROM observation o WHERE o.entity_id = e.id
2309 ), json('[]'))
2310 ) ORDER BY e.id), json('[]'))
2311 FROM entity e
2312 JOIN type_dict t ON t.id = e.type_id
2313 WHERE e.id IN ({}) AND e.flags = 0",
2314 int_csv(&ids)
2315 );
2316 conn.query_row(&sql, [], |row| row.get::<_, String>(0))
2317 .map_err(sqlite_err)?
2318 };
2319
2320 let relations_json: String = if all_rels.is_empty() {
2321 "[]".to_string()
2322 } else {
2323 let sql = format!(
2324 "WITH r(from_id, to_id, type_id) AS (VALUES {})
2325 SELECT COALESCE(json_group_array(json_object(
2326 'from', e1.name,
2327 'to', e2.name,
2328 'relationType', t.name
2329 )), json('[]'))
2330 FROM r
2331 JOIN entity e1 ON e1.id = r.from_id
2332 JOIN entity e2 ON e2.id = r.to_id
2333 JOIN type_dict t ON t.id = r.type_id
2334 WHERE e1.flags = 0 AND e2.flags = 0",
2335 rel_values_literal(&all_rels)
2336 );
2337 conn.query_row(&sql, [], |row| row.get::<_, String>(0))
2338 .map_err(sqlite_err)?
2339 };
2340
2341 let mut out = String::with_capacity(32 + entities_json.len() + relations_json.len());
2342 out.push_str("{\"entities\":");
2343 out.push_str(&entities_json);
2344 out.push_str(",\"relations\":");
2345 out.push_str(&relations_json);
2346 out.push('}');
2347 Ok(out)
2348 }
2349}
2350
2351#[cfg(test)]
2354mod tests {
2355 use super::*;
2356 use crate::types::EntityInput as Entity;
2357 use serde_json::Value;
2358 use std::ops::Deref;
2359 use std::path::PathBuf;
2360
2361 struct TestKg(GraphHandle, PathBuf);
2362
2363 impl Deref for TestKg {
2364 type Target = GraphHandle;
2365 fn deref(&self) -> &GraphHandle {
2366 &self.0
2367 }
2368 }
2369
2370 impl Drop for TestKg {
2371 fn drop(&mut self) {
2372 cleanup_db(&self.1);
2373 }
2374 }
2375
2376 fn cleanup_db(path: &std::path::Path) {
2377 let _ = std::fs::remove_file(path);
2378 let _ = std::fs::remove_file(path.with_extension("db-wal"));
2379 let _ = std::fs::remove_file(path.with_extension("db-shm"));
2380 }
2381
2382 fn new_kg() -> TestKg {
2383 use std::sync::atomic::AtomicU64;
2384 use std::sync::atomic::Ordering;
2385 static COUNTER: AtomicU64 = AtomicU64::new(0);
2386 let n = COUNTER.fetch_add(1, Ordering::SeqCst);
2387 let dir = std::env::temp_dir();
2388 let path = dir.join(format!("kg_test_{}_{}.db", std::process::id(), n));
2389 cleanup_db(&path);
2390 let kg = GraphHandle::new(
2391 &path,
2392 Durability::Async,
2393 SqliteTuning::default(),
2394 NonZeroUsize::new(10000).unwrap(),
2395 4,
2396 )
2397 .expect("create KG");
2398 TestKg(kg, path)
2399 }
2400
2401 #[test]
2402 fn test_create_and_get_entity() {
2403 let kg = new_kg();
2404 let entities = vec![Entity {
2405 name: "test".into(),
2406 entity_type: "person".into(),
2407 observations: vec!["obs1".into(), "obs2".into()],
2408 attributes: None,
2409 }];
2410 let created = kg.create_entities(&entities).unwrap();
2411 assert_eq!(created.len(), 1);
2412
2413 let got = kg.get_entity("test").unwrap().unwrap();
2414 assert_eq!(got.name, "test");
2415 assert_eq!(got.entity_type, "person");
2416 assert_eq!(
2417 got.observations
2418 .iter()
2419 .map(|o| o.body.as_str())
2420 .collect::<Vec<_>>(),
2421 vec!["obs1", "obs2"]
2422 );
2423 }
2424
2425 #[test]
2426 fn test_get_nonexistent() {
2427 let kg = new_kg();
2428 assert!(kg.get_entity("nonexistent").unwrap().is_none());
2429 }
2430
2431 #[test]
2432 fn test_delete_entity() {
2433 let kg = new_kg();
2434 kg.create_entities(&[Entity {
2435 name: "del".into(),
2436 entity_type: "t".into(),
2437 observations: vec![],
2438 attributes: None,
2439 }])
2440 .unwrap();
2441 assert!(kg.get_entity("del").unwrap().is_some());
2442 kg.delete_entities(&["del".to_string()]).unwrap();
2443 assert!(kg.get_entity("del").unwrap().is_none());
2444 }
2445
2446 #[test]
2447 fn test_add_and_delete_observations() {
2448 let kg = new_kg();
2449 kg.create_entities(&[Entity {
2450 name: "obs_test".into(),
2451 entity_type: "t".into(),
2452 observations: vec!["a".into()],
2453 attributes: None,
2454 }])
2455 .unwrap();
2456
2457 let added = kg
2458 .add_observations("obs_test", &["b".into(), "c".into()])
2459 .unwrap();
2460 assert_eq!(added.len(), 2);
2461
2462 let ent = kg.get_entity("obs_test").unwrap().unwrap();
2463 assert!(ent.observations.iter().any(|o| o.body == "b"));
2464 assert!(ent.observations.iter().any(|o| o.body == "c"));
2465
2466 kg.delete_observations("obs_test", &["b".into()]).unwrap();
2467 let ent = kg.get_entity("obs_test").unwrap().unwrap();
2468 assert!(!ent.observations.iter().any(|o| o.body == "b"));
2469 assert!(ent.observations.iter().any(|o| o.body == "c"));
2470 assert!(ent.observations.iter().any(|o| o.body == "a"));
2471 }
2472
2473 #[test]
2474 fn test_create_relations() {
2475 let kg = new_kg();
2476 kg.create_entities(&[
2477 Entity {
2478 name: "A".into(),
2479 entity_type: "node".into(),
2480 observations: vec![],
2481 attributes: None,
2482 },
2483 Entity {
2484 name: "B".into(),
2485 entity_type: "node".into(),
2486 observations: vec![],
2487 attributes: None,
2488 },
2489 ])
2490 .unwrap();
2491
2492 let rels = kg.create_relations(&[rel_input("A", "B", "edge")]).unwrap();
2493 assert_eq!(rels.len(), 1);
2494
2495 assert_eq!(kg.get_entity_count().unwrap(), 2);
2496 assert_eq!(kg.get_relation_count().unwrap(), 1);
2497 }
2498
2499 #[test]
2500 fn test_search_nodes() {
2501 let kg = new_kg();
2502 kg.create_entities(&[Entity {
2503 name: "Einstein".into(),
2504 entity_type: "scientist".into(),
2505 observations: vec!["physics".into(), "relativity".into()],
2506 attributes: None,
2507 }])
2508 .unwrap();
2509
2510 let results = kg.search_nodes_filtered("physics", None, 0, 10);
2511 assert_eq!(results.len(), 1);
2512 assert_eq!(results[0].name, "Einstein");
2513
2514 let results = kg.search_nodes_filtered("physics", Some("scientist"), 0, 10);
2515 assert_eq!(results.len(), 1);
2516
2517 let results = kg.search_nodes_filtered("physics", Some("nonexistent"), 0, 10);
2518 assert_eq!(results.len(), 0);
2519 }
2520
2521 #[test]
2522 fn test_find_path() {
2523 let kg = new_kg();
2524 kg.create_entities(&[
2525 Entity {
2526 name: "A".into(),
2527 entity_type: "n".into(),
2528 observations: vec![],
2529 attributes: None,
2530 },
2531 Entity {
2532 name: "B".into(),
2533 entity_type: "n".into(),
2534 observations: vec![],
2535 attributes: None,
2536 },
2537 Entity {
2538 name: "C".into(),
2539 entity_type: "n".into(),
2540 observations: vec![],
2541 attributes: None,
2542 },
2543 ])
2544 .unwrap();
2545
2546 kg.create_relations(&[rel_input("A", "B", "e"), rel_input("B", "C", "e")])
2547 .unwrap();
2548
2549 let path = kg.find_path("A", "C").unwrap().unwrap();
2550 assert_eq!(path, vec!["A", "B", "C"]);
2551 }
2552
2553 #[test]
2554 fn test_degree() {
2555 let kg = new_kg();
2556 kg.create_entities(&[
2557 Entity {
2558 name: "A".into(),
2559 entity_type: "n".into(),
2560 observations: vec![],
2561 attributes: None,
2562 },
2563 Entity {
2564 name: "B".into(),
2565 entity_type: "n".into(),
2566 observations: vec![],
2567 attributes: None,
2568 },
2569 Entity {
2570 name: "C".into(),
2571 entity_type: "n".into(),
2572 observations: vec![],
2573 attributes: None,
2574 },
2575 ])
2576 .unwrap();
2577
2578 kg.create_relations(&[rel_input("A", "B", "e"), rel_input("A", "C", "e")])
2579 .unwrap();
2580
2581 assert_eq!(kg.degree("A", Direction::Outgoing).unwrap(), 2);
2582 assert_eq!(kg.degree("A", Direction::Incoming).unwrap(), 0);
2583 assert_eq!(kg.degree("B", Direction::Incoming).unwrap(), 1);
2584 }
2585
2586 #[test]
2587 fn test_neighbors() {
2588 let kg = new_kg();
2589 kg.create_entities(&[
2590 Entity {
2591 name: "A".into(),
2592 entity_type: "n".into(),
2593 observations: vec![],
2594 attributes: None,
2595 },
2596 Entity {
2597 name: "B".into(),
2598 entity_type: "n".into(),
2599 observations: vec![],
2600 attributes: None,
2601 },
2602 ])
2603 .unwrap();
2604
2605 kg.create_relations(&[rel_input("A", "B", "e")]).unwrap();
2606
2607 let result = kg.neighbors("A", Direction::Outgoing, None, 1).unwrap();
2608 let v: Value = serde_json::from_str(&result).unwrap();
2609 assert_eq!(v["entities"].as_array().unwrap().len(), 2);
2610 assert_eq!(v["relations"].as_array().unwrap().len(), 1);
2611 }
2612
2613 #[test]
2614 fn test_open_nodes() {
2615 let kg = new_kg();
2616 kg.create_entities(&[
2617 Entity {
2618 name: "X".into(),
2619 entity_type: "n".into(),
2620 observations: vec!["obs_x".into()],
2621 attributes: None,
2622 },
2623 Entity {
2624 name: "Y".into(),
2625 entity_type: "n".into(),
2626 observations: vec!["obs_y".into()],
2627 attributes: None,
2628 },
2629 ])
2630 .unwrap();
2631
2632 kg.create_relations(&[rel_input("X", "Y", "e")]).unwrap();
2633
2634 let result = kg.open_nodes(&["X".into()]);
2635 let v: Value = serde_json::from_str(&result).unwrap();
2636 assert_eq!(v["entities"].as_array().unwrap().len(), 1);
2637 assert_eq!(v["relations"].as_array().unwrap().len(), 1);
2638 }
2639
2640 #[test]
2641 fn test_entities_exist() {
2642 let kg = new_kg();
2643 kg.create_entities(&[Entity {
2644 name: "exists".into(),
2645 entity_type: "t".into(),
2646 observations: vec![],
2647 attributes: None,
2648 }])
2649 .unwrap();
2650
2651 let res = kg
2652 .entities_exist(&["exists".into(), "missing".into()])
2653 .unwrap();
2654 assert_eq!(res, vec![true, false]);
2655 }
2656
2657 #[test]
2658 fn test_describe_entity() {
2659 let kg = new_kg();
2660 kg.create_entities(&[
2661 Entity {
2662 name: "A".into(),
2663 entity_type: "t".into(),
2664 observations: vec!["o".into()],
2665 attributes: None,
2666 },
2667 Entity {
2668 name: "B".into(),
2669 entity_type: "t".into(),
2670 observations: vec![],
2671 attributes: None,
2672 },
2673 Entity {
2674 name: "C".into(),
2675 entity_type: "t".into(),
2676 observations: vec![],
2677 attributes: None,
2678 },
2679 ])
2680 .unwrap();
2681
2682 kg.create_relations(&[
2683 rel_input("B", "A", "inbound"),
2684 rel_input("A", "B", "outbound"),
2685 rel_input("A", "C", "other"),
2686 rel_input("A", "A", "self"),
2687 ])
2688 .unwrap();
2689
2690 let entity = kg.describe_entity("A").unwrap();
2691 assert_eq!(entity.name, "A");
2692 assert_eq!(entity.entity_type, "t");
2693 assert_eq!(
2694 entity
2695 .observations
2696 .iter()
2697 .map(|o| o.body.as_str())
2698 .collect::<Vec<_>>(),
2699 ["o"]
2700 );
2701 assert_eq!(entity.relations.len(), 4);
2702 assert_eq!(
2703 entity.relations,
2704 vec![
2705 Relation {
2706 from: "A".into(),
2707 to: "A".into(),
2708 relation_type: "self".into(),
2709 },
2710 Relation {
2711 from: "A".into(),
2712 to: "B".into(),
2713 relation_type: "outbound".into(),
2714 },
2715 Relation {
2716 from: "A".into(),
2717 to: "C".into(),
2718 relation_type: "other".into(),
2719 },
2720 Relation {
2721 from: "B".into(),
2722 to: "A".into(),
2723 relation_type: "inbound".into(),
2724 },
2725 ]
2726 );
2727 assert_eq!(entity.neighbors, ["A", "B", "C"]);
2728 assert_eq!(entity.degree.incoming, 2);
2729 assert_eq!(entity.degree.outgoing, 3);
2730
2731 kg.writer
2735 .lock()
2736 .execute(
2737 "UPDATE entity SET out_deg = 99, in_deg = 88 WHERE name = 'A'",
2738 [],
2739 )
2740 .unwrap();
2741 let entity = kg.describe_entity("A").unwrap();
2742 assert_eq!(entity.degree.incoming, 2);
2743 assert_eq!(entity.degree.outgoing, 3);
2744 assert!(kg.describe_entity("missing").is_err());
2745 }
2746
2747 #[test]
2748 fn test_entity_type_counts() {
2749 let kg = new_kg();
2750 kg.create_entities(&[
2751 Entity {
2752 name: "a".into(),
2753 entity_type: "person".into(),
2754 observations: vec![],
2755 attributes: None,
2756 },
2757 Entity {
2758 name: "b".into(),
2759 entity_type: "person".into(),
2760 observations: vec![],
2761 attributes: None,
2762 },
2763 Entity {
2764 name: "c".into(),
2765 entity_type: "place".into(),
2766 observations: vec![],
2767 attributes: None,
2768 },
2769 ])
2770 .unwrap();
2771
2772 let counts = kg.entity_type_counts();
2773 let map: FxHashMap<_, _> = counts.into_iter().collect();
2774 assert_eq!(map.get("person"), Some(&2));
2775 assert_eq!(map.get("place"), Some(&1));
2776 }
2777
2778 #[test]
2779 fn test_relation_type_counts() {
2780 let kg = new_kg();
2781 kg.create_entities(&[
2782 Entity {
2783 name: "a".into(),
2784 entity_type: "n".into(),
2785 observations: vec![],
2786 attributes: None,
2787 },
2788 Entity {
2789 name: "b".into(),
2790 entity_type: "n".into(),
2791 observations: vec![],
2792 attributes: None,
2793 },
2794 Entity {
2795 name: "c".into(),
2796 entity_type: "n".into(),
2797 observations: vec![],
2798 attributes: None,
2799 },
2800 ])
2801 .unwrap();
2802
2803 kg.create_relations(&[rel_input("a", "b", "knows"), rel_input("a", "c", "knows")])
2804 .unwrap();
2805
2806 let counts = kg.relation_type_counts();
2807 let map: FxHashMap<_, _> = counts.into_iter().collect();
2808 assert_eq!(map.get("knows"), Some(&2));
2809 }
2810
2811 #[test]
2812 fn test_upsert_entities() {
2813 let kg = new_kg();
2814 kg.create_entities(&[Entity {
2815 name: "A".into(),
2816 entity_type: "OldType".into(),
2817 observations: vec!["old".into()],
2818 attributes: None,
2819 }])
2820 .unwrap();
2821 kg.create_relations(&[rel_input("A", "A", "self")]).unwrap();
2822
2823 kg.upsert_entities(&[Entity {
2825 name: "A".into(),
2826 entity_type: "NewType".into(),
2827 observations: vec!["old".into(), "new".into()],
2828 attributes: None,
2829 }])
2830 .unwrap();
2831
2832 assert_eq!(kg.get_entity_count().unwrap(), 1);
2833 let ent = kg.get_entity("A").unwrap().unwrap();
2834 assert_eq!(ent.entity_type, "NewType");
2835 assert_eq!(
2836 ent.observations
2837 .iter()
2838 .map(|o| o.body.as_str())
2839 .collect::<Vec<_>>(),
2840 ["old", "new"]
2841 );
2842
2843 let type_counts: FxHashMap<_, _> = kg.entity_type_counts().into_iter().collect();
2844 assert_eq!(type_counts.get("OldType"), None);
2845 assert_eq!(type_counts.get("NewType"), Some(&1));
2846
2847 assert_eq!(
2848 kg.search_relations(Some("A"), Some("A"), Some("self"), None, None)
2849 .unwrap(),
2850 [RelationDetail {
2851 from: "A".into(),
2852 to: "A".into(),
2853 relation_type: "self".into(),
2854 observations: vec![],
2855 attributes: BTreeMap::new(),
2856 }]
2857 );
2858 }
2859
2860 #[test]
2861 fn test_merge_entities() {
2862 let kg = new_kg();
2863 kg.create_entities(&[
2864 Entity {
2865 name: "source".into(),
2866 entity_type: "t".into(),
2867 observations: vec!["src_obs".into()],
2868 attributes: None,
2869 },
2870 Entity {
2871 name: "target".into(),
2872 entity_type: "t".into(),
2873 observations: vec!["tgt_obs".into()],
2874 attributes: None,
2875 },
2876 ])
2877 .unwrap();
2878
2879 kg.create_relations(&[rel_input("source", "target", "e")])
2880 .unwrap();
2881
2882 let merged = kg.merge_entities("source", "target").unwrap();
2883 assert_eq!(merged.name, "target");
2884 assert!(kg.get_entity("source").unwrap().is_none());
2885 }
2886
2887 #[test]
2888 fn test_find_all_paths() {
2889 let kg = new_kg();
2890 kg.create_entities(&[
2891 Entity {
2892 name: "A".into(),
2893 entity_type: "n".into(),
2894 observations: vec![],
2895 attributes: None,
2896 },
2897 Entity {
2898 name: "B".into(),
2899 entity_type: "n".into(),
2900 observations: vec![],
2901 attributes: None,
2902 },
2903 Entity {
2904 name: "C".into(),
2905 entity_type: "n".into(),
2906 observations: vec![],
2907 attributes: None,
2908 },
2909 ])
2910 .unwrap();
2911
2912 kg.create_relations(&[
2913 rel_input("A", "B", "e"),
2914 rel_input("B", "C", "e"),
2915 rel_input("A", "C", "e"),
2916 ])
2917 .unwrap();
2918
2919 let paths = kg.find_all_paths("A", "C", 5, 10).unwrap();
2920 assert!(paths.len() >= 2);
2921 }
2922
2923 #[test]
2924 fn test_batch_get_entities() {
2925 let kg = new_kg();
2926 kg.create_entities(&[
2927 Entity {
2928 name: "a".into(),
2929 entity_type: "t".into(),
2930 observations: vec![],
2931 attributes: None,
2932 },
2933 Entity {
2934 name: "b".into(),
2935 entity_type: "t".into(),
2936 observations: vec![],
2937 attributes: None,
2938 },
2939 ])
2940 .unwrap();
2941
2942 let results = kg.batch_get_entities(&["a".into(), "missing".into(), "b".into()]);
2943 assert_eq!(results.len(), 3);
2944 assert!(results[0].is_some());
2945 assert!(results[1].is_none());
2946 assert!(results[2].is_some());
2947 }
2948
2949 #[test]
2950 fn test_export_graph() {
2951 let kg = new_kg();
2952 kg.create_entities(&[Entity {
2953 name: "exp".into(),
2954 entity_type: "t".into(),
2955 observations: vec!["o".into()],
2956 attributes: None,
2957 }])
2958 .unwrap();
2959
2960 let exported = kg.export("json", i64::MAX).unwrap();
2961 assert!(exported.contains("exp"));
2962 assert!(exported.contains("o"));
2963 }
2964
2965 #[test]
2966 fn test_graph_stats() {
2967 let kg = new_kg();
2968 assert_eq!(kg.get_entity_count().unwrap(), 0);
2969 assert_eq!(kg.get_relation_count().unwrap(), 0);
2970
2971 kg.create_entities(&[Entity {
2972 name: "s".into(),
2973 entity_type: "t".into(),
2974 observations: vec![],
2975 attributes: None,
2976 }])
2977 .unwrap();
2978
2979 assert_eq!(kg.get_entity_count().unwrap(), 1);
2980 }
2981
2982 #[test]
2983 fn test_read_graph_filtered() {
2984 let kg = new_kg();
2985 kg.create_entities(&[
2986 Entity {
2987 name: "p1".into(),
2988 entity_type: "person".into(),
2989 observations: vec![],
2990 attributes: None,
2991 },
2992 Entity {
2993 name: "p2".into(),
2994 entity_type: "place".into(),
2995 observations: vec![],
2996 attributes: None,
2997 },
2998 ])
2999 .unwrap();
3000
3001 let out = kg.read_graph_filtered(Some("person"), 0, 10).unwrap();
3002 let v: Value = serde_json::from_str(&out).unwrap();
3003 assert_eq!(v["entities"].as_array().unwrap().len(), 1);
3004 assert_eq!(v["entities"][0]["name"], "p1");
3005 }
3006
3007 #[test]
3008 fn test_wipe() {
3009 let kg = new_kg();
3010 kg.create_entities(&[Entity {
3011 name: "w".into(),
3012 entity_type: "t".into(),
3013 observations: vec!["o".into()],
3014 attributes: None,
3015 }])
3016 .unwrap();
3017 assert_eq!(kg.get_entity_count().unwrap(), 1);
3018
3019 kg.wipe().unwrap();
3020 assert_eq!(kg.get_entity_count().unwrap(), 0);
3021 }
3022
3023 #[test]
3024 fn test_push_json_str() {
3025 let mut buf = String::new();
3026 push_json_str(&mut buf, "hello");
3027 assert_eq!(buf, "\"hello\"");
3028 let mut buf = String::new();
3029 push_json_str(&mut buf, "he\"llo");
3030 assert_eq!(buf, "\"he\\\"llo\"");
3031 }
3032
3033 #[test]
3036 fn test_create_entities_empty_input() {
3037 let kg = new_kg();
3038 let created = kg.create_entities(&[]).unwrap();
3039 assert!(created.is_empty());
3040 }
3041
3042 #[test]
3043 fn test_create_entities_skip_empty_name() {
3044 let kg = new_kg();
3045 let created = kg
3046 .create_entities(&[Entity {
3047 name: "".into(),
3048 entity_type: "t".into(),
3049 observations: vec![],
3050 attributes: None,
3051 }])
3052 .unwrap();
3053 assert!(created.is_empty());
3054 assert_eq!(kg.get_entity_count().unwrap(), 0);
3055 }
3056
3057 #[test]
3058 fn test_create_entities_duplicate_names() {
3059 let kg = new_kg();
3060 let e = Entity {
3061 name: "dup".into(),
3062 entity_type: "t".into(),
3063 observations: vec!["obs".into()],
3064 attributes: None,
3065 };
3066 let first = kg.create_entities(std::slice::from_ref(&e)).unwrap();
3067 assert_eq!(first.len(), 1);
3068 let second = kg.create_entities(&[e]).unwrap();
3069 assert!(second.is_empty());
3070 assert_eq!(kg.get_entity_count().unwrap(), 1);
3071 }
3072
3073 #[test]
3074 fn test_create_entities_partial_duplicates() {
3075 let kg = new_kg();
3076 let created = kg
3077 .create_entities(&[
3078 Entity {
3079 name: "a".into(),
3080 entity_type: "t".into(),
3081 observations: vec![],
3082 attributes: None,
3083 },
3084 Entity {
3085 name: "b".into(),
3086 entity_type: "t".into(),
3087 observations: vec![],
3088 attributes: None,
3089 },
3090 ])
3091 .unwrap();
3092 assert_eq!(created.len(), 2);
3093
3094 let second = kg
3095 .create_entities(&[
3096 Entity {
3097 name: "b".into(),
3098 entity_type: "t".into(),
3099 observations: vec![],
3100 attributes: None,
3101 },
3102 Entity {
3103 name: "c".into(),
3104 entity_type: "t".into(),
3105 observations: vec![],
3106 attributes: None,
3107 },
3108 ])
3109 .unwrap();
3110 assert_eq!(second.len(), 1); assert_eq!(second[0].name, "c");
3112 assert_eq!(kg.get_entity_count().unwrap(), 3);
3113 }
3114
3115 #[test]
3116 fn test_create_entities_mixed_empty_and_valid() {
3117 let kg = new_kg();
3118 let created = kg
3119 .create_entities(&[
3120 Entity {
3121 name: "".into(),
3122 entity_type: "t".into(),
3123 observations: vec![],
3124 attributes: None,
3125 },
3126 Entity {
3127 name: "valid".into(),
3128 entity_type: "t".into(),
3129 observations: vec![],
3130 attributes: None,
3131 },
3132 Entity {
3133 name: "".into(),
3134 entity_type: "t".into(),
3135 observations: vec![],
3136 attributes: None,
3137 },
3138 ])
3139 .unwrap();
3140 assert_eq!(created.len(), 1);
3141 assert_eq!(created[0].name, "valid");
3142 assert_eq!(kg.get_entity_count().unwrap(), 1);
3143 }
3144
3145 #[test]
3146 fn test_create_entities_same_name_in_batch() {
3147 let kg = new_kg();
3148 let created = kg
3149 .create_entities(&[
3150 Entity {
3151 name: "dup_in_batch".into(),
3152 entity_type: "t".into(),
3153 observations: vec![],
3154 attributes: None,
3155 },
3156 Entity {
3157 name: "dup_in_batch".into(),
3158 entity_type: "t".into(),
3159 observations: vec![],
3160 attributes: None,
3161 },
3162 ])
3163 .unwrap();
3164 assert_eq!(created.len(), 1);
3165 assert_eq!(kg.get_entity_count().unwrap(), 1);
3166 }
3167
3168 #[test]
3171 fn test_create_relations_empty_input() {
3172 let kg = new_kg();
3173 let rels = kg.create_relations(&[]).unwrap();
3174 assert!(rels.is_empty());
3175 }
3176
3177 #[test]
3178 fn test_create_relations_nonexistent_from() {
3179 let kg = new_kg();
3180 kg.create_entities(&[Entity {
3181 name: "B".into(),
3182 entity_type: "t".into(),
3183 observations: vec![],
3184 attributes: None,
3185 }])
3186 .unwrap();
3187
3188 let rels = kg.create_relations(&[rel_input("A", "B", "e")]).unwrap();
3189 assert!(rels.is_empty());
3190 assert_eq!(kg.get_relation_count().unwrap(), 0);
3191 }
3192
3193 #[test]
3194 fn test_create_relations_nonexistent_to() {
3195 let kg = new_kg();
3196 kg.create_entities(&[Entity {
3197 name: "A".into(),
3198 entity_type: "t".into(),
3199 observations: vec![],
3200 attributes: None,
3201 }])
3202 .unwrap();
3203
3204 let rels = kg.create_relations(&[rel_input("A", "B", "e")]).unwrap();
3205 assert!(rels.is_empty());
3206 assert_eq!(kg.get_relation_count().unwrap(), 0);
3207 }
3208
3209 #[test]
3210 fn test_create_relations_both_nonexistent() {
3211 let kg = new_kg();
3212 let rels = kg.create_relations(&[rel_input("A", "B", "e")]).unwrap();
3213 assert!(rels.is_empty());
3214 }
3215
3216 #[test]
3217 fn test_create_relations_self_loop() {
3218 let kg = new_kg();
3219 kg.create_entities(&[Entity {
3220 name: "self".into(),
3221 entity_type: "t".into(),
3222 observations: vec![],
3223 attributes: None,
3224 }])
3225 .unwrap();
3226
3227 let rels = kg
3228 .create_relations(&[rel_input("self", "self", "loop")])
3229 .unwrap();
3230 assert_eq!(rels.len(), 1);
3231 assert_eq!(kg.get_relation_count().unwrap(), 1);
3232 assert_eq!(kg.degree("self", Direction::Outgoing).unwrap(), 1);
3233 assert_eq!(kg.degree("self", Direction::Incoming).unwrap(), 1);
3234 }
3235
3236 #[test]
3237 fn test_create_relations_duplicate() {
3238 let kg = new_kg();
3239 kg.create_entities(&[
3240 Entity {
3241 name: "A".into(),
3242 entity_type: "t".into(),
3243 observations: vec![],
3244 attributes: None,
3245 },
3246 Entity {
3247 name: "B".into(),
3248 entity_type: "t".into(),
3249 observations: vec![],
3250 attributes: None,
3251 },
3252 ])
3253 .unwrap();
3254
3255 let r = RelationInput {
3256 from: "A".into(),
3257 to: "B".into(),
3258 relation_type: "e".into(),
3259 observations: vec![],
3260 attributes: None,
3261 };
3262 let first = kg.create_relations(std::slice::from_ref(&r)).unwrap();
3263 assert_eq!(first.len(), 1);
3264
3265 let second = kg.create_relations(&[r]).unwrap();
3266 assert!(second.is_empty());
3267 assert_eq!(kg.get_relation_count().unwrap(), 1);
3268 }
3269
3270 #[test]
3271 fn test_create_relations_new_type_auto_created() {
3272 let kg = new_kg();
3273 kg.create_entities(&[
3274 Entity {
3275 name: "A".into(),
3276 entity_type: "t".into(),
3277 observations: vec![],
3278 attributes: None,
3279 },
3280 Entity {
3281 name: "B".into(),
3282 entity_type: "t".into(),
3283 observations: vec![],
3284 attributes: None,
3285 },
3286 ])
3287 .unwrap();
3288
3289 let rels = kg
3290 .create_relations(&[rel_input("A", "B", "brand_new_type")])
3291 .unwrap();
3292 assert_eq!(rels.len(), 1);
3293
3294 let counts = kg.relation_type_counts();
3295 let map: FxHashMap<_, _> = counts.into_iter().collect();
3296 assert_eq!(map.get("brand_new_type"), Some(&1));
3297 }
3298
3299 #[test]
3300 fn test_create_relations_degree_updates() {
3301 let kg = new_kg();
3302 kg.create_entities(&[
3303 Entity {
3304 name: "A".into(),
3305 entity_type: "t".into(),
3306 observations: vec![],
3307 attributes: None,
3308 },
3309 Entity {
3310 name: "B".into(),
3311 entity_type: "t".into(),
3312 observations: vec![],
3313 attributes: None,
3314 },
3315 Entity {
3316 name: "C".into(),
3317 entity_type: "t".into(),
3318 observations: vec![],
3319 attributes: None,
3320 },
3321 ])
3322 .unwrap();
3323
3324 kg.create_relations(&[rel_input("A", "B", "e"), rel_input("A", "C", "e")])
3325 .unwrap();
3326
3327 assert_eq!(kg.degree("A", Direction::Outgoing).unwrap(), 2);
3328 assert_eq!(kg.degree("A", Direction::Incoming).unwrap(), 0);
3329 assert_eq!(kg.degree("B", Direction::Incoming).unwrap(), 1);
3330 assert_eq!(kg.degree("C", Direction::Incoming).unwrap(), 1);
3331 assert_eq!(kg.degree("A", Direction::Both).unwrap(), 2);
3332 }
3333
3334 #[test]
3335 fn test_create_relations_delete_and_recreate() {
3336 let kg = new_kg();
3337 kg.create_entities(&[
3338 Entity {
3339 name: "A".into(),
3340 entity_type: "t".into(),
3341 observations: vec![],
3342 attributes: None,
3343 },
3344 Entity {
3345 name: "B".into(),
3346 entity_type: "t".into(),
3347 observations: vec![],
3348 attributes: None,
3349 },
3350 ])
3351 .unwrap();
3352
3353 let r = Relation {
3354 from: "A".into(),
3355 to: "B".into(),
3356 relation_type: "e".into(),
3357 };
3358 let input = RelationInput {
3359 from: r.from.clone(),
3360 to: r.to.clone(),
3361 relation_type: r.relation_type.clone(),
3362 observations: vec![],
3363 attributes: None,
3364 };
3365 kg.create_relations(std::slice::from_ref(&input)).unwrap();
3366 assert_eq!(kg.get_relation_count().unwrap(), 1);
3367
3368 kg.delete_relations(std::slice::from_ref(&r)).unwrap();
3369 assert_eq!(kg.get_relation_count().unwrap(), 0);
3370
3371 let re = kg.create_relations(&[input]).unwrap();
3373 assert_eq!(re.len(), 1);
3374 assert_eq!(kg.get_relation_count().unwrap(), 1);
3375 }
3376
3377 #[test]
3380 fn test_create_entities_then_relations_then_delete_entity_with_relations() {
3381 let kg = new_kg();
3382 kg.create_entities(&[
3383 Entity {
3384 name: "A".into(),
3385 entity_type: "t".into(),
3386 observations: vec![],
3387 attributes: None,
3388 },
3389 Entity {
3390 name: "B".into(),
3391 entity_type: "t".into(),
3392 observations: vec![],
3393 attributes: None,
3394 },
3395 ])
3396 .unwrap();
3397 kg.create_relations(&[rel_input("A", "B", "e")]).unwrap();
3398
3399 assert_eq!(kg.get_relation_count().unwrap(), 1);
3400
3401 kg.delete_entities(&["A".into()]).unwrap();
3403 assert!(kg.get_entity("A").unwrap().is_none());
3404 assert_eq!(kg.get_relation_count().unwrap(), 0);
3405 }
3406
3407 #[test]
3408 fn test_graph_stats_after_entity_with_observations() {
3409 let kg = new_kg();
3410 kg.create_entities(&[Entity {
3411 name: "stat".into(),
3412 entity_type: "t".into(),
3413 observations: vec!["o1".into(), "o2".into(), "o3".into()],
3414 attributes: None,
3415 }])
3416 .unwrap();
3417
3418 let ecount = kg.get_entity_count().unwrap();
3419 assert_eq!(ecount, 1);
3421
3422 kg.delete_entities(&["stat".into()]).unwrap();
3424 assert_eq!(kg.get_entity_count().unwrap(), 0);
3425 }
3426
3427 fn new_kg_with_pool(read_pool_size: usize) -> TestKg {
3430 use std::sync::atomic::AtomicU64;
3431 static COUNTER: AtomicU64 = AtomicU64::new(1_000_000);
3432 let n = COUNTER.fetch_add(1, Ordering::SeqCst);
3433 let path = std::env::temp_dir().join(format!("kg_pool_{}_{}.db", std::process::id(), n));
3434 cleanup_db(&path);
3435 let kg = GraphHandle::new(
3436 &path,
3437 Durability::Async,
3438 SqliteTuning::default(),
3439 NonZeroUsize::new(10_000).unwrap(),
3440 read_pool_size,
3441 )
3442 .expect("create KG");
3443 TestKg(kg, path)
3444 }
3445
3446 fn seed_line(kg: &GraphHandle, n: usize) {
3447 let entities: Vec<Entity> = (0..n)
3448 .map(|i| Entity {
3449 name: format!("n{i}"),
3450 entity_type: "node".into(),
3451 observations: vec![format!("obs of n{i}").into()],
3452 attributes: None,
3453 })
3454 .collect();
3455 kg.create_entities(&entities).unwrap();
3456 let rels: Vec<RelationInput> = (0..n.saturating_sub(1))
3457 .map(|i| RelationInput {
3458 from: format!("n{i}"),
3459 to: format!("n{}", i + 1),
3460 relation_type: "edge".into(),
3461 observations: vec![],
3462 attributes: None,
3463 })
3464 .collect();
3465 if !rels.is_empty() {
3466 kg.create_relations(&rels).unwrap();
3467 }
3468 }
3469
3470 fn count_relations(graph_json: &str) -> usize {
3471 let v: Value = serde_json::from_str(graph_json).unwrap();
3472 v["relations"].as_array().unwrap().len()
3473 }
3474
3475 fn count_entities(graph_json: &str) -> usize {
3476 let v: Value = serde_json::from_str(graph_json).unwrap();
3477 v["entities"].as_array().unwrap().len()
3478 }
3479
3480 #[test]
3483 fn test_pool_size_one_still_works() {
3484 let kg = new_kg_with_pool(1);
3485 seed_line(&kg, 5);
3486 assert_eq!(kg.get_entity_count().unwrap(), 5);
3487 assert!(kg.get_entity("n2").unwrap().is_some());
3488 let g = kg.read_graph_filtered(None, 0, usize::MAX).unwrap();
3489 assert_eq!(count_entities(&g), 5);
3490 }
3491
3492 #[test]
3493 fn test_reads_see_committed_writes() {
3494 let kg = new_kg_with_pool(4);
3497 kg.create_entities(&[Entity {
3498 name: "fresh".into(),
3499 entity_type: "t".into(),
3500 observations: vec!["v".into()],
3501 attributes: None,
3502 }])
3503 .unwrap();
3504 let got = kg.get_entity("fresh").unwrap().unwrap();
3506 assert_eq!(
3507 got.observations
3508 .iter()
3509 .map(|o| o.body.as_str())
3510 .collect::<Vec<_>>(),
3511 vec!["v"]
3512 );
3513 }
3514
3515 #[test]
3516 fn test_concurrent_readers_consistent() {
3517 let kg = new_kg_with_pool(4);
3520 seed_line(&kg, 50);
3521
3522 std::thread::scope(|s| {
3523 for _ in 0..8 {
3525 s.spawn(|| {
3526 for _ in 0..200 {
3527 let _ = kg.get_entity("n10");
3528 let _ = kg.search_nodes_filtered("obs", None, 0, 10);
3529 let _ = kg.read_graph_filtered(None, 0, 100);
3530 let _ = kg.get_entity_count();
3531 let _ = kg.neighbors("n10", Direction::Both, None, 2);
3532 }
3533 });
3534 }
3535 s.spawn(|| {
3537 for i in 100..160 {
3538 kg.create_entities(&[Entity {
3539 name: format!("w{i}"),
3540 entity_type: "node".into(),
3541 observations: vec![format!("w obs {i}").into()],
3542 attributes: None,
3543 }])
3544 .unwrap();
3545 }
3546 });
3547 });
3548
3549 assert_eq!(kg.get_entity_count().unwrap(), 110);
3551 assert!(kg.get_entity("w159").unwrap().is_some());
3552 }
3553
3554 #[test]
3555 fn test_reader_pool_rejects_writes_internally() {
3556 let kg = new_kg_with_pool(1);
3562 seed_line(&kg, 3);
3563 std::thread::scope(|s| {
3564 for _ in 0..4 {
3565 s.spawn(|| {
3566 for _ in 0..100 {
3567 let _ = kg.read_graph_filtered(None, 0, 10);
3568 }
3569 });
3570 }
3571 });
3572 assert_eq!(kg.get_entity_count().unwrap(), 3);
3573 }
3574
3575 #[test]
3578 fn test_read_graph_relations_scoped_to_page() {
3579 let kg = new_kg_with_pool(2);
3580 seed_line(&kg, 4);
3582
3583 let full = kg.read_graph_filtered(None, 0, usize::MAX).unwrap();
3585 assert_eq!(count_entities(&full), 4);
3586 assert_eq!(count_relations(&full), 3);
3587
3588 let page1 = kg.read_graph_filtered(None, 0, 1).unwrap();
3591 assert_eq!(count_entities(&page1), 1);
3592 assert_eq!(count_relations(&page1), 0);
3593
3594 let page2 = kg.read_graph_filtered(None, 0, 2).unwrap();
3597 assert_eq!(count_entities(&page2), 2);
3598 assert_eq!(count_relations(&page2), 1);
3599 }
3600
3601 #[test]
3602 fn test_read_graph_pagination_offset() {
3603 let kg = new_kg_with_pool(2);
3604 seed_line(&kg, 5);
3605 let g = kg.read_graph_filtered(None, 2, 2).unwrap();
3606 assert_eq!(count_entities(&g), 2);
3607 assert!(!g.contains("\"n0\""));
3609 assert!(!g.contains("\"n1\""));
3610 assert!(g.contains("\"n2\""));
3611 assert!(g.contains("\"n3\""));
3612 }
3613
3614 #[test]
3615 fn test_read_graph_empty() {
3616 let kg = new_kg_with_pool(2);
3617 let g = kg.read_graph_filtered(None, 0, usize::MAX).unwrap();
3618 assert_eq!(g, r#"{"entities":[],"relations":[]}"#);
3619 }
3620
3621 #[test]
3622 fn test_read_graph_filtered_by_type() {
3623 let kg = new_kg_with_pool(2);
3624 kg.create_entities(&[
3625 Entity {
3626 name: "p1".into(),
3627 entity_type: "person".into(),
3628 observations: vec![],
3629 attributes: None,
3630 },
3631 Entity {
3632 name: "q1".into(),
3633 entity_type: "place".into(),
3634 observations: vec![],
3635 attributes: None,
3636 },
3637 Entity {
3638 name: "p2".into(),
3639 entity_type: "person".into(),
3640 observations: vec![],
3641 attributes: None,
3642 },
3643 ])
3644 .unwrap();
3645 let g = kg
3646 .read_graph_filtered(Some("person"), 0, usize::MAX)
3647 .unwrap();
3648 assert_eq!(count_entities(&g), 2);
3649 assert!(g.contains("\"p1\""));
3650 assert!(g.contains("\"p2\""));
3651 assert!(!g.contains("\"q1\""));
3652 }
3653
3654 #[test]
3655 fn test_export_respects_max_rows() {
3656 let kg = new_kg_with_pool(2);
3657 seed_line(&kg, 5);
3658
3659 let full = kg.export("json", i64::MAX).unwrap();
3661 assert_eq!(count_entities(&full), 5);
3662 assert_eq!(count_relations(&full), 4);
3663
3664 let capped = kg.export("json", 2).unwrap();
3666 assert_eq!(count_entities(&capped), 2);
3667 assert_eq!(count_relations(&capped), 2);
3668 }
3669
3670 #[test]
3671 fn test_export_negative_max_rows_is_unbounded() {
3672 let kg = new_kg_with_pool(2);
3673 seed_line(&kg, 3);
3674 let out = kg.export("json", -1).unwrap();
3676 assert_eq!(count_entities(&out), 3);
3677 }
3678
3679 #[test]
3682 fn test_many_small_write_batches_stay_consistent() {
3683 let kg = new_kg_with_pool(2);
3684 for i in 0..100 {
3685 kg.create_entities(&[Entity {
3686 name: format!("e{i}"),
3687 entity_type: "t".into(),
3688 observations: vec![format!("o{i}").into()],
3689 attributes: None,
3690 }])
3691 .unwrap();
3692 }
3693 assert_eq!(kg.get_entity_count().unwrap(), 100);
3694 let hits = kg.search_nodes_filtered("e57", None, 0, 10);
3697 assert!(hits.iter().any(|e| e.name == "e57"));
3698 }
3699
3700 #[test]
3703 fn test_wipe_clears_name_and_obs_fts() {
3704 let kg = new_kg_with_pool(2);
3705 kg.create_entities(&[Entity {
3706 name: "Einstein".into(),
3707 entity_type: "scientist".into(),
3708 observations: vec!["physics".into()],
3709 attributes: None,
3710 }])
3711 .unwrap();
3712
3713 assert_eq!(kg.search_nodes_filtered("Einstein", None, 0, 10).len(), 1);
3715 assert_eq!(kg.search_nodes_filtered("physics", None, 0, 10).len(), 1);
3716
3717 kg.wipe().unwrap();
3718
3719 assert_eq!(kg.get_entity_count().unwrap(), 0);
3722 assert!(kg.search_nodes_filtered("Einstein", None, 0, 10).is_empty());
3723 assert!(kg.search_nodes_filtered("physics", None, 0, 10).is_empty());
3724 }
3725
3726 #[test]
3727 fn test_wipe_then_recreate_search_works() {
3728 let kg = new_kg_with_pool(2);
3731 kg.create_entities(&[Entity {
3732 name: "Einstein".into(),
3733 entity_type: "scientist".into(),
3734 observations: vec!["physics".into()],
3735 attributes: None,
3736 }])
3737 .unwrap();
3738 kg.wipe().unwrap();
3739
3740 kg.create_entities(&[Entity {
3741 name: "Einstein".into(),
3742 entity_type: "scientist".into(),
3743 observations: vec!["physics".into(), "relativity".into()],
3744 attributes: None,
3745 }])
3746 .unwrap();
3747
3748 let by_name = kg.search_nodes_filtered("Einstein", None, 0, 10);
3749 assert_eq!(by_name.len(), 1, "exactly one Einstein after recreate");
3750 let by_obs = kg.search_nodes_filtered("relativity", None, 0, 10);
3751 assert_eq!(by_obs.len(), 1);
3752 assert_eq!(kg.get_entity_count().unwrap(), 1);
3753 }
3754
3755 #[test]
3758 fn test_search_relations_missing_type_returns_empty() {
3759 let kg = new_kg_with_pool(2);
3760 seed_line(&kg, 3); let r = kg
3764 .search_relations(None, None, Some("does_not_exist"), None, None)
3765 .unwrap();
3766 assert!(r.is_empty());
3767 let types = kg.relation_type_counts();
3769 assert!(types.iter().all(|(t, _)| t != "does_not_exist"));
3770 }
3771
3772 #[test]
3773 fn test_entity_type_exists() {
3774 let kg = new_kg_with_pool(2);
3775 kg.create_entities(&[Entity {
3776 name: "a".into(),
3777 entity_type: "person".into(),
3778 observations: vec![],
3779 attributes: None,
3780 }])
3781 .unwrap();
3782 assert!(kg.entity_type_exists("person"));
3783 assert!(!kg.entity_type_exists("persn"));
3784 let types = kg.entity_type_counts();
3786 assert!(types.iter().all(|(t, _)| t != "persn"));
3787 }
3788
3789 #[test]
3790 fn test_type_descriptions_registry_and_catalog() {
3791 let kg = new_kg_with_pool(2);
3792 kg.set_type_description(0, "person", Some("A human being or persona"))
3794 .unwrap();
3795 let catalog = kg.entity_type_catalog();
3796 assert_eq!(
3797 catalog,
3798 vec![(
3799 "person".into(),
3800 0usize,
3801 Some("A human being or persona".into())
3802 )]
3803 );
3804 assert!(kg.entity_type_exists("person"));
3806 assert!(
3807 kg.entity_type_counts().is_empty(),
3808 "count 0 types stay out of counts"
3809 );
3810
3811 kg.set_type_description(0, "person", Some("A living human"))
3813 .unwrap();
3814 let catalog = kg.entity_type_catalog();
3815 assert_eq!(
3816 catalog,
3817 vec![("person".into(), 0usize, Some("A living human".into()))]
3818 );
3819
3820 kg.set_type_description(0, "person", None).unwrap();
3823 assert!(kg.entity_type_catalog().is_empty());
3824 assert!(
3825 kg.entity_type_exists("person"),
3826 "the row stays; only the desc clears"
3827 );
3828
3829 kg.create_entities(&[Entity {
3832 name: "alice".into(),
3833 entity_type: "person".into(),
3834 observations: vec![],
3835 attributes: None,
3836 }])
3837 .unwrap();
3838 assert_eq!(
3839 kg.entity_type_catalog(),
3840 vec![("person".into(), 1usize, None)]
3841 );
3842 assert_eq!(kg.entity_type_counts(), vec![("person".into(), 1usize)]);
3843
3844 kg.set_type_description(1, "works_at", Some("Employment link"))
3846 .unwrap();
3847 assert_eq!(
3848 kg.relation_type_catalog(),
3849 vec![("works_at".into(), 0usize, Some("Employment link".into()))]
3850 );
3851 assert!(kg.relation_type_exists("works_at"));
3852 }
3853
3854 #[test]
3855 fn test_type_description_length_not_capped_at_core() {
3856 let kg = new_kg_with_pool(2);
3858 let long = "x".repeat(20_000);
3859 kg.set_type_description(0, "person", Some(long.as_str()))
3860 .unwrap();
3861 assert_eq!(kg.entity_type_catalog().len(), 1);
3862 }
3863
3864 #[test]
3865 fn test_relation_type_exists() {
3866 let kg = new_kg_with_pool(2);
3867 kg.create_entities(&[
3868 Entity {
3869 name: "a".into(),
3870 entity_type: "person".into(),
3871 observations: vec![],
3872 attributes: None,
3873 },
3874 Entity {
3875 name: "b".into(),
3876 entity_type: "person".into(),
3877 observations: vec![],
3878 attributes: None,
3879 },
3880 ])
3881 .unwrap();
3882 kg.create_relations(&[rel_input("a", "b", "knows")])
3883 .unwrap();
3884 assert!(kg.relation_type_exists("knows"));
3885 assert!(!kg.relation_type_exists("unknown_kind"));
3886 let types = kg.relation_type_counts();
3888 assert!(types.iter().all(|(t, _)| t != "unknown_kind"));
3889 }
3890
3891 #[test]
3892 fn test_search_relations_missing_from_returns_empty() {
3893 let kg = new_kg_with_pool(2);
3894 seed_line(&kg, 3);
3895 let r = kg
3896 .search_relations(Some("ghost"), None, None, None, None)
3897 .unwrap();
3898 assert!(r.is_empty(), "missing 'from' must not match every relation");
3899 }
3900
3901 #[test]
3902 fn test_search_relations_existing_filters_still_work() {
3903 let kg = new_kg_with_pool(2);
3904 seed_line(&kg, 3);
3905 let r = kg
3906 .search_relations(Some("n0"), None, Some("edge"), None, None)
3907 .unwrap();
3908 assert_eq!(r.len(), 1);
3909 assert_eq!(r[0].from, "n0");
3910 assert_eq!(r[0].to, "n1");
3911 }
3912
3913 #[test]
3914 fn test_neighbors_missing_type_returns_only_start() {
3915 let kg = new_kg_with_pool(2);
3916 seed_line(&kg, 3);
3917 let json = kg
3918 .neighbors("n0", Direction::Both, Some("nonexistent"), 2)
3919 .unwrap();
3920 assert_eq!(count_entities(&json), 1);
3922 assert_eq!(count_relations(&json), 0);
3923 }
3924
3925 #[test]
3926 fn test_neighbors_existing_type_filters() {
3927 let kg = new_kg_with_pool(2);
3928 kg.create_entities(&[
3929 Entity {
3930 name: "a".into(),
3931 entity_type: "n".into(),
3932 observations: vec![],
3933 attributes: None,
3934 },
3935 Entity {
3936 name: "b".into(),
3937 entity_type: "n".into(),
3938 observations: vec![],
3939 attributes: None,
3940 },
3941 Entity {
3942 name: "c".into(),
3943 entity_type: "n".into(),
3944 observations: vec![],
3945 attributes: None,
3946 },
3947 ])
3948 .unwrap();
3949 kg.create_relations(&[rel_input("a", "b", "knows"), rel_input("a", "c", "likes")])
3950 .unwrap();
3951 let json = kg
3952 .neighbors("a", Direction::Outgoing, Some("knows"), 1)
3953 .unwrap();
3954 assert!(json.contains("\"b\""));
3955 assert!(!json.contains("\"c\""));
3956 assert_eq!(count_relations(&json), 1);
3957 }
3958
3959 #[test]
3960 fn test_sqlite_tuning_applied_to_fresh_db() {
3961 use std::sync::atomic::AtomicU64;
3962 static COUNTER: AtomicU64 = AtomicU64::new(2_000_000);
3963 let n = COUNTER.fetch_add(1, Ordering::SeqCst);
3964 let path = std::env::temp_dir().join(format!("kg_tuning_{}_{}.db", std::process::id(), n));
3965 cleanup_db(&path);
3966
3967 let tuning = SqliteTuning {
3968 page_size: 8192,
3969 ..SqliteTuning::default()
3970 };
3971 let kg = TestKg(
3972 GraphHandle::new(
3973 &path,
3974 Durability::Async,
3975 tuning,
3976 NonZeroUsize::new(64).unwrap(),
3977 2,
3978 )
3979 .expect("create KG"),
3980 path.clone(),
3981 );
3982 kg.create_entities(&[Entity {
3983 name: "a".into(),
3984 entity_type: "n".into(),
3985 observations: vec!["o".into()],
3986 attributes: None,
3987 }])
3988 .unwrap();
3989
3990 let probe = Connection::open(&path).unwrap();
3993 let page_size: i64 = probe
3994 .query_row("PRAGMA page_size", [], |r| r.get(0))
3995 .unwrap();
3996 assert_eq!(page_size, 8192);
3997 let auto_vacuum: i64 = probe
3998 .query_row("PRAGMA auto_vacuum", [], |r| r.get(0))
3999 .unwrap();
4000 assert_eq!(auto_vacuum, 2, "expected INCREMENTAL auto_vacuum");
4001 let journal: String = probe
4002 .query_row("PRAGMA journal_mode", [], |r| r.get(0))
4003 .unwrap();
4004 assert_eq!(journal.to_lowercase(), "wal");
4005 }
4006
4007 #[test]
4008 fn test_checkpoint_passive_is_noop_safe() {
4009 let kg = new_kg();
4010 kg.checkpoint_passive().unwrap();
4012 kg.create_entities(&[Entity {
4013 name: "a".into(),
4014 entity_type: "n".into(),
4015 observations: vec!["o".into()],
4016 attributes: None,
4017 }])
4018 .unwrap();
4019 kg.checkpoint_passive().unwrap();
4021 kg.checkpoint_passive().unwrap();
4022 assert!(kg.get_entity("a").unwrap().is_some());
4024 }
4025
4026 fn rel_input(from: &str, to: &str, relation_type: &str) -> RelationInput {
4029 RelationInput {
4030 from: from.into(),
4031 to: to.into(),
4032 relation_type: relation_type.into(),
4033 observations: vec![],
4034 attributes: None,
4035 }
4036 }
4037
4038 fn seed_relation_obs_attrs(kg: &GraphHandle) -> i64 {
4039 kg.create_entities(&[
4040 Entity {
4041 name: "a".into(),
4042 entity_type: "n".into(),
4043 observations: vec![],
4044 attributes: None,
4045 },
4046 Entity {
4047 name: "b".into(),
4048 entity_type: "n".into(),
4049 observations: vec![],
4050 attributes: None,
4051 },
4052 ])
4053 .unwrap();
4054 kg.create_relations(&[RelationInput {
4055 from: "a".into(),
4056 to: "b".into(),
4057 relation_type: "uses".into(),
4058 observations: vec!["contract #12".into(), "legacy".into()],
4059 attributes: Some(std::collections::BTreeMap::from([("k".into(), "v".into())])),
4060 }])
4061 .unwrap();
4062 let conn = kg.writer.lock();
4063 conn.query_row(
4064 "SELECT m.id FROM taxonomy_relation m
4065 JOIN entity f ON f.id = m.from_id AND f.name = 'a'",
4066 [],
4067 |row| row.get(0),
4068 )
4069 .unwrap()
4070 }
4071
4072 #[test]
4073 fn get_entity_and_describe_include_attributes() {
4074 let kg = new_kg();
4075 kg.create_entities(&[Entity {
4076 name: "a".into(),
4077 entity_type: "t".into(),
4078 observations: vec![],
4079 attributes: None,
4080 }])
4081 .unwrap();
4082 let conn = kg.writer.lock();
4083 let id: i64 = conn
4084 .query_row("SELECT id FROM entity WHERE name='a'", [], |row| row.get(0))
4085 .unwrap();
4086 conn.execute(
4087 "INSERT INTO attribute(owner_kind,owner_id,key,value,created_us,updated_us)
4088 VALUES('entity',?1,'k','v',1,1)",
4089 [id],
4090 )
4091 .unwrap();
4092 drop(conn);
4093
4094 let got = kg.get_entity("a").unwrap().unwrap();
4095 assert_eq!(
4096 got.attributes,
4097 Some(std::collections::BTreeMap::from([("k".into(), "v".into())]))
4098 );
4099 let described = kg.describe_entity("a").unwrap();
4100 assert_eq!(
4101 described.attributes,
4102 Some(std::collections::BTreeMap::from([("k".into(), "v".into())]))
4103 );
4104
4105 kg.create_entities(&[Entity {
4107 name: "plain".into(),
4108 entity_type: "t".into(),
4109 observations: vec![],
4110 attributes: None,
4111 }])
4112 .unwrap();
4113 assert_eq!(kg.get_entity("plain").unwrap().unwrap().attributes, None);
4114 assert_eq!(kg.describe_entity("plain").unwrap().attributes, None);
4115 }
4116
4117 #[test]
4118 fn batch_get_entities_include_attributes() {
4119 let kg = new_kg();
4120 kg.create_entities(&[
4121 Entity {
4122 name: "a".into(),
4123 entity_type: "t".into(),
4124 observations: vec![],
4125 attributes: None,
4126 },
4127 Entity {
4128 name: "b".into(),
4129 entity_type: "t".into(),
4130 observations: vec![],
4131 attributes: None,
4132 },
4133 ])
4134 .unwrap();
4135 let conn = kg.writer.lock();
4136 conn.execute(
4137 "INSERT INTO attribute(owner_kind,owner_id,key,value,created_us,updated_us)
4138 SELECT 'entity', e.id, 'k', 'v', 1, 1 FROM entity e WHERE e.name='a'",
4139 [],
4140 )
4141 .unwrap();
4142 drop(conn);
4143
4144 let results = kg.batch_get_entities(&["a".into(), "missing".into(), "b".into()]);
4145 assert_eq!(results.len(), 3);
4146 assert_eq!(
4147 results[0].as_ref().unwrap().attributes,
4148 Some(std::collections::BTreeMap::from([("k".into(), "v".into())]))
4149 );
4150 assert!(results[1].is_none());
4151 assert_eq!(results[2].as_ref().unwrap().attributes, None);
4152
4153 let dupes = kg.batch_get_entities(&["a".into(), "a".into()]);
4156 assert_eq!(dupes.len(), 2);
4157 assert!(dupes[0].is_some());
4158 assert!(dupes[1].is_some(), "duplicate input names each resolve");
4159 }
4160
4161 #[test]
4162 fn search_relations_query_matches_relation_observation_bodies() {
4163 let kg = new_kg();
4164 seed_relation_obs_attrs(&kg);
4165
4166 let detail = kg
4168 .search_relations(None, None, None, Some("contract"), Some(10usize))
4169 .unwrap();
4170 assert_eq!(detail.len(), 1);
4171 assert_eq!(detail[0].from, "a");
4172 assert_eq!(detail[0].to, "b");
4173 assert_eq!(detail[0].relation_type, "uses");
4174 assert_eq!(
4175 detail[0]
4176 .observations
4177 .iter()
4178 .map(|o| o.body.as_str())
4179 .collect::<Vec<_>>(),
4180 ["contract #12", "legacy"],
4181 "owner-level detail carries the full observation list in idx order"
4182 );
4183 assert_eq!(
4184 detail[0].attributes,
4185 std::collections::BTreeMap::from([("k".into(), "v".into())])
4186 );
4187
4188 let filtered = kg
4190 .search_relations(
4191 Some("a"),
4192 Some("b"),
4193 Some("uses"),
4194 Some("legacy"),
4195 Some(10usize),
4196 )
4197 .unwrap();
4198 assert_eq!(filtered.len(), 1);
4199 let no_match = kg
4200 .search_relations(Some("ghost"), None, None, Some("contract"), Some(10usize))
4201 .unwrap();
4202 assert!(no_match.is_empty(), "missing entity filter yields nothing");
4203
4204 let all = kg
4206 .search_relations(None, None, None, None, Some(10usize))
4207 .unwrap();
4208 assert_eq!(all.len(), 1);
4209 assert_eq!(
4210 all[0]
4211 .observations
4212 .iter()
4213 .map(|o| o.body.as_str())
4214 .collect::<Vec<_>>(),
4215 ["contract #12", "legacy"],
4216 "observations always present, ordered by idx"
4217 );
4218 assert_eq!(
4219 all[0].attributes,
4220 std::collections::BTreeMap::from([("k".into(), "v".into())]),
4221 "attributes always present"
4222 );
4223 }
4224
4225 #[test]
4226 fn export_contains_relation_observations_and_attributes() {
4227 let kg = new_kg();
4228 seed_relation_obs_attrs(&kg);
4229 kg.set_attributes(&[AttributeSet {
4231 owner_kind: "entity".into(),
4232 entity_name: Some("a".into()),
4233 from: None,
4234 to: None,
4235 relation_type: None,
4236 attributes: std::collections::BTreeMap::from([("ea".into(), "ev".into())]),
4237 }])
4238 .unwrap();
4239
4240 let exported = kg.export("json", 100).unwrap();
4241 assert!(exported.contains("\"contract #12\""), "observation body");
4242 assert!(exported.contains("\"legacy\""), "second observation body");
4243 assert!(exported.contains("\"k\":\"v\""), "relation attribute map");
4244 assert!(exported.contains("\"ea\":\"ev\""), "entity attribute map");
4245 }
4246
4247 #[test]
4248 fn wipe_clears_relation_observations_attributes_and_fts() {
4249 let kg = new_kg();
4250 seed_relation_obs_attrs(&kg);
4251 kg.set_attributes(&[AttributeSet {
4252 owner_kind: "entity".into(),
4253 entity_name: Some("a".into()),
4254 from: None,
4255 to: None,
4256 relation_type: None,
4257 attributes: std::collections::BTreeMap::from([("ea".into(), "ev".into())]),
4258 }])
4259 .unwrap();
4260 let conn = kg.writer.lock();
4261 let obs_hits: i64 = conn
4262 .query_row(
4263 "SELECT COUNT(*) FROM rel_obs_fts WHERE rel_obs_fts MATCH 'contract'",
4264 [],
4265 |row| row.get(0),
4266 )
4267 .unwrap();
4268 assert_eq!(obs_hits, 1, "non-empty rel_obs_fts fixture");
4269 drop(conn);
4270
4271 kg.wipe().unwrap();
4272
4273 let conn = kg.writer.lock();
4274 for table in ["relation_observation", "attribute"] {
4275 let count: i64 = conn
4276 .query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |row| {
4277 row.get(0)
4278 })
4279 .unwrap();
4280 assert_eq!(count, 0, "{table} must be empty after wipe");
4281 }
4282 let fts_hits: i64 = conn
4283 .query_row(
4284 "SELECT COUNT(*) FROM rel_obs_fts WHERE rel_obs_fts MATCH 'contract'",
4285 [],
4286 |row| row.get(0),
4287 )
4288 .unwrap();
4289 drop(conn);
4290 assert_eq!(fts_hits, 0, "rel_obs_fts must not retain postings");
4291 }
4292}