diff --git a/src/retrieval/graph.rs b/src/retrieval/graph.rs index 27aa961..ed13376 100644 --- a/src/retrieval/graph.rs +++ b/src/retrieval/graph.rs @@ -1,7 +1,10 @@ use surrealdb::{engine::remote::ws::Client, Surreal}; use tracing::debug; -use crate::{error::ProcessingError, storage::types::knowledge_entity::KnowledgeEntity}; +use crate::{ + error::ProcessingError, + storage::types::{knowledge_entity::KnowledgeEntity, StoredObject}, +}; /// Retrieves database entries that match a specific source identifier. /// @@ -67,37 +70,38 @@ where Ok(matching_entities) } -pub async fn find_entities_by_relationship_by_source_ids( - db_client: &Surreal, - source_ids: &[String], -) -> Result, ProcessingError> { - let ids = source_ids - .iter() - // .map(|id| format!("`{}`", id)) - .map(|id| format!("knowledge_entity:`{}`", id)) - .collect::>() - .join(", "); +// pub async fn find_entities_by_relationship_by_source_ids( +// db_client: &Surreal, +// source_ids: &[String], +// ) -> Result, ProcessingError> { +// let ids = source_ids +// .iter() +// .map(|id| format!("knowledge_entity:`{}`", id)) +// .collect::>() +// .join(", "); - debug!("{:?}", ids); +// debug!("{:?}", ids); - let query = format!( - "SELECT *, <-> relates_to <-> knowledge_entity AS related FROM [{}]", - ids - ); +// let query = format!( +// "SELECT *, <-> relates_to <-> knowledge_entity AS related FROM [{}]", +// ids +// ); - debug!("{}", query); +// debug!("{}", query); - let result: Vec = db_client.query(query).await?.take(0)?; +// let result: Vec = db_client.query(query).await?.take(0)?; - Ok(result) -} +// Ok(result) +// } + +/// Find entities by their relationship to the id pub async fn find_entities_by_relationship_by_id( db_client: &Surreal, - source_id: String, + entity_id: String, ) -> Result, ProcessingError> { let query = format!( "SELECT *, <-> relates_to <-> knowledge_entity AS related FROM knowledge_entity:`{}`", - source_id + entity_id ); debug!("{}", query); @@ -106,3 +110,23 @@ pub async fn find_entities_by_relationship_by_id( Ok(result) } + +/// Get a specific KnowledgeEntity by its id +pub async fn get_entity_by_id( + db_client: &Surreal, + entity_id: &str, +) -> Result, ProcessingError> { + let response: Option = db_client + .select((KnowledgeEntity::table_name(), entity_id)) + .await?; + + Ok(response) +} + +pub async fn get_all_stored_items(db_client: &Surreal) -> Result, ProcessingError> +where + T: for<'de> StoredObject, +{ + let response: Vec = db_client.select(T::table_name()).await?; + Ok(response) +}