#![doc = include_str!("storage.md")]
use async_trait::async_trait;
pub trait StorageProviderPool: Send + Sync + 'static {}
#[derive(Default)]
pub struct StorageRegistry {
map: tokio::sync::Mutex<
std::collections::HashMap<
std::any::TypeId,
std::sync::Arc<dyn std::any::Any + Send + Sync>,
>,
>,
}
impl StorageRegistry {
pub async fn get_or_init<T: StorageProviderPool, F, Fut, E>(
&self,
init: F,
) -> Result<std::sync::Arc<T>, E>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<T, E>>,
{
let mut map = self.map.lock().await;
let type_id = std::any::TypeId::of::<T>();
if let Some(resource) = map.get(&type_id) {
return Ok(resource
.clone()
.downcast::<T>()
.unwrap_or_else(|_| unreachable!("TypeId mismatch in StorageRegistry")));
}
let resource = std::sync::Arc::new(init().await?);
map.insert(type_id, resource.clone());
Ok(resource)
}
}
pub trait StorageConfigResolver: Send + Sync + 'static {
fn resolve_config(
&self,
crate_name: &str,
storage_type_name: &str,
) -> Option<serde_json::Value>;
}
#[async_trait]
pub trait StorageConnection: Send + Sync + Sized + 'static {
type Config: serde::de::DeserializeOwned + Send + Sync;
async fn connect(
config: Self::Config,
storage_registry: std::sync::Arc<StorageRegistry>,
data_dir: &std::path::Path,
plugin_namespace: &str,
) -> Result<Self, String>;
}
use serde::{Deserialize, Serialize, de::DeserializeOwned};
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct EmptyStorageConfig {}
#[async_trait]
pub trait RecordStore: Send + Sync + 'static {
async fn upsert_record<T>(&self, collection: &str, key: &str, value: T) -> Result<(), String>
where
T: Serialize + Send + Sync + 'static;
async fn get_ordered_records<T>(
&self,
collection: &str,
limit: Option<usize>,
reverse: bool,
) -> Result<Vec<(String, T)>, String>
where
T: DeserializeOwned + Send + Sync + 'static;
async fn delete_record(&self, collection: &str, key: &str) -> Result<(), String>;
async fn trim_records_before(&self, collection: &str, cutoff_key: &str) -> Result<(), String>;
}
#[async_trait]
pub trait KeyValueStore: Send + Sync + 'static {
async fn set<T>(&self, collection: &str, key: &str, value: T) -> Result<(), String>
where
T: Serialize + Send + Sync + 'static;
async fn get<T>(&self, collection: &str, key: &str) -> Result<Option<T>, String>
where
T: DeserializeOwned + Send + Sync + 'static;
async fn delete(&self, collection: &str, key: &str) -> Result<(), String>;
async fn get_all<T>(&self, collection: &str) -> Result<Vec<T>, String>
where
T: DeserializeOwned + Send + Sync + 'static;
}
#[async_trait]
pub trait FileStore: Send + Sync + 'static {
async fn save_file(
&self,
collection: &str,
file_id: &str,
content: Vec<u8>,
) -> Result<(), String>;
async fn get_file(&self, collection: &str, file_id: &str) -> Result<Option<Vec<u8>>, String>;
async fn delete_file(&self, collection: &str, file_id: &str) -> Result<(), String>;
}
#[async_trait]
pub trait VectorStore: Send + Sync + 'static {
async fn setup_collection(&self, _collection: &str, _dimension: u32) -> Result<(), String> {
Ok(())
}
async fn insert_vectors<T>(&self, collection: &str, records: Vec<T>) -> Result<(), String>
where
T: Serialize + Send + Sync + 'static;
async fn search_vectors<T>(
&self,
collection: &str,
vector: Vec<f32>,
limit: u32,
) -> Result<Vec<T>, String>
where
T: DeserializeOwned + Send + Sync + 'static;
async fn delete_vectors(
&self,
collection: &str,
filter_field: &str,
filter_value: &str,
) -> Result<(), String>;
}