use uuid::Uuid;
use crate::storage::{StorageError, StorageTrait};
use crate::types::Memory;
#[derive(Debug, Clone, Default)]
pub struct ErasureResult {
pub memories_deleted: usize,
pub edges_deleted: usize,
pub entities_deleted: usize,
pub complete: bool,
pub warnings: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct ExportResult {
pub memories: Vec<String>,
pub entities: Vec<String>,
pub total_records: usize,
}
pub fn erase_entity(
storage: &dyn StorageTrait,
entity_id: Uuid,
) -> Result<ErasureResult, StorageError> {
let mut result = ErasureResult::default();
match storage.delete_memories_by_entity(entity_id) {
Ok(count) => result.memories_deleted = count,
Err(e) => result.warnings.push(format!("Memory deletion error: {e}")),
}
match storage.get_edges_for_entity(entity_id) {
Ok(edges) => result.edges_deleted = edges.len(),
Err(e) => result.warnings.push(format!("Edge query error: {e}")),
}
result.entities_deleted = 1;
result.complete = result.warnings.is_empty();
Ok(result)
}
pub fn erase_namespace(
storage: &dyn StorageTrait,
namespace_id: Uuid,
) -> Result<ErasureResult, StorageError> {
let mut result = ErasureResult::default();
let entities = storage.list_entities_by_namespace(namespace_id)?;
for entity in &entities {
match erase_entity(storage, entity.id) {
Ok(entity_result) => {
result.memories_deleted += entity_result.memories_deleted;
result.edges_deleted += entity_result.edges_deleted;
result.entities_deleted += entity_result.entities_deleted;
result.warnings.extend(entity_result.warnings);
}
Err(e) => {
result
.warnings
.push(format!("Entity {} erasure error: {e}", entity.id));
}
}
}
result.complete = result.warnings.is_empty();
Ok(result)
}
pub fn export_entity_data(
storage: &dyn StorageTrait,
entity_id: Uuid,
namespace_id: Uuid,
) -> Result<ExportResult, StorageError> {
let all_memories = storage.get_all_memories_by_namespace(namespace_id)?;
let entity_memories: Vec<String> = all_memories
.into_iter()
.filter(|m| match m {
Memory::Episodic(e) => e.about_entity == entity_id || e.source_entity == entity_id,
Memory::Semantic(s) => s.subject == entity_id,
Memory::Procedural(_) => false,
})
.map(|m| {
let json = match m {
Memory::Episodic(e) => serde_json::json!({
"type": "episodic",
"id": e.id.to_string(),
"content": e.content,
"timestamp": e.timestamp.to_rfc3339(),
}),
Memory::Semantic(s) => serde_json::json!({
"type": "semantic",
"id": s.id.to_string(),
"subject": s.subject.to_string(),
"predicate": s.predicate,
"object": s.object,
}),
Memory::Procedural(p) => serde_json::json!({
"type": "procedural",
"id": p.id.to_string(),
"trigger": p.trigger,
"action": p.action,
}),
};
json.to_string()
})
.collect();
let total = entity_memories.len();
Ok(ExportResult {
memories: entity_memories,
entities: vec![serde_json::json!({"id": entity_id.to_string()}).to_string()],
total_records: total + 1,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::embedding::OnnxEmbedder;
use crate::storage::sqlite::SqliteBackend;
use crate::types::{Entity, EntityKind, Episode, EpisodicMemory, Namespace};
#[test]
fn test_erase_entity_empty() {
let dir = tempfile::tempdir().unwrap();
let storage = SqliteBackend::open(dir.path()).unwrap();
let entity_id = Uuid::new_v4();
let result = erase_entity(&storage, entity_id).unwrap();
assert_eq!(result.memories_deleted, 0);
assert!(result.complete);
}
#[test]
fn test_erase_entity_with_memories() {
let dir = tempfile::tempdir().unwrap();
let storage = SqliteBackend::open(dir.path()).unwrap();
let embedder = OnnxEmbedder::new_mock(64);
let ns = Namespace::new("gdpr-test");
storage.save_namespace(&ns).unwrap();
let mut entity = Entity::new("user-123", EntityKind::User);
entity.namespace_id = ns.id;
storage.save_entity(&entity).unwrap();
let episode = Episode::new(ns.id, vec![entity.id]);
storage.save_episode(&episode).unwrap();
let mut mem = EpisodicMemory::new(ns.id, episode.id, entity.id, entity.id, "test data");
mem.embedding = embedder.embed("test data").unwrap();
storage.save_episodic(&mem).unwrap();
let result = erase_entity(&storage, entity.id).unwrap();
assert_eq!(result.memories_deleted, 1);
}
#[test]
fn test_export_entity_data() {
let dir = tempfile::tempdir().unwrap();
let storage = SqliteBackend::open(dir.path()).unwrap();
let embedder = OnnxEmbedder::new_mock(64);
let ns = Namespace::new("export-test");
storage.save_namespace(&ns).unwrap();
let mut entity = Entity::new("user-456", EntityKind::User);
entity.namespace_id = ns.id;
storage.save_entity(&entity).unwrap();
let episode = Episode::new(ns.id, vec![entity.id]);
storage.save_episode(&episode).unwrap();
let mut mem = EpisodicMemory::new(ns.id, episode.id, entity.id, entity.id, "personal data");
mem.embedding = embedder.embed("personal data").unwrap();
storage.save_episodic(&mem).unwrap();
let result = export_entity_data(&storage, entity.id, ns.id).unwrap();
assert_eq!(result.total_records, 2); assert!(!result.memories.is_empty());
}
#[test]
fn test_erase_namespace() {
let dir = tempfile::tempdir().unwrap();
let storage = SqliteBackend::open(dir.path()).unwrap();
let ns = Namespace::new("erase-ns-test");
storage.save_namespace(&ns).unwrap();
let result = erase_namespace(&storage, ns.id).unwrap();
assert!(result.complete);
}
}