Skip to main content

mocra_core/utils/
storage.rs

1use crate::errors::Result;
2use async_trait::async_trait;
3use std::sync::Arc;
4
5#[async_trait]
6pub trait BlobStorage: Send + Sync {
7    /// Upload data to storage, returns the access path/key/url
8    async fn put(&self, key: &str, data: &[u8]) -> Result<String>;
9    /// Download data from storage
10    async fn get(&self, key: &str) -> Result<Vec<u8>>;
11}
12
13#[async_trait]
14pub trait Offloadable {
15    fn should_offload(&self, threshold: usize) -> bool;
16    async fn offload(&mut self, storage: &Arc<dyn BlobStorage>) -> Result<()>;
17    async fn reload(&mut self, storage: &Arc<dyn BlobStorage>) -> Result<()>;
18}
19
20pub struct FileSystemBlobStorage {
21    root_path: std::path::PathBuf,
22}
23
24impl FileSystemBlobStorage {
25    pub fn new(root_path: impl Into<std::path::PathBuf>) -> Self {
26        Self {
27            root_path: root_path.into(),
28        }
29    }
30}
31
32#[async_trait]
33impl BlobStorage for FileSystemBlobStorage {
34    async fn put(&self, key: &str, data: &[u8]) -> Result<String> {
35        let path = self.root_path.join(key);
36        if let Some(parent) = path.parent() {
37            tokio::fs::create_dir_all(parent).await?;
38        }
39        tokio::fs::write(&path, data).await?;
40        Ok(path.to_string_lossy().to_string())
41    }
42
43    async fn get(&self, key: &str) -> Result<Vec<u8>> {
44        let path = std::path::PathBuf::from(key);
45        let data = tokio::fs::read(path).await?;
46        Ok(data)
47    }
48}