1use async_trait::async_trait;
26use chrono::{DateTime, Utc};
27use serde::{de::DeserializeOwned, Serialize};
28use sled::{Batch as SledBatch, CompareAndSwapError, Db, IVec, Tree};
29use std::collections::HashMap;
30use std::path::Path;
31use std::sync::Arc;
32use tokio::sync::RwLock;
33use tracing::{debug, error, info, instrument, warn};
34
35use super::config::SledConfig;
36use super::error::{StorageError, StorageResult};
37use super::traits::{KeyValueStorage, Storage};
38use super::types::{
39 Batch, Operation, Query, StorageBackend, StorageEntry, StorageHealth, StorageMetadata,
40 StorageStats, Transaction, TransactionState,
41};
42
43const METADATA_TREE: &str = "__metadata__";
45const STATS_KEY: &str = "stats";
47const VERSION_COUNTER_KEY: &str = "version_counter";
49
50#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
52struct InternalEntry<T> {
53 key: String,
54 value: T,
55 version: u64,
56 created_at: DateTime<Utc>,
57 updated_at: DateTime<Utc>,
58 expires_at: Option<DateTime<Utc>>,
59 metadata: HashMap<String, String>,
60}
61
62impl<T> From<StorageEntry<T>> for InternalEntry<T> {
63 fn from(entry: StorageEntry<T>) -> Self {
64 Self {
65 key: entry.key,
66 value: entry.value,
67 version: entry.version,
68 created_at: entry.created_at,
69 updated_at: entry.updated_at,
70 expires_at: entry.expires_at,
71 metadata: entry.metadata,
72 }
73 }
74}
75
76impl<T> From<InternalEntry<T>> for StorageEntry<T> {
77 fn from(entry: InternalEntry<T>) -> Self {
78 Self {
79 key: entry.key,
80 value: entry.value,
81 version: entry.version,
82 created_at: entry.created_at,
83 updated_at: entry.updated_at,
84 expires_at: entry.expires_at,
85 metadata: entry.metadata,
86 }
87 }
88}
89
90#[derive(Debug)]
92pub struct SledStorage {
93 db: Arc<Db>,
95 config: SledConfig,
97 stats: Arc<RwLock<StorageStats>>,
99 trees: Arc<RwLock<HashMap<String, Tree>>>,
101 metadata_tree: Tree,
103 initialized: Arc<RwLock<bool>>,
105}
106
107impl SledStorage {
108 #[instrument(skip(config))]
110 pub async fn new(config: SledConfig) -> StorageResult<Self> {
111 info!("Initializing Sled storage at {:?}", config.path);
112
113 if config.create_dir {
115 if let Some(parent) = config.path.parent() {
116 tokio::fs::create_dir_all(parent).await.map_err(|e| {
117 StorageError::InvalidConfiguration {
118 message: format!("Failed to create directory: {}", e),
119 }
120 })?;
121 }
122 }
123
124 let mut sled_config = sled::Config::new().path(&config.path);
126
127 if config.use_compression {
128 sled_config = sled_config.use_compression(true);
129 }
130
131 if let Some(flush_ms) = config.flush_every_ms {
132 sled_config = sled_config.flush_every_ms(Some(flush_ms));
133 }
134
135 sled_config = sled_config.cache_capacity(config.cache_capacity_bytes);
136
137 if config.temporary {
138 sled_config = sled_config.temporary(true);
139 }
140
141 let db = tokio::task::spawn_blocking(move || sled_config.open())
143 .await
144 .map_err(|e| StorageError::InternalError {
145 message: format!("Task join error: {}", e),
146 })?
147 .map_err(|e| StorageError::ConnectionError {
148 source: Box::new(e),
149 })?;
150
151 let db = Arc::new(db);
152
153 let metadata_tree = db
155 .open_tree(METADATA_TREE)
156 .map_err(|e| StorageError::ConnectionError {
157 source: Box::new(e),
158 })?;
159
160 let stats = Arc::new(RwLock::new(StorageStats::new()));
162
163 Ok(Self {
164 db,
165 config,
166 stats,
167 trees: Arc::new(RwLock::new(HashMap::new())),
168 metadata_tree,
169 initialized: Arc::new(RwLock::new(false)),
170 })
171 }
172
173 #[instrument(skip(self))]
175 async fn get_tree(&self, table: &str) -> StorageResult<Tree> {
176 {
178 let trees = self.trees.read().await;
179 if let Some(tree) = trees.get(table) {
180 return Ok(tree.clone());
181 }
182 }
183
184 let db = self.db.clone();
186 let table_name = table.to_string();
187 let tree = tokio::task::spawn_blocking(move || db.open_tree(table_name))
188 .await
189 .map_err(|e| StorageError::InternalError {
190 message: format!("Task join error: {}", e),
191 })?
192 .map_err(|e| StorageError::ConnectionError {
193 source: Box::new(e),
194 })?;
195
196 let mut trees = self.trees.write().await;
198 trees.insert(table.to_string(), tree.clone());
199
200 debug!("Opened tree for table: {}", table);
201 Ok(tree)
202 }
203
204 fn serialize<T: Serialize>(value: &T) -> StorageResult<Vec<u8>> {
206 bincode::serialize(value).map_err(|e| StorageError::SerializationError {
207 source: Box::new(e),
208 })
209 }
210
211 fn deserialize<T: DeserializeOwned>(bytes: &[u8]) -> StorageResult<T> {
213 bincode::deserialize(bytes).map_err(|e| StorageError::DeserializationError {
214 source: Box::new(e),
215 })
216 }
217
218 async fn next_version(&self) -> StorageResult<u64> {
220 let metadata_tree = self.metadata_tree.clone();
221
222 tokio::task::spawn_blocking(move || {
223 let current = metadata_tree
224 .get(VERSION_COUNTER_KEY)
225 .map_err(|e| StorageError::InternalError {
226 message: format!("Failed to get version counter: {}", e),
227 })?
228 .map(|v| {
229 let bytes: [u8; 8] = v.as_ref().try_into().unwrap_or([0u8; 8]);
230 u64::from_le_bytes(bytes)
231 })
232 .unwrap_or(0);
233
234 let next = current + 1;
235 metadata_tree
236 .insert(VERSION_COUNTER_KEY, &next.to_le_bytes())
237 .map_err(|e| StorageError::InternalError {
238 message: format!("Failed to update version counter: {}", e),
239 })?;
240
241 Ok(next)
242 })
243 .await
244 .map_err(|e| StorageError::InternalError {
245 message: format!("Task join error: {}", e),
246 })?
247 }
248
249 async fn persist_stats(&self) -> StorageResult<()> {
251 let stats = self.stats.read().await.clone();
252 let stats_bytes = Self::serialize(&stats)?;
253 let metadata_tree = self.metadata_tree.clone();
254
255 tokio::task::spawn_blocking(move || {
256 metadata_tree
257 .insert(STATS_KEY, stats_bytes)
258 .map_err(|e| StorageError::InternalError {
259 message: format!("Failed to persist stats: {}", e),
260 })?;
261 Ok(())
262 })
263 .await
264 .map_err(|e| StorageError::InternalError {
265 message: format!("Task join error: {}", e),
266 })?
267 }
268
269 async fn load_stats(&self) -> StorageResult<()> {
271 let metadata_tree = self.metadata_tree.clone();
272
273 let stats_bytes = tokio::task::spawn_blocking(move || {
274 metadata_tree
275 .get(STATS_KEY)
276 .map_err(|e| StorageError::InternalError {
277 message: format!("Failed to load stats: {}", e),
278 })
279 })
280 .await
281 .map_err(|e| StorageError::InternalError {
282 message: format!("Task join error: {}", e),
283 })??;
284
285 if let Some(bytes) = stats_bytes {
286 let loaded_stats: StorageStats = Self::deserialize(&bytes)?;
287 let mut stats = self.stats.write().await;
288 *stats = loaded_stats;
289 debug!("Loaded statistics from metadata");
290 }
291
292 Ok(())
293 }
294
295 fn is_expired<T>(entry: &InternalEntry<T>) -> bool {
297 if let Some(expires_at) = entry.expires_at {
298 Utc::now() > expires_at
299 } else {
300 false
301 }
302 }
303
304 async fn record_operation(&self, success: bool, latency_ms: f64, bytes: u64) {
306 let mut stats = self.stats.write().await;
307 if success {
308 stats.record_success(latency_ms, bytes);
309 } else {
310 stats.record_failure();
311 }
312 }
313
314 async fn with_metrics<F, T>(&self, operation: F) -> StorageResult<T>
316 where
317 F: std::future::Future<Output = StorageResult<T>>,
318 {
319 let start = std::time::Instant::now();
320 let result = operation.await;
321 let latency_ms = start.elapsed().as_secs_f64() * 1000.0;
322
323 let success = result.is_ok();
324 self.record_operation(success, latency_ms, 0).await;
325
326 result
327 }
328
329 async fn cleanup_tree(&self, tree: &Tree) -> StorageResult<usize> {
331 let tree = tree.clone();
332
333 tokio::task::spawn_blocking(move || {
334 let mut removed = 0;
335 let mut to_remove = Vec::new();
336
337 for result in tree.iter() {
339 let (key, value) = result.map_err(|e| StorageError::InternalError {
340 message: format!("Iterator error: {}", e),
341 })?;
342
343 if let Ok(entry) = bincode::deserialize::<InternalEntry<serde_json::Value>>(&value)
345 {
346 if Self::is_expired(&entry) {
347 to_remove.push(key);
348 }
349 }
350 }
351
352 for key in to_remove {
354 tree.remove(&key).map_err(|e| StorageError::InternalError {
355 message: format!("Failed to remove key: {}", e),
356 })?;
357 removed += 1;
358 }
359
360 Ok(removed)
361 })
362 .await
363 .map_err(|e| StorageError::InternalError {
364 message: format!("Task join error: {}", e),
365 })?
366 }
367
368 fn apply_filters<T: DeserializeOwned>(
370 entries: Vec<StorageEntry<T>>,
371 query: &Query,
372 ) -> Vec<StorageEntry<T>> {
373 let mut filtered = entries;
374
375 if let Some(offset) = query.offset {
377 filtered = filtered.into_iter().skip(offset).collect();
378 }
379
380 if let Some(limit) = query.limit {
381 filtered.truncate(limit);
382 }
383
384 filtered
385 }
386}
387
388#[async_trait]
389impl Storage for SledStorage {
390 fn backend(&self) -> StorageBackend {
391 StorageBackend::Sled
392 }
393
394 async fn metadata(&self) -> StorageResult<StorageMetadata> {
395 let tree_names = self.db.tree_names();
396 let capabilities = vec![
397 "embedded".to_string(),
398 "key-value".to_string(),
399 "prefix-scan".to_string(),
400 "range-query".to_string(),
401 "compare-and-swap".to_string(),
402 "compression".to_string(),
403 "transactions".to_string(),
404 ];
405
406 Ok(StorageMetadata {
407 backend: StorageBackend::Sled,
408 version: "0.34.7".to_string(),
409 capabilities,
410 status: StorageHealth::Healthy,
411 last_health_check: Utc::now(),
412 })
413 }
414
415 #[instrument(skip(self))]
416 async fn health_check(&self) -> StorageResult<StorageHealth> {
417 let test_key = "__health_check__";
419 let test_value = Utc::now().timestamp();
420
421 match self.metadata_tree.insert(test_key, &test_value.to_le_bytes()) {
422 Ok(_) => {
423 self.metadata_tree.remove(test_key).ok();
424 Ok(StorageHealth::Healthy)
425 }
426 Err(e) => {
427 warn!("Health check failed: {}", e);
428 Ok(StorageHealth::Unhealthy)
429 }
430 }
431 }
432
433 async fn stats(&self) -> StorageResult<StorageStats> {
434 let mut stats = self.stats.read().await.clone();
435
436 stats.active_connections = self.trees.read().await.len();
438
439 Ok(stats)
440 }
441
442 #[instrument(skip(self, value))]
443 async fn insert<T: Serialize + Send + Sync>(
444 &self,
445 table: &str,
446 key: &str,
447 value: &T,
448 ) -> StorageResult<()> {
449 self.with_metrics(async {
450 let tree = self.get_tree(table).await?;
451 let version = self.next_version().await?;
452
453 let entry = InternalEntry {
454 key: key.to_string(),
455 value,
456 version,
457 created_at: Utc::now(),
458 updated_at: Utc::now(),
459 expires_at: None,
460 metadata: HashMap::new(),
461 };
462
463 let entry_bytes = Self::serialize(&entry)?;
464 let tree_clone = tree.clone();
465 let key_bytes = key.as_bytes().to_vec();
466
467 tokio::task::spawn_blocking(move || {
468 tree_clone
469 .insert(key_bytes, entry_bytes)
470 .map_err(|e| StorageError::InternalError {
471 message: format!("Insert failed: {}", e),
472 })?;
473 Ok(())
474 })
475 .await
476 .map_err(|e| StorageError::InternalError {
477 message: format!("Task join error: {}", e),
478 })?
479 })
480 .await
481 }
482
483 #[instrument(skip(self, value))]
484 async fn insert_with_ttl<T: Serialize + Send + Sync>(
485 &self,
486 table: &str,
487 key: &str,
488 value: &T,
489 ttl_seconds: i64,
490 ) -> StorageResult<()> {
491 self.with_metrics(async {
492 let tree = self.get_tree(table).await?;
493 let version = self.next_version().await?;
494 let now = Utc::now();
495
496 let entry = InternalEntry {
497 key: key.to_string(),
498 value,
499 version,
500 created_at: now,
501 updated_at: now,
502 expires_at: Some(now + chrono::Duration::seconds(ttl_seconds)),
503 metadata: HashMap::new(),
504 };
505
506 let entry_bytes = Self::serialize(&entry)?;
507 let tree_clone = tree.clone();
508 let key_bytes = key.as_bytes().to_vec();
509
510 tokio::task::spawn_blocking(move || {
511 tree_clone
512 .insert(key_bytes, entry_bytes)
513 .map_err(|e| StorageError::InternalError {
514 message: format!("Insert failed: {}", e),
515 })?;
516 Ok(())
517 })
518 .await
519 .map_err(|e| StorageError::InternalError {
520 message: format!("Task join error: {}", e),
521 })?
522 })
523 .await
524 }
525
526 #[instrument(skip(self))]
527 async fn get<T: DeserializeOwned + Serialize + Send + Sync>(
528 &self,
529 table: &str,
530 key: &str,
531 ) -> StorageResult<Option<StorageEntry<T>>> {
532 self.with_metrics(async {
533 let tree = self.get_tree(table).await?;
534 let key_bytes = key.as_bytes().to_vec();
535
536 let result = tokio::task::spawn_blocking(move || {
537 tree.get(key_bytes).map_err(|e| StorageError::InternalError {
538 message: format!("Get failed: {}", e),
539 })
540 })
541 .await
542 .map_err(|e| StorageError::InternalError {
543 message: format!("Task join error: {}", e),
544 })??;
545
546 if let Some(bytes) = result {
547 let entry: InternalEntry<T> = Self::deserialize(&bytes)?;
548
549 if Self::is_expired(&entry) {
551 let tree = self.get_tree(table).await?;
553 let key_bytes = key.as_bytes().to_vec();
554 tokio::task::spawn_blocking(move || tree.remove(key_bytes))
555 .await
556 .ok();
557 return Ok(None);
558 }
559
560 Ok(Some(entry.into()))
561 } else {
562 Ok(None)
563 }
564 })
565 .await
566 }
567
568 #[instrument(skip(self, value))]
569 async fn update<T: Serialize + Send + Sync>(
570 &self,
571 table: &str,
572 key: &str,
573 value: &T,
574 ) -> StorageResult<()> {
575 self.with_metrics(async {
576 let tree = self.get_tree(table).await?;
577 let key_bytes = key.as_bytes().to_vec();
578
579 let existing = tokio::task::spawn_blocking({
581 let tree = tree.clone();
582 let key_bytes = key_bytes.clone();
583 move || tree.get(key_bytes)
584 })
585 .await
586 .map_err(|e| StorageError::InternalError {
587 message: format!("Task join error: {}", e),
588 })?
589 .map_err(|e| StorageError::InternalError {
590 message: format!("Get failed: {}", e),
591 })?;
592
593 let (version, created_at) = if let Some(bytes) = existing {
594 let entry: InternalEntry<serde_json::Value> = Self::deserialize(&bytes)?;
595 (entry.version + 1, entry.created_at)
596 } else {
597 return Err(StorageError::NotFound {
598 entity: table.to_string(),
599 key: key.to_string(),
600 });
601 };
602
603 let entry = InternalEntry {
604 key: key.to_string(),
605 value,
606 version,
607 created_at,
608 updated_at: Utc::now(),
609 expires_at: None,
610 metadata: HashMap::new(),
611 };
612
613 let entry_bytes = Self::serialize(&entry)?;
614
615 tokio::task::spawn_blocking(move || {
616 tree.insert(key_bytes, entry_bytes)
617 .map_err(|e| StorageError::InternalError {
618 message: format!("Update failed: {}", e),
619 })?;
620 Ok(())
621 })
622 .await
623 .map_err(|e| StorageError::InternalError {
624 message: format!("Task join error: {}", e),
625 })?
626 })
627 .await
628 }
629
630 #[instrument(skip(self, value))]
631 async fn update_versioned<T: Serialize + Send + Sync>(
632 &self,
633 table: &str,
634 key: &str,
635 value: &T,
636 expected_version: u64,
637 ) -> StorageResult<()> {
638 self.with_metrics(async {
639 let tree = self.get_tree(table).await?;
640 let key_bytes = key.as_bytes().to_vec();
641
642 let existing = tokio::task::spawn_blocking({
644 let tree = tree.clone();
645 let key_bytes = key_bytes.clone();
646 move || tree.get(key_bytes)
647 })
648 .await
649 .map_err(|e| StorageError::InternalError {
650 message: format!("Task join error: {}", e),
651 })?
652 .map_err(|e| StorageError::InternalError {
653 message: format!("Get failed: {}", e),
654 })?;
655
656 let existing_bytes = existing.ok_or_else(|| StorageError::NotFound {
657 entity: table.to_string(),
658 key: key.to_string(),
659 })?;
660
661 let existing_entry: InternalEntry<serde_json::Value> =
662 Self::deserialize(&existing_bytes)?;
663
664 if existing_entry.version != expected_version {
666 return Err(StorageError::LockError {
667 message: format!(
668 "Version mismatch: expected {}, got {}",
669 expected_version, existing_entry.version
670 ),
671 });
672 }
673
674 let new_entry = InternalEntry {
675 key: key.to_string(),
676 value,
677 version: expected_version + 1,
678 created_at: existing_entry.created_at,
679 updated_at: Utc::now(),
680 expires_at: existing_entry.expires_at,
681 metadata: existing_entry.metadata,
682 };
683
684 let entry_bytes = Self::serialize(&new_entry)?;
685
686 tokio::task::spawn_blocking(move || {
688 match tree.compare_and_swap(key_bytes, Some(existing_bytes), Some(entry_bytes)) {
689 Ok(Ok(_)) => Ok(()),
690 Ok(Err(_)) => Err(StorageError::LockError {
691 message: "Concurrent modification detected".to_string(),
692 }),
693 Err(e) => Err(StorageError::InternalError {
694 message: format!("CAS failed: {}", e),
695 }),
696 }
697 })
698 .await
699 .map_err(|e| StorageError::InternalError {
700 message: format!("Task join error: {}", e),
701 })?
702 })
703 .await
704 }
705
706 #[instrument(skip(self))]
707 async fn delete(&self, table: &str, key: &str) -> StorageResult<()> {
708 self.with_metrics(async {
709 let tree = self.get_tree(table).await?;
710 let key_bytes = key.as_bytes().to_vec();
711
712 tokio::task::spawn_blocking(move || {
713 tree.remove(key_bytes)
714 .map_err(|e| StorageError::InternalError {
715 message: format!("Delete failed: {}", e),
716 })?;
717 Ok(())
718 })
719 .await
720 .map_err(|e| StorageError::InternalError {
721 message: format!("Task join error: {}", e),
722 })?
723 })
724 .await
725 }
726
727 #[instrument(skip(self))]
728 async fn exists(&self, table: &str, key: &str) -> StorageResult<bool> {
729 self.with_metrics(async {
730 let tree = self.get_tree(table).await?;
731 let key_bytes = key.as_bytes().to_vec();
732
733 tokio::task::spawn_blocking(move || {
734 tree.contains_key(key_bytes)
735 .map_err(|e| StorageError::InternalError {
736 message: format!("Exists check failed: {}", e),
737 })
738 })
739 .await
740 .map_err(|e| StorageError::InternalError {
741 message: format!("Task join error: {}", e),
742 })?
743 })
744 .await
745 }
746
747 #[instrument(skip(self, batch))]
748 async fn execute_batch(&self, batch: Batch) -> StorageResult<()> {
749 self.with_metrics(async {
750 for operation in batch.operations {
751 match operation {
752 Operation::Insert { table, data } => {
753 let entry: InternalEntry<serde_json::Value> = Self::deserialize(&data)?;
755 let tree = self.get_tree(&table).await?;
756 let key_bytes = entry.key.as_bytes().to_vec();
757
758 tokio::task::spawn_blocking(move || {
759 tree.insert(key_bytes, data)
760 .map_err(|e| StorageError::InternalError {
761 message: format!("Batch insert failed: {}", e),
762 })
763 })
764 .await
765 .map_err(|e| StorageError::InternalError {
766 message: format!("Task join error: {}", e),
767 })??;
768 }
769 Operation::Update { table, key, data } => {
770 let tree = self.get_tree(&table).await?;
771 let key_bytes = key.as_bytes().to_vec();
772
773 tokio::task::spawn_blocking(move || {
774 tree.insert(key_bytes, data)
775 .map_err(|e| StorageError::InternalError {
776 message: format!("Batch update failed: {}", e),
777 })
778 })
779 .await
780 .map_err(|e| StorageError::InternalError {
781 message: format!("Task join error: {}", e),
782 })??;
783 }
784 Operation::Delete { table, key } => {
785 let tree = self.get_tree(&table).await?;
786 let key_bytes = key.as_bytes().to_vec();
787
788 tokio::task::spawn_blocking(move || {
789 tree.remove(key_bytes)
790 .map_err(|e| StorageError::InternalError {
791 message: format!("Batch delete failed: {}", e),
792 })
793 })
794 .await
795 .map_err(|e| StorageError::InternalError {
796 message: format!("Task join error: {}", e),
797 })??;
798 }
799 }
800 }
801
802 Ok(())
803 })
804 .await
805 }
806
807 #[instrument(skip(self))]
808 async fn query<T: DeserializeOwned + Send + Sync + 'static>(
809 &self,
810 query: Query,
811 ) -> StorageResult<Vec<StorageEntry<T>>> {
812 self.with_metrics(async {
813 let tree = self.get_tree(&query.table).await?;
814 let tree_clone = tree.clone();
815
816 let entries = tokio::task::spawn_blocking(move || {
817 let mut results = Vec::new();
818
819 for result in tree_clone.iter() {
820 let (_key, value) = result.map_err(|e| StorageError::InternalError {
821 message: format!("Iterator error: {}", e),
822 })?;
823
824 if let Ok(entry) = Self::deserialize::<InternalEntry<T>>(&value) {
825 if !Self::is_expired(&entry) {
826 results.push(entry.into());
827 }
828 }
829 }
830
831 Ok::<_, StorageError>(results)
832 })
833 .await
834 .map_err(|e| StorageError::InternalError {
835 message: format!("Task join error: {}", e),
836 })??;
837
838 Ok(Self::apply_filters(entries, &query))
839 })
840 .await
841 }
842
843 #[instrument(skip(self))]
844 async fn count(&self, query: Query) -> StorageResult<usize> {
845 let entries: Vec<StorageEntry<serde_json::Value>> = self.query(query).await?;
846 Ok(entries.len())
847 }
848
849 async fn begin_transaction(&self) -> StorageResult<Transaction> {
850 Ok(Transaction::new())
851 }
852
853 #[instrument(skip(self, transaction))]
854 async fn commit_transaction(&self, mut transaction: Transaction) -> StorageResult<()> {
855 self.with_metrics(async {
856 for operation in &transaction.operations {
858 match operation {
859 Operation::Insert { table, data } => {
860 let entry: InternalEntry<serde_json::Value> = Self::deserialize(data)?;
861 let tree = self.get_tree(table).await?;
862 let key_bytes = entry.key.as_bytes().to_vec();
863 let data = data.clone();
864
865 tokio::task::spawn_blocking(move || {
866 tree.insert(key_bytes, data)
867 .map_err(|e| StorageError::TransactionError {
868 message: format!("Transaction insert failed: {}", e),
869 })
870 })
871 .await
872 .map_err(|e| StorageError::InternalError {
873 message: format!("Task join error: {}", e),
874 })??;
875 }
876 Operation::Update { table, key, data } => {
877 let tree = self.get_tree(table).await?;
878 let key_bytes = key.as_bytes().to_vec();
879 let data = data.clone();
880
881 tokio::task::spawn_blocking(move || {
882 tree.insert(key_bytes, data)
883 .map_err(|e| StorageError::TransactionError {
884 message: format!("Transaction update failed: {}", e),
885 })
886 })
887 .await
888 .map_err(|e| StorageError::InternalError {
889 message: format!("Task join error: {}", e),
890 })??;
891 }
892 Operation::Delete { table, key } => {
893 let tree = self.get_tree(table).await?;
894 let key_bytes = key.as_bytes().to_vec();
895
896 tokio::task::spawn_blocking(move || {
897 tree.remove(key_bytes)
898 .map_err(|e| StorageError::TransactionError {
899 message: format!("Transaction delete failed: {}", e),
900 })
901 })
902 .await
903 .map_err(|e| StorageError::InternalError {
904 message: format!("Task join error: {}", e),
905 })??;
906 }
907 }
908 }
909
910 transaction.state = TransactionState::Committed;
911 Ok(())
912 })
913 .await
914 }
915
916 async fn rollback_transaction(&self, mut transaction: Transaction) -> StorageResult<()> {
917 transaction.state = TransactionState::RolledBack;
918 Ok(())
919 }
920
921 #[instrument(skip(self))]
922 async fn scan_prefix(&self, table: &str, prefix: &str) -> StorageResult<Vec<String>> {
923 self.with_metrics(async {
924 let tree = self.get_tree(table).await?;
925 let prefix_bytes = prefix.as_bytes().to_vec();
926
927 tokio::task::spawn_blocking(move || {
928 let mut keys = Vec::new();
929
930 for result in tree.scan_prefix(prefix_bytes) {
931 let (key, _value) = result.map_err(|e| StorageError::InternalError {
932 message: format!("Scan prefix error: {}", e),
933 })?;
934
935 if let Ok(key_str) = String::from_utf8(key.to_vec()) {
936 keys.push(key_str);
937 }
938 }
939
940 Ok::<_, StorageError>(keys)
941 })
942 .await
943 .map_err(|e| StorageError::InternalError {
944 message: format!("Task join error: {}", e),
945 })?
946 })
947 .await
948 }
949
950 #[instrument(skip(self, keys))]
951 async fn get_multi<T: DeserializeOwned + Serialize + Send + Sync>(
952 &self,
953 table: &str,
954 keys: Vec<String>,
955 ) -> StorageResult<Vec<Option<StorageEntry<T>>>> {
956 self.with_metrics(async {
957 let mut results = Vec::new();
958
959 for key in keys {
960 let entry = self.get(table, &key).await?;
961 results.push(entry);
962 }
963
964 Ok(results)
965 })
966 .await
967 }
968
969 #[instrument(skip(self, keys))]
970 async fn delete_multi(&self, table: &str, keys: Vec<String>) -> StorageResult<()> {
971 self.with_metrics(async {
972 let tree = self.get_tree(table).await?;
973 let keys: Vec<Vec<u8>> = keys.into_iter().map(|k| k.as_bytes().to_vec()).collect();
974
975 tokio::task::spawn_blocking(move || {
976 for key in keys {
977 tree.remove(key).map_err(|e| StorageError::InternalError {
978 message: format!("Multi-delete failed: {}", e),
979 })?;
980 }
981 Ok(())
982 })
983 .await
984 .map_err(|e| StorageError::InternalError {
985 message: format!("Task join error: {}", e),
986 })?
987 })
988 .await
989 }
990
991 #[instrument(skip(self, new_value))]
992 async fn compare_and_swap<T: Serialize + Send + Sync>(
993 &self,
994 table: &str,
995 key: &str,
996 expected_version: u64,
997 new_value: &T,
998 ) -> StorageResult<bool> {
999 self.with_metrics(async {
1000 let tree = self.get_tree(table).await?;
1001 let key_bytes = key.as_bytes().to_vec();
1002
1003 let existing = tokio::task::spawn_blocking({
1005 let tree = tree.clone();
1006 let key_bytes = key_bytes.clone();
1007 move || tree.get(key_bytes)
1008 })
1009 .await
1010 .map_err(|e| StorageError::InternalError {
1011 message: format!("Task join error: {}", e),
1012 })?
1013 .map_err(|e| StorageError::InternalError {
1014 message: format!("Get failed: {}", e),
1015 })?;
1016
1017 if let Some(existing_bytes) = existing {
1018 let existing_entry: InternalEntry<serde_json::Value> =
1019 Self::deserialize(&existing_bytes)?;
1020
1021 if existing_entry.version != expected_version {
1023 return Ok(false);
1024 }
1025
1026 let new_entry = InternalEntry {
1027 key: key.to_string(),
1028 value: new_value,
1029 version: expected_version + 1,
1030 created_at: existing_entry.created_at,
1031 updated_at: Utc::now(),
1032 expires_at: existing_entry.expires_at,
1033 metadata: existing_entry.metadata,
1034 };
1035
1036 let new_bytes = Self::serialize(&new_entry)?;
1037
1038 let result = tokio::task::spawn_blocking(move || {
1039 tree.compare_and_swap(key_bytes, Some(existing_bytes), Some(new_bytes))
1040 })
1041 .await
1042 .map_err(|e| StorageError::InternalError {
1043 message: format!("Task join error: {}", e),
1044 })?
1045 .map_err(|e| StorageError::InternalError {
1046 message: format!("CAS failed: {}", e),
1047 })?;
1048
1049 Ok(result.is_ok())
1050 } else {
1051 Ok(false)
1052 }
1053 })
1054 .await
1055 }
1056
1057 #[instrument(skip(self))]
1058 async fn get_modified_since<T: DeserializeOwned + Send + Sync + 'static>(
1059 &self,
1060 table: &str,
1061 since: DateTime<Utc>,
1062 ) -> StorageResult<Vec<StorageEntry<T>>> {
1063 self.with_metrics(async {
1064 let tree = self.get_tree(table).await?;
1065
1066 let entries = tokio::task::spawn_blocking(move || {
1067 let mut results = Vec::new();
1068
1069 for result in tree.iter() {
1070 let (_key, value) = result.map_err(|e| StorageError::InternalError {
1071 message: format!("Iterator error: {}", e),
1072 })?;
1073
1074 if let Ok(entry) = Self::deserialize::<InternalEntry<T>>(&value) {
1075 if entry.updated_at > since && !Self::is_expired(&entry) {
1076 results.push(entry.into());
1077 }
1078 }
1079 }
1080
1081 Ok::<_, StorageError>(results)
1082 })
1083 .await
1084 .map_err(|e| StorageError::InternalError {
1085 message: format!("Task join error: {}", e),
1086 })??;
1087
1088 Ok(entries)
1089 })
1090 .await
1091 }
1092
1093 #[instrument(skip(self))]
1094 async fn get_expiring_before<T: DeserializeOwned + Send + Sync + 'static>(
1095 &self,
1096 table: &str,
1097 before: DateTime<Utc>,
1098 ) -> StorageResult<Vec<StorageEntry<T>>> {
1099 self.with_metrics(async {
1100 let tree = self.get_tree(table).await?;
1101
1102 let entries = tokio::task::spawn_blocking(move || {
1103 let mut results = Vec::new();
1104
1105 for result in tree.iter() {
1106 let (_key, value) = result.map_err(|e| StorageError::InternalError {
1107 message: format!("Iterator error: {}", e),
1108 })?;
1109
1110 if let Ok(entry) = Self::deserialize::<InternalEntry<T>>(&value) {
1111 if let Some(expires_at) = entry.expires_at {
1112 if expires_at < before {
1113 results.push(entry.into());
1114 }
1115 }
1116 }
1117 }
1118
1119 Ok::<_, StorageError>(results)
1120 })
1121 .await
1122 .map_err(|e| StorageError::InternalError {
1123 message: format!("Task join error: {}", e),
1124 })??;
1125
1126 Ok(entries)
1127 })
1128 .await
1129 }
1130
1131 #[instrument(skip(self))]
1132 async fn cleanup_expired(&self, table: &str) -> StorageResult<usize> {
1133 self.with_metrics(async {
1134 let tree = self.get_tree(table).await?;
1135 self.cleanup_tree(&tree).await
1136 })
1137 .await
1138 }
1139
1140 #[instrument(skip(self))]
1141 async fn compact(&self) -> StorageResult<()> {
1142 self.with_metrics(async {
1143 let db = self.db.clone();
1144
1145 info!("Starting database compaction");
1146
1147 let tree_names: Vec<_> = db.tree_names().into_iter().collect();
1149
1150 for name in tree_names {
1152 let name_str = String::from_utf8_lossy(&name).to_string();
1153 debug!("Compacting tree: {}", name_str);
1154
1155 let tree = db.open_tree(&name).map_err(|e| StorageError::InternalError {
1156 message: format!("Failed to open tree for compaction: {}", e),
1157 })?;
1158
1159 let tree_clone = tree.clone();
1161 tokio::task::spawn_blocking(move || {
1162 for _entry in tree_clone.iter() {
1163 }
1165 })
1166 .await
1167 .map_err(|e| StorageError::InternalError {
1168 message: format!("Task join error: {}", e),
1169 })?;
1170 }
1171
1172 info!("Database compaction completed");
1173 Ok(())
1174 })
1175 .await
1176 }
1177
1178 #[instrument(skip(self))]
1179 async fn flush(&self) -> StorageResult<()> {
1180 self.with_metrics(async {
1181 let db = self.db.clone();
1182
1183 tokio::task::spawn_blocking(move || {
1184 db.flush()
1185 .map_err(|e| StorageError::InternalError {
1186 message: format!("Flush failed: {}", e),
1187 })?;
1188 Ok(())
1189 })
1190 .await
1191 .map_err(|e| StorageError::InternalError {
1192 message: format!("Task join error: {}", e),
1193 })?
1194 })
1195 .await
1196 }
1197
1198 #[instrument(skip(self))]
1199 async fn initialize(&mut self) -> StorageResult<()> {
1200 let mut initialized = self.initialized.write().await;
1201 if *initialized {
1202 return Ok(());
1203 }
1204
1205 info!("Initializing Sled storage");
1206
1207 self.load_stats().await?;
1209
1210 self.health_check().await?;
1212
1213 *initialized = true;
1214 info!("Sled storage initialized successfully");
1215
1216 Ok(())
1217 }
1218
1219 #[instrument(skip(self))]
1220 async fn shutdown(&mut self) -> StorageResult<()> {
1221 let mut initialized = self.initialized.write().await;
1222 if !*initialized {
1223 return Ok(());
1224 }
1225
1226 info!("Shutting down Sled storage");
1227
1228 self.persist_stats().await?;
1230
1231 self.flush().await?;
1233
1234 *initialized = false;
1235 info!("Sled storage shutdown complete");
1236
1237 Ok(())
1238 }
1239}
1240
1241#[async_trait]
1242impl KeyValueStorage for SledStorage {
1243 #[instrument(skip(self))]
1244 async fn get_bytes(&self, table: &str, key: &str) -> StorageResult<Option<Vec<u8>>> {
1245 self.with_metrics(async {
1246 let tree = self.get_tree(table).await?;
1247 let key_bytes = key.as_bytes().to_vec();
1248
1249 tokio::task::spawn_blocking(move || {
1250 tree.get(key_bytes)
1251 .map_err(|e| StorageError::InternalError {
1252 message: format!("Get bytes failed: {}", e),
1253 })
1254 .map(|opt| opt.map(|v| v.to_vec()))
1255 })
1256 .await
1257 .map_err(|e| StorageError::InternalError {
1258 message: format!("Task join error: {}", e),
1259 })?
1260 })
1261 .await
1262 }
1263
1264 #[instrument(skip(self, value))]
1265 async fn set_bytes(&self, table: &str, key: &str, value: Vec<u8>) -> StorageResult<()> {
1266 self.with_metrics(async {
1267 let tree = self.get_tree(table).await?;
1268 let key_bytes = key.as_bytes().to_vec();
1269
1270 tokio::task::spawn_blocking(move || {
1271 tree.insert(key_bytes, value)
1272 .map_err(|e| StorageError::InternalError {
1273 message: format!("Set bytes failed: {}", e),
1274 })?;
1275 Ok(())
1276 })
1277 .await
1278 .map_err(|e| StorageError::InternalError {
1279 message: format!("Task join error: {}", e),
1280 })?
1281 })
1282 .await
1283 }
1284
1285 #[instrument(skip(self))]
1286 async fn iter_prefix(
1287 &self,
1288 table: &str,
1289 prefix: &str,
1290 ) -> StorageResult<Vec<(String, Vec<u8>)>> {
1291 self.with_metrics(async {
1292 let tree = self.get_tree(table).await?;
1293 let prefix_bytes = prefix.as_bytes().to_vec();
1294
1295 tokio::task::spawn_blocking(move || {
1296 let mut results = Vec::new();
1297
1298 for result in tree.scan_prefix(prefix_bytes) {
1299 let (key, value) = result.map_err(|e| StorageError::InternalError {
1300 message: format!("Iter prefix error: {}", e),
1301 })?;
1302
1303 if let Ok(key_str) = String::from_utf8(key.to_vec()) {
1304 results.push((key_str, value.to_vec()));
1305 }
1306 }
1307
1308 Ok::<_, StorageError>(results)
1309 })
1310 .await
1311 .map_err(|e| StorageError::InternalError {
1312 message: format!("Task join error: {}", e),
1313 })?
1314 })
1315 .await
1316 }
1317
1318 #[instrument(skip(self))]
1319 async fn get_range(
1320 &self,
1321 table: &str,
1322 start: &str,
1323 end: &str,
1324 ) -> StorageResult<Vec<(String, Vec<u8>)>> {
1325 self.with_metrics(async {
1326 let tree = self.get_tree(table).await?;
1327 let start_bytes = start.as_bytes().to_vec();
1328 let end_bytes = end.as_bytes().to_vec();
1329
1330 tokio::task::spawn_blocking(move || {
1331 let mut results = Vec::new();
1332
1333 for result in tree.range(start_bytes..end_bytes) {
1334 let (key, value) = result.map_err(|e| StorageError::InternalError {
1335 message: format!("Range query error: {}", e),
1336 })?;
1337
1338 if let Ok(key_str) = String::from_utf8(key.to_vec()) {
1339 results.push((key_str, value.to_vec()));
1340 }
1341 }
1342
1343 Ok::<_, StorageError>(results)
1344 })
1345 .await
1346 .map_err(|e| StorageError::InternalError {
1347 message: format!("Task join error: {}", e),
1348 })?
1349 })
1350 .await
1351 }
1352
1353 #[instrument(skip(self, value))]
1354 async fn merge(&self, table: &str, key: &str, value: Vec<u8>) -> StorageResult<()> {
1355 self.with_metrics(async {
1356 let tree = self.get_tree(table).await?;
1357 let key_bytes = key.as_bytes().to_vec();
1358
1359 tokio::task::spawn_blocking(move || {
1360 tree.merge(key_bytes, value)
1361 .map_err(|e| StorageError::InternalError {
1362 message: format!("Merge failed: {}", e),
1363 })?;
1364 Ok(())
1365 })
1366 .await
1367 .map_err(|e| StorageError::InternalError {
1368 message: format!("Task join error: {}", e),
1369 })?
1370 })
1371 .await
1372 }
1373}
1374
1375#[cfg(test)]
1376mod tests {
1377 use super::*;
1378 use serde::{Deserialize, Serialize};
1379 use tempfile::TempDir;
1380
1381 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1382 struct TestData {
1383 id: String,
1384 name: String,
1385 value: i32,
1386 }
1387
1388 async fn create_test_storage() -> (SledStorage, TempDir) {
1389 let temp_dir = TempDir::new().unwrap();
1390 let config = SledConfig {
1391 path: temp_dir.path().to_path_buf(),
1392 cache_capacity_bytes: 1024 * 1024,
1393 use_compression: true,
1394 flush_every_ms: Some(100),
1395 mode: super::super::config::SledMode::HighThroughput,
1396 temporary: false,
1397 create_dir: true,
1398 };
1399
1400 let mut storage = SledStorage::new(config).await.unwrap();
1401 storage.initialize().await.unwrap();
1402 (storage, temp_dir)
1403 }
1404
1405 #[tokio::test]
1406 async fn test_insert_and_get() {
1407 let (storage, _temp) = create_test_storage().await;
1408 let data = TestData {
1409 id: "1".to_string(),
1410 name: "Test".to_string(),
1411 value: 42,
1412 };
1413
1414 storage.insert("test_table", "key1", &data).await.unwrap();
1415
1416 let retrieved: Option<StorageEntry<TestData>> =
1417 storage.get("test_table", "key1").await.unwrap();
1418
1419 assert!(retrieved.is_some());
1420 let entry = retrieved.unwrap();
1421 assert_eq!(entry.value, data);
1422 assert_eq!(entry.version, 1);
1423 }
1424
1425 #[tokio::test]
1426 async fn test_insert_with_ttl() {
1427 let (storage, _temp) = create_test_storage().await;
1428 let data = TestData {
1429 id: "1".to_string(),
1430 name: "Test".to_string(),
1431 value: 42,
1432 };
1433
1434 storage
1435 .insert_with_ttl("test_table", "key1", &data, 1)
1436 .await
1437 .unwrap();
1438
1439 let retrieved: Option<StorageEntry<TestData>> =
1441 storage.get("test_table", "key1").await.unwrap();
1442 assert!(retrieved.is_some());
1443
1444 tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
1446
1447 let expired: Option<StorageEntry<TestData>> =
1449 storage.get("test_table", "key1").await.unwrap();
1450 assert!(expired.is_none());
1451 }
1452
1453 #[tokio::test]
1454 async fn test_update() {
1455 let (storage, _temp) = create_test_storage().await;
1456 let data = TestData {
1457 id: "1".to_string(),
1458 name: "Test".to_string(),
1459 value: 42,
1460 };
1461
1462 storage.insert("test_table", "key1", &data).await.unwrap();
1463
1464 let updated_data = TestData {
1465 id: "1".to_string(),
1466 name: "Updated".to_string(),
1467 value: 100,
1468 };
1469
1470 storage
1471 .update("test_table", "key1", &updated_data)
1472 .await
1473 .unwrap();
1474
1475 let retrieved: Option<StorageEntry<TestData>> =
1476 storage.get("test_table", "key1").await.unwrap();
1477
1478 assert!(retrieved.is_some());
1479 let entry = retrieved.unwrap();
1480 assert_eq!(entry.value.name, "Updated");
1481 assert_eq!(entry.value.value, 100);
1482 assert_eq!(entry.version, 2);
1483 }
1484
1485 #[tokio::test]
1486 async fn test_update_versioned() {
1487 let (storage, _temp) = create_test_storage().await;
1488 let data = TestData {
1489 id: "1".to_string(),
1490 name: "Test".to_string(),
1491 value: 42,
1492 };
1493
1494 storage.insert("test_table", "key1", &data).await.unwrap();
1495
1496 let updated_data = TestData {
1497 id: "1".to_string(),
1498 name: "Updated".to_string(),
1499 value: 100,
1500 };
1501
1502 storage
1504 .update_versioned("test_table", "key1", &updated_data, 1)
1505 .await
1506 .unwrap();
1507
1508 let result = storage
1510 .update_versioned("test_table", "key1", &updated_data, 1)
1511 .await;
1512 assert!(result.is_err());
1513 }
1514
1515 #[tokio::test]
1516 async fn test_delete() {
1517 let (storage, _temp) = create_test_storage().await;
1518 let data = TestData {
1519 id: "1".to_string(),
1520 name: "Test".to_string(),
1521 value: 42,
1522 };
1523
1524 storage.insert("test_table", "key1", &data).await.unwrap();
1525 storage.delete("test_table", "key1").await.unwrap();
1526
1527 let retrieved: Option<StorageEntry<TestData>> =
1528 storage.get("test_table", "key1").await.unwrap();
1529 assert!(retrieved.is_none());
1530 }
1531
1532 #[tokio::test]
1533 async fn test_exists() {
1534 let (storage, _temp) = create_test_storage().await;
1535 let data = TestData {
1536 id: "1".to_string(),
1537 name: "Test".to_string(),
1538 value: 42,
1539 };
1540
1541 assert!(!storage.exists("test_table", "key1").await.unwrap());
1542
1543 storage.insert("test_table", "key1", &data).await.unwrap();
1544 assert!(storage.exists("test_table", "key1").await.unwrap());
1545
1546 storage.delete("test_table", "key1").await.unwrap();
1547 assert!(!storage.exists("test_table", "key1").await.unwrap());
1548 }
1549
1550 #[tokio::test]
1551 async fn test_scan_prefix() {
1552 let (storage, _temp) = create_test_storage().await;
1553
1554 for i in 0..5 {
1555 let data = TestData {
1556 id: i.to_string(),
1557 name: format!("Test {}", i),
1558 value: i,
1559 };
1560 storage
1561 .insert("test_table", &format!("prefix_{}", i), &data)
1562 .await
1563 .unwrap();
1564 }
1565
1566 let keys = storage.scan_prefix("test_table", "prefix_").await.unwrap();
1567 assert_eq!(keys.len(), 5);
1568 }
1569
1570 #[tokio::test]
1571 async fn test_get_multi() {
1572 let (storage, _temp) = create_test_storage().await;
1573
1574 for i in 0..3 {
1575 let data = TestData {
1576 id: i.to_string(),
1577 name: format!("Test {}", i),
1578 value: i,
1579 };
1580 storage
1581 .insert("test_table", &format!("key{}", i), &data)
1582 .await
1583 .unwrap();
1584 }
1585
1586 let keys = vec!["key0".to_string(), "key1".to_string(), "key2".to_string()];
1587 let results: Vec<Option<StorageEntry<TestData>>> =
1588 storage.get_multi("test_table", keys).await.unwrap();
1589
1590 assert_eq!(results.len(), 3);
1591 assert!(results.iter().all(|r| r.is_some()));
1592 }
1593
1594 #[tokio::test]
1595 async fn test_delete_multi() {
1596 let (storage, _temp) = create_test_storage().await;
1597
1598 for i in 0..3 {
1599 let data = TestData {
1600 id: i.to_string(),
1601 name: format!("Test {}", i),
1602 value: i,
1603 };
1604 storage
1605 .insert("test_table", &format!("key{}", i), &data)
1606 .await
1607 .unwrap();
1608 }
1609
1610 let keys = vec!["key0".to_string(), "key1".to_string(), "key2".to_string()];
1611 storage.delete_multi("test_table", keys).await.unwrap();
1612
1613 for i in 0..3 {
1614 let exists = storage
1615 .exists("test_table", &format!("key{}", i))
1616 .await
1617 .unwrap();
1618 assert!(!exists);
1619 }
1620 }
1621
1622 #[tokio::test]
1623 async fn test_compare_and_swap() {
1624 let (storage, _temp) = create_test_storage().await;
1625 let data = TestData {
1626 id: "1".to_string(),
1627 name: "Test".to_string(),
1628 value: 42,
1629 };
1630
1631 storage.insert("test_table", "key1", &data).await.unwrap();
1632
1633 let new_data = TestData {
1634 id: "1".to_string(),
1635 name: "Updated".to_string(),
1636 value: 100,
1637 };
1638
1639 let result = storage
1641 .compare_and_swap("test_table", "key1", 1, &new_data)
1642 .await
1643 .unwrap();
1644 assert!(result);
1645
1646 let result = storage
1648 .compare_and_swap("test_table", "key1", 1, &new_data)
1649 .await
1650 .unwrap();
1651 assert!(!result);
1652 }
1653
1654 #[tokio::test]
1655 async fn test_batch_operations() {
1656 let (storage, _temp) = create_test_storage().await;
1657
1658 let data1 = TestData {
1659 id: "1".to_string(),
1660 name: "Test1".to_string(),
1661 value: 1,
1662 };
1663
1664 let data2 = TestData {
1665 id: "2".to_string(),
1666 name: "Test2".to_string(),
1667 value: 2,
1668 };
1669
1670 storage.insert("test_table", "key1", &data1).await.unwrap();
1672
1673 let mut batch = Batch::new(10);
1674
1675 let entry1 = InternalEntry {
1677 key: "key2".to_string(),
1678 value: data2.clone(),
1679 version: 1,
1680 created_at: Utc::now(),
1681 updated_at: Utc::now(),
1682 expires_at: None,
1683 metadata: HashMap::new(),
1684 };
1685
1686 batch
1687 .add(Operation::Insert {
1688 table: "test_table".to_string(),
1689 data: SledStorage::serialize(&entry1).unwrap(),
1690 })
1691 .unwrap();
1692
1693 batch
1694 .add(Operation::Delete {
1695 table: "test_table".to_string(),
1696 key: "key1".to_string(),
1697 })
1698 .unwrap();
1699
1700 storage.execute_batch(batch).await.unwrap();
1701
1702 assert!(storage.exists("test_table", "key2").await.unwrap());
1704 assert!(!storage.exists("test_table", "key1").await.unwrap());
1706 }
1707
1708 #[tokio::test]
1709 async fn test_query() {
1710 let (storage, _temp) = create_test_storage().await;
1711
1712 for i in 0..10 {
1713 let data = TestData {
1714 id: i.to_string(),
1715 name: format!("Test {}", i),
1716 value: i,
1717 };
1718 storage
1719 .insert("test_table", &format!("key{}", i), &data)
1720 .await
1721 .unwrap();
1722 }
1723
1724 let query = Query::new("test_table").limit(5);
1725 let results: Vec<StorageEntry<TestData>> = storage.query(query).await.unwrap();
1726
1727 assert_eq!(results.len(), 5);
1728 }
1729
1730 #[tokio::test]
1731 async fn test_count() {
1732 let (storage, _temp) = create_test_storage().await;
1733
1734 for i in 0..10 {
1735 let data = TestData {
1736 id: i.to_string(),
1737 name: format!("Test {}", i),
1738 value: i,
1739 };
1740 storage
1741 .insert("test_table", &format!("key{}", i), &data)
1742 .await
1743 .unwrap();
1744 }
1745
1746 let query = Query::new("test_table");
1747 let count = storage.count(query).await.unwrap();
1748
1749 assert_eq!(count, 10);
1750 }
1751
1752 #[tokio::test]
1753 async fn test_health_check() {
1754 let (storage, _temp) = create_test_storage().await;
1755 let health = storage.health_check().await.unwrap();
1756 assert_eq!(health, StorageHealth::Healthy);
1757 }
1758
1759 #[tokio::test]
1760 async fn test_metadata() {
1761 let (storage, _temp) = create_test_storage().await;
1762 let metadata = storage.metadata().await.unwrap();
1763 assert_eq!(metadata.backend, StorageBackend::Sled);
1764 assert!(!metadata.capabilities.is_empty());
1765 }
1766
1767 #[tokio::test]
1768 async fn test_stats() {
1769 let (storage, _temp) = create_test_storage().await;
1770 let stats = storage.stats().await.unwrap();
1771 assert_eq!(stats.total_operations, 0);
1772 }
1773
1774 #[tokio::test]
1775 async fn test_get_bytes_and_set_bytes() {
1776 let (storage, _temp) = create_test_storage().await;
1777 let data = vec![1, 2, 3, 4, 5];
1778
1779 storage
1780 .set_bytes("test_table", "key1", data.clone())
1781 .await
1782 .unwrap();
1783
1784 let retrieved = storage.get_bytes("test_table", "key1").await.unwrap();
1785 assert_eq!(retrieved, Some(data));
1786 }
1787
1788 #[tokio::test]
1789 async fn test_iter_prefix_kv() {
1790 let (storage, _temp) = create_test_storage().await;
1791
1792 for i in 0..5 {
1793 storage
1794 .set_bytes("test_table", &format!("prefix_{}", i), vec![i as u8])
1795 .await
1796 .unwrap();
1797 }
1798
1799 let results = storage.iter_prefix("test_table", "prefix_").await.unwrap();
1800 assert_eq!(results.len(), 5);
1801 }
1802
1803 #[tokio::test]
1804 async fn test_get_range() {
1805 let (storage, _temp) = create_test_storage().await;
1806
1807 for i in 0..10 {
1808 storage
1809 .set_bytes("test_table", &format!("key{:02}", i), vec![i as u8])
1810 .await
1811 .unwrap();
1812 }
1813
1814 let results = storage
1815 .get_range("test_table", "key00", "key05")
1816 .await
1817 .unwrap();
1818 assert_eq!(results.len(), 5);
1819 }
1820
1821 #[tokio::test]
1822 async fn test_merge() {
1823 let (storage, _temp) = create_test_storage().await;
1824
1825 storage
1826 .merge("test_table", "key1", vec![1, 2, 3])
1827 .await
1828 .unwrap();
1829
1830 storage
1831 .merge("test_table", "key1", vec![4, 5, 6])
1832 .await
1833 .unwrap();
1834
1835 let exists = storage.exists("test_table", "key1").await.unwrap();
1837 assert!(exists);
1838 }
1839
1840 #[tokio::test]
1841 async fn test_flush() {
1842 let (storage, _temp) = create_test_storage().await;
1843 let result = storage.flush().await;
1844 assert!(result.is_ok());
1845 }
1846
1847 #[tokio::test]
1848 async fn test_compact() {
1849 let (storage, _temp) = create_test_storage().await;
1850
1851 for i in 0..10 {
1853 let data = TestData {
1854 id: i.to_string(),
1855 name: format!("Test {}", i),
1856 value: i,
1857 };
1858 storage
1859 .insert("test_table", &format!("key{}", i), &data)
1860 .await
1861 .unwrap();
1862 }
1863
1864 let result = storage.compact().await;
1865 assert!(result.is_ok());
1866 }
1867
1868 #[tokio::test]
1869 async fn test_get_modified_since() {
1870 let (storage, _temp) = create_test_storage().await;
1871 let now = Utc::now();
1872
1873 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
1874
1875 let data = TestData {
1876 id: "1".to_string(),
1877 name: "Test".to_string(),
1878 value: 42,
1879 };
1880 storage.insert("test_table", "key1", &data).await.unwrap();
1881
1882 let results: Vec<StorageEntry<TestData>> = storage
1883 .get_modified_since("test_table", now)
1884 .await
1885 .unwrap();
1886
1887 assert_eq!(results.len(), 1);
1888 }
1889
1890 #[tokio::test]
1891 async fn test_cleanup_expired() {
1892 let (storage, _temp) = create_test_storage().await;
1893
1894 let data = TestData {
1896 id: "1".to_string(),
1897 name: "Test".to_string(),
1898 value: 42,
1899 };
1900 storage
1901 .insert_with_ttl("test_table", "key1", &data, 1)
1902 .await
1903 .unwrap();
1904
1905 tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
1906
1907 let removed = storage.cleanup_expired("test_table").await.unwrap();
1908 assert_eq!(removed, 1);
1909 }
1910}