1#![doc = include_str!("storage.md")]
2
3use async_trait::async_trait;
4
5pub trait StorageProviderPool: Send + Sync + 'static {}
7
8#[derive(Default)]
9pub struct StorageRegistry {
10 map: tokio::sync::Mutex<
11 std::collections::HashMap<
12 std::any::TypeId,
13 std::sync::Arc<dyn std::any::Any + Send + Sync>,
14 >,
15 >,
16}
17
18impl StorageRegistry {
19 pub async fn get_or_init<T: StorageProviderPool, F, Fut, E>(
23 &self,
24 init: F,
25 ) -> Result<std::sync::Arc<T>, E>
26 where
27 F: FnOnce() -> Fut,
28 Fut: std::future::Future<Output = Result<T, E>>,
29 {
30 let mut map = self.map.lock().await;
31 let type_id = std::any::TypeId::of::<T>();
32
33 if let Some(resource) = map.get(&type_id) {
34 return Ok(resource
37 .clone()
38 .downcast::<T>()
39 .unwrap_or_else(|_| unreachable!("TypeId mismatch in StorageRegistry")));
40 }
41
42 let resource = std::sync::Arc::new(init().await?);
43 map.insert(type_id, resource.clone());
44 Ok(resource)
45 }
46}
47
48pub trait StorageConfigResolver: Send + Sync + 'static {
49 fn resolve_config(
50 &self,
51 crate_name: &str,
52 storage_type_name: &str,
53 ) -> Option<serde_json::Value>;
54}
55
56#[async_trait]
60pub trait StorageConnection: Send + Sync + Sized + 'static {
61 type Config: serde::de::DeserializeOwned + Send + Sync;
62
63 async fn connect(
64 config: Self::Config,
65 storage_registry: std::sync::Arc<StorageRegistry>,
66 plugin_namespace: &str,
67 ) -> Result<Self, String>;
68}
69
70use serde::{Deserialize, Serialize, de::DeserializeOwned};
71
72#[derive(Debug, Clone, Serialize, Deserialize, Default)]
73pub struct EmptyStorageConfig {}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
76pub enum SortOrder {
77 #[default]
78 Ascending,
79 Descending,
80}
81
82#[async_trait]
83pub trait RecordStore: Send + Sync + 'static {
84 async fn upsert_record<T>(&self, collection: &str, key: &str, value: T) -> Result<(), String>
87 where
88 T: Serialize + Send + Sync + 'static;
89
90 async fn get_ordered_records<T>(
93 &self,
94 collection: &str,
95 limit: Option<usize>,
96 order: SortOrder,
97 ) -> Result<Vec<(String, T)>, String>
98 where
99 T: DeserializeOwned + Send + Sync + 'static;
100
101 async fn delete_record(&self, collection: &str, key: &str) -> Result<(), String>;
103
104 async fn trim_records_before(&self, collection: &str, cutoff_key: &str) -> Result<(), String>;
107}
108
109#[async_trait]
111pub trait KeyValueStore: Send + Sync + 'static {
112 async fn set<T>(&self, collection: &str, key: &str, value: T) -> Result<(), String>
113 where
114 T: Serialize + Send + Sync + 'static;
115
116 async fn get<T>(&self, collection: &str, key: &str) -> Result<Option<T>, String>
117 where
118 T: DeserializeOwned + Send + Sync + 'static;
119
120 async fn delete(&self, collection: &str, key: &str) -> Result<(), String>;
121
122 async fn get_all<T>(&self, collection: &str) -> Result<Vec<T>, String>
123 where
124 T: DeserializeOwned + Send + Sync + 'static;
125}
126
127#[async_trait]
129pub trait FileStore: Send + Sync + 'static {
130 async fn save_file(
132 &self,
133 collection: &str,
134 file_id: &str,
135 content: Vec<u8>,
136 ) -> Result<(), String>;
137
138 async fn get_file(&self, collection: &str, file_id: &str) -> Result<Option<Vec<u8>>, String>;
140
141 async fn delete_file(&self, collection: &str, file_id: &str) -> Result<(), String>;
143}
144#[async_trait]
145pub trait VectorStore: Send + Sync + 'static {
146 async fn setup_collection(&self, _collection: &str, _dimension: u32) -> Result<(), String> {
159 Ok(())
160 }
161
162 async fn insert_vectors<T>(&self, collection: &str, records: Vec<T>) -> Result<(), String>
163 where
164 T: Serialize + Send + Sync + 'static;
165
166 async fn search_vectors<T>(
167 &self,
168 collection: &str,
169 vector: Vec<f32>,
170 limit: u32,
171 ) -> Result<Vec<T>, String>
172 where
173 T: DeserializeOwned + Send + Sync + 'static;
174
175 async fn delete_vectors(
176 &self,
177 collection: &str,
178 filter_field: &str,
179 filter_value: &str,
180 ) -> Result<(), String>;
181}