retrieval functions

This commit is contained in:
Per Stark
2024-11-27 12:38:50 +01:00
parent bb57844a68
commit b7c228e340

View File

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