1use super::{
10 btree_index, Analyzer, Arc, BTreeMap, DocId, DocumentStore, InvertedIndex, KeyValueCatalog,
11 KeyValueDocumentStore, KeyValueHNSWIndex, KeyValueIVFIndex, KeyValueInvertedIndex,
12 KeyValueStore, KeyValueVectorIndex, PersistentStorageBackend, PersistentStorageIdentity,
13 StorageBackendResult, Value, VectorIndex, VectorIndexOpenMode, VectorIndexSpec,
14};
15use crate::{CatalogFacade, PersistentStorageSession};
16
17#[derive(Clone)]
19pub struct KeyValueStorageBackend {
20 store: Arc<dyn KeyValueStore>,
21}
22
23impl KeyValueStorageBackend {
24 pub fn new(store: Arc<dyn KeyValueStore>) -> Self {
25 Self { store }
26 }
27
28 pub fn store(&self) -> Arc<dyn KeyValueStore> {
29 Arc::clone(&self.store)
30 }
31}
32
33impl PersistentStorageBackend for KeyValueStorageBackend {
34 fn auxiliary_encryption_key(&self) -> Option<crate::StorageEncryptionKey> {
35 self.store.auxiliary_encryption_key()
36 }
37
38 fn retention_control(&self) -> Option<crate::read_control::StorageReadControl> {
39 self.store.retention_control()
40 }
41
42 fn notification_publications(
43 &self,
44 ) -> Option<&dyn crate::notifications::NotificationPublicationStore> {
45 self.store.notification_publications()
46 }
47
48 fn serializable_session(&self) -> Option<&dyn crate::mvcc::SerializableSession> {
49 self.store.serializable_session()
50 }
51
52 fn reclaim_obsolete(&self) -> StorageBackendResult<()> {
53 self.store.reclaim_obsolete()
54 }
55
56 fn vacuum(&self) -> StorageBackendResult<()> {
57 self.store.vacuum()
58 }
59
60 fn write_cancellation(&self) -> Option<uqa_core::CancellationToken> {
61 self.store.write_cancellation()
62 }
63
64 fn open_session_with_cancellation(
65 &self,
66 cancellation: &uqa_core::CancellationToken,
67 ) -> StorageBackendResult<PersistentStorageSession> {
68 let store = self.store.open_session_with_cancellation(cancellation)?;
69 let catalog: Arc<dyn CatalogFacade> = Arc::new(KeyValueCatalog::new(Arc::clone(&store)));
70 let backend: Arc<dyn PersistentStorageBackend> = Arc::new(Self::new(store));
71 Ok(PersistentStorageSession::new(catalog, backend))
72 }
73
74 fn open_controlled_session(
75 &self,
76 control: &crate::read_control::StorageReadControl,
77 ) -> StorageBackendResult<PersistentStorageSession> {
78 let store = self.store.open_controlled_session(control)?;
79 let catalog: Arc<dyn CatalogFacade> = Arc::new(KeyValueCatalog::new(Arc::clone(&store)));
80 let backend: Arc<dyn PersistentStorageBackend> = Arc::new(Self::new(store));
81 Ok(PersistentStorageSession::new(catalog, backend))
82 }
83
84 fn open_retained_read_session(
85 &self,
86 cancellation: &uqa_core::CancellationToken,
87 ) -> StorageBackendResult<PersistentStorageSession> {
88 let store = self.store.open_retained_read_session(cancellation)?;
89 let catalog: Arc<dyn CatalogFacade> = Arc::new(KeyValueCatalog::new(Arc::clone(&store)));
90 let backend: Arc<dyn PersistentStorageBackend> = Arc::new(Self::new(store));
91 Ok(PersistentStorageSession::new(catalog, backend))
92 }
93
94 fn transaction_model(&self) -> crate::StorageTransactionModel {
95 self.store.transaction_model()
96 }
97
98 fn supports_concurrent_pinned_read_and_write(&self) -> bool {
99 self.store.transaction_model().is_versioned()
100 }
101
102 fn identifier_allocator(&self) -> Option<&dyn crate::mvcc::IdentifierAllocator> {
103 self.store.identifier_allocator()
104 }
105
106 fn transaction_affinity(&self) -> Option<crate::StorageSessionAffinity> {
107 self.store.transaction_affinity()
108 }
109
110 fn storage_identity(&self) -> StorageBackendResult<Option<PersistentStorageIdentity>> {
111 self.store.storage_identity()
112 }
113
114 fn open_session(&self) -> StorageBackendResult<PersistentStorageSession> {
115 let store = self.store.open_session()?;
116 let catalog: Arc<dyn CatalogFacade> = Arc::new(KeyValueCatalog::new(Arc::clone(&store)));
117 let backend: Arc<dyn PersistentStorageBackend> = Arc::new(Self::new(store));
118 Ok(PersistentStorageSession::new(catalog, backend))
119 }
120
121 fn document_store(&self, table: &str) -> Box<dyn DocumentStore> {
122 Box::new(KeyValueDocumentStore::new(Arc::clone(&self.store), table))
123 }
124
125 fn migrate_document_storage(&self) -> StorageBackendResult<()> {
126 KeyValueDocumentStore::migrate_legacy_storage(self.store.as_ref())
127 }
128
129 fn inverted_index(&self, table: &str, analyzer: Analyzer) -> Box<dyn InvertedIndex> {
130 Box::new(KeyValueInvertedIndex::new(
131 Arc::clone(&self.store),
132 table,
133 analyzer,
134 ))
135 }
136
137 fn migrate_inverted_index_storage(&self) -> StorageBackendResult<()> {
138 KeyValueInvertedIndex::migrate_legacy_storage(self.store.as_ref())
139 }
140
141 fn vector_index(
142 &self,
143 table: &str,
144 field: &str,
145 dimensions: u32,
146 spec: VectorIndexSpec,
147 mode: VectorIndexOpenMode,
148 ) -> StorageBackendResult<Box<dyn VectorIndex>> {
149 match spec {
150 VectorIndexSpec::DiskANN(_) => Err(crate::StorageBackendError::Other(
151 "DiskANN requires an actual catalog binding".into(),
152 )),
153 VectorIndexSpec::BruteForce => Ok(Box::new(KeyValueVectorIndex::new(
154 Arc::clone(&self.store),
155 table,
156 field,
157 dimensions,
158 ))),
159 VectorIndexSpec::IVF(params) => match mode {
160 VectorIndexOpenMode::Create => Ok(Box::new(KeyValueIVFIndex::create(
161 Arc::clone(&self.store),
162 table,
163 field,
164 dimensions,
165 params,
166 )?)),
167 VectorIndexOpenMode::Restore => Ok(Box::new(KeyValueIVFIndex::restore(
168 Arc::clone(&self.store),
169 table,
170 field,
171 dimensions,
172 params,
173 )?)),
174 },
175 VectorIndexSpec::HNSW(params) => match mode {
176 VectorIndexOpenMode::Create => Ok(Box::new(KeyValueHNSWIndex::create(
177 Arc::clone(&self.store),
178 table,
179 field,
180 dimensions,
181 params,
182 )?)),
183 VectorIndexOpenMode::Restore => Ok(Box::new(KeyValueHNSWIndex::restore(
184 Arc::clone(&self.store),
185 table,
186 field,
187 dimensions,
188 params,
189 )?)),
190 },
191 }
192 }
193
194 fn diskann_index(
195 &self,
196 binding: crate::diskann_index::DiskANNIndexBinding<'_>,
197 options: crate::diskann_index::DiskANNIndexOptions,
198 temporary: &crate::diskann_index::build::DiskANNTemporaryBudget,
199 mode: VectorIndexOpenMode,
200 ) -> StorageBackendResult<Box<dyn VectorIndex>> {
201 let canonical = super::KeyValueDiskANNCanonical::new(
202 self.store.clone(),
203 binding.table,
204 binding.field,
205 binding.dimensions,
206 )?;
207 if mode == VectorIndexOpenMode::Create {
208 canonical.create_index(
209 binding.index,
210 &*binding.resolver,
211 options,
212 temporary,
213 binding.control,
214 )?;
215 }
216 let handle = canonical.bind(
217 binding.index.clone(),
218 binding.resolver,
219 options.read,
220 binding.control,
221 )?;
222 Ok(Box::new(crate::diskann_index::PersistentDiskANNIndex::new(
223 handle, options, temporary,
224 )?))
225 }
226
227 fn diskann_journal_pruner(
228 &self,
229 binding: crate::diskann_index::DiskANNIndexBinding<'_>,
230 max_record_bytes: usize,
231 ) -> StorageBackendResult<Box<dyn crate::diskann_index::changes::DiskANNJournalPruner>> {
232 super::KeyValueDiskANNCanonical::new(
233 self.store.clone(),
234 binding.table,
235 binding.field,
236 binding.dimensions,
237 )?
238 .journal_pruner(
239 binding.index,
240 binding.resolver,
241 max_record_bytes,
242 binding.control,
243 )
244 }
245
246 fn diskann_maintenance_source(
247 &self,
248 binding: crate::diskann_index::DiskANNIndexBinding<'_>,
249 max_record_bytes: usize,
250 ) -> StorageBackendResult<Box<dyn crate::diskann_index::maintenance::DiskANNMaintenanceSource>>
251 {
252 super::KeyValueDiskANNCanonical::new(
253 self.store.clone(),
254 binding.table,
255 binding.field,
256 binding.dimensions,
257 )?
258 .maintenance_source(
259 binding.index,
260 binding.resolver,
261 max_record_bytes,
262 binding.control,
263 )
264 }
265
266 fn retire_diskann_index(
267 &self,
268 binding: crate::diskann_index::DiskANNIndexBinding<'_>,
269 ) -> StorageBackendResult<()> {
270 super::KeyValueDiskANNCanonical::new(
271 self.store.clone(),
272 binding.table,
273 binding.field,
274 binding.dimensions,
275 )?
276 .retire_index(binding.index, &*binding.resolver, binding.control)
277 }
278
279 fn drop_vector_index_metadata(&self, table: &str, field: &str) -> StorageBackendResult<()> {
280 KeyValueIVFIndex::drop_metadata(self.store.as_ref(), table, field)
281 }
282
283 fn persists_btree_indexes(&self) -> bool {
284 true
285 }
286
287 fn load_btree_index(
288 &self,
289 table: &str,
290 field: &crate::ValueIndexKey,
291 ) -> StorageBackendResult<Option<Vec<(DocId, Value)>>> {
292 btree_index::load(self.store.as_ref(), table, field)
293 }
294
295 fn btree_index_fields(&self, table: &str) -> StorageBackendResult<Vec<crate::ValueIndexKey>> {
296 btree_index::fields(self.store.as_ref(), table)
297 }
298
299 fn read_btree_index_entry(
300 &self,
301 table: &str,
302 field: &crate::ValueIndexKey,
303 doc_id: DocId,
304 ) -> StorageBackendResult<crate::ValueIndexEntry> {
305 btree_index::read_entry(self.store.as_ref(), table, field, doc_id)
306 }
307
308 fn replace_btree_index(
309 &self,
310 table: &str,
311 field: &crate::ValueIndexKey,
312 values: &[(DocId, Value)],
313 ) -> StorageBackendResult<()> {
314 btree_index::replace(self.store.as_ref(), table, field, values)
315 }
316
317 fn replace_btree_indexes(
318 &self,
319 table: &str,
320 indexes: &[(&crate::ValueIndexKey, &[(DocId, Value)])],
321 ) -> StorageBackendResult<()> {
322 btree_index::replace_many(self.store.as_ref(), table, indexes)
323 }
324
325 fn apply_btree_index_write(
326 &self,
327 table: &str,
328 doc_id: DocId,
329 values: Option<&BTreeMap<crate::ValueIndexKey, Value>>,
330 ) -> StorageBackendResult<()> {
331 btree_index::apply_write(self.store.as_ref(), table, doc_id, values)
332 }
333
334 fn drop_btree_index(
335 &self,
336 table: &str,
337 field: &crate::ValueIndexKey,
338 ) -> StorageBackendResult<()> {
339 btree_index::drop_index(self.store.as_ref(), table, field)
340 }
341
342 fn clear_btree_indexes(&self, table: &str) -> StorageBackendResult<()> {
343 btree_index::clear_entries(self.store.as_ref(), table)
344 }
345
346 fn begin_transaction(&self) -> StorageBackendResult<()> {
347 self.store.begin_transaction()
348 }
349
350 fn begin_read_transaction(&self) -> StorageBackendResult<()> {
351 self.store.begin_upgradeable_transaction()
353 }
354
355 fn begin_upgradeable_transaction(&self) -> StorageBackendResult<()> {
356 self.store.begin_upgradeable_transaction()
357 }
358
359 fn refresh_transaction_snapshot(
360 &self,
361 cancellation: &uqa_core::CancellationToken,
362 ) -> StorageBackendResult<()> {
363 self.store.refresh_transaction_snapshot(cancellation)
364 }
365
366 fn in_transaction(&self) -> bool {
367 self.store.in_transaction()
368 }
369
370 fn transaction_has_written(&self) -> StorageBackendResult<bool> {
371 self.store.transaction_has_written()
372 }
373
374 fn change_version(&self) -> StorageBackendResult<Option<u64>> {
375 self.store.change_version()
376 }
377
378 fn read_view_revision(
379 &self,
380 ) -> StorageBackendResult<Option<crate::key_value::KeyValueReadRevision>> {
381 self.store.read_view_revision()
382 }
383
384 fn change_version_monitor_is_nonblocking(&self) -> StorageBackendResult<bool> {
385 self.store.change_version_monitor_is_nonblocking()
386 }
387
388 fn pin_transaction_snapshot(&self) -> StorageBackendResult<()> {
389 self.store.pin_transaction_snapshot()
390 }
391
392 fn commit_transaction(&self) -> StorageBackendResult<()> {
393 self.store.commit_transaction()
394 }
395
396 fn rollback_transaction(&self) -> StorageBackendResult<()> {
397 self.store.rollback_transaction()
398 }
399
400 fn savepoint(&self, id: crate::StorageSavepointId) -> StorageBackendResult<()> {
401 self.store.savepoint(&id.backend_name())
402 }
403
404 fn release_savepoint(&self, id: crate::StorageSavepointId) -> StorageBackendResult<()> {
405 self.store.release_savepoint(&id.backend_name())
406 }
407
408 fn rollback_to_savepoint(&self, id: crate::StorageSavepointId) -> StorageBackendResult<()> {
409 self.store.rollback_to_savepoint(&id.backend_name())
410 }
411}